Add --file_write_strategy to Bazel

src/main/java/com/google/devtools/build/lib/actions/BUILD

@@ -457,6 +457,7 @@ java_library(
     deps = [
         ":artifacts",
         ":has_digest",
+        "//src/main/java/com/google/devtools/build/lib/analysis:actions/deterministic_writer",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/io:inconsistent_filesystem_exception",
         "//src/main/java/com/google/devtools/build/lib/skyframe:filesystem_keys",

src/main/java/com/google/devtools/build/lib/actions/FileArtifactValue.java

@@ -21,12 +21,16 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import com.google.common.hash.HashFunction;
+import com.google.common.hash.HashingOutputStream;
 import com.google.common.io.BaseEncoding;
+import com.google.common.io.CountingOutputStream;
+import com.google.devtools.build.lib.analysis.actions.DeterministicWriter;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.HashCodes;
+import com.google.devtools.build.lib.util.StreamWriter;
 import com.google.devtools.build.lib.vfs.DigestUtils;
 import com.google.devtools.build.lib.vfs.FileStatus;
 import com.google.devtools.build.lib.vfs.Path;
@@ -36,8 +40,10 @@
 import com.google.devtools.build.lib.vfs.XattrProvider;
 import com.google.devtools.build.skyframe.SkyValue;
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.Objects;
@@ -67,7 +73,7 @@
  */
 @Immutable
 @ThreadSafe
-public abstract class FileArtifactValue implements SkyValue, HasDigest {
+public abstract class FileArtifactValue implements SkyValue, HasDigest, StreamWriter {
   /**
    * The type of the underlying file system object. If it is a regular file, then it is guaranteed
    * to have a digest. Otherwise it does not have a digest.
@@ -148,11 +154,30 @@ public InputStream getInputStream() {
     throw new UnsupportedOperationException();
   }
 
+  /**
+   * Writes the inline file contents to the given {@link OutputStream}.
+   *
+   * @throws UnsupportedOperationException if the file contents are not inline.
+   */
+  public void writeTo(OutputStream out) throws IOException {
+    try (var in = getInputStream()) {
+      in.transferTo(out);
+    }
+  }
+
   /** Returns whether the file contents exist remotely. */
   public boolean isRemote() {
     return false;
   }
 
+  /**
+   * Returns whether the file contents are materialized lazily, for example because they exist
+   * remotely.
+   */
+  public final boolean isLazy() {
+    return isRemote() || isInline();
+  }
+
   /** Returns the location index for remote files. For non-remote files, returns 0. */
   public int getLocationIndex() {
     return 0;
@@ -398,6 +423,20 @@ public static FileArtifactValue createForRemoteFileWithMaterializationData(
         digest, size, locationIndex, expirationTime);
   }
 
+  public static FileArtifactValue createForFileWriteActionOutput(
+      DeterministicWriter writer, HashFunction hashFunction) throws IOException {
+    long size;
+    byte[] digest;
+    try (CountingOutputStream countingOut =
+            new CountingOutputStream(OutputStream.nullOutputStream());
+        HashingOutputStream hashingOut = new HashingOutputStream(hashFunction, countingOut)) {
+      writer.writeOutputFile(hashingOut);
+      size = countingOut.getCount();
+      digest = hashingOut.hash().asBytes();
+    }
+    return new FileWriteOutputArtifactValue(writer, size, digest);
+  }
+
   public static FileArtifactValue createFromExistingWithResolvedPath(
       FileArtifactValue delegate, PathFragment resolvedPath) {
     return new ResolvedSymlinkArtifactValue(delegate, resolvedPath);
@@ -1004,6 +1043,105 @@ public boolean wasModifiedSinceDigest(Path path) {
     }
   }
 
+  private static final class FileWriteOutputArtifactValue extends FileArtifactValue {
+    private final DeterministicWriter writer;
+    private final long size;
+    private final byte[] digest;
+    @Nullable private FileContentsProxy proxy;
+
+    private FileWriteOutputArtifactValue(DeterministicWriter writer, long size, byte[] digest) {
+      this.writer = writer;
+      this.size = size;
+      this.digest = digest;
+    }
+
+    @Override
+    public FileStateType getType() {
+      return FileStateType.REGULAR_FILE;
+    }
+
+    @Override
+    public byte[] getDigest() {
+      return digest;
+    }
+
+    @Override
+    public long getSize() {
+      return size;
+    }
+
+    @Override
+    public boolean isInline() {
+      return true;
+    }
+
+    @Override
+    public InputStream getInputStream() {
+      // TODO: Avoid materializing the full content in memory by using a variant of
+      //  Piped{Input,Output}Stream that works well with virtual threads.
+      var out = new ByteArrayOutputStream(Math.clamp(getSize(), 0, Integer.MAX_VALUE));
+      try {
+        writeTo(out);
+      } catch (IOException e) {
+        // writer is not expected to throw if out doesn't.
+        throw new IllegalStateException(e);
+      }
+      return new ByteArrayInputStream(out.toByteArray());
+    }
+
+    @Override
+    public void writeTo(OutputStream out) throws IOException {
+      writer.writeOutputFile(out);
+    }
+
+    @Override
+    public long getModifiedTime() {
+      throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Sets the {@link FileContentsProxy} if the output backed by this metadata is materialized
+     * later.
+     */
+    @Override
+    public void setContentsProxy(FileContentsProxy proxy) {
+      this.proxy = proxy;
+    }
+
+    /**
+     * Returns a non-null {@link FileContentsProxy} if this metadata is backed by a local file, e.g.
+     * the file is materialized after action execution.
+     */
+    @Override
+    @Nullable
+    public FileContentsProxy getContentsProxy() {
+      return proxy;
+    }
+
+    @Override
+    public boolean wasModifiedSinceDigest(Path path) {
+      return false;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof FileWriteOutputArtifactValue that)) {
+        return false;
+      }
+      return Objects.equals(writer, that.writer)
+          && size == that.size
+          && Arrays.equals(digest, that.digest);
+    }
+
+    @Override
+    public int hashCode() {
+      return HashCodes.hashObjects(writer, size, Arrays.hashCode(digest));
+    }
+  }
+
   /** Metadata for files whose contents are available in memory. */
   private static final class InlineFileArtifactValue extends FileArtifactValue {
     private final byte[] data;

src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java

@@ -114,19 +114,6 @@ protected byte[] writeTo(Path target) throws IOException {
     return digest;
   }
 
-  /**
-   * Gets a {@link ByteString} representation of the fake file. Used to avoid copying if the fake
-   * file is internally represented as a {@link ByteString}.
-   *
-   * <p>Prefer {@link #writeTo} to this method to avoid materializing the entire file in memory. The
-   * return value should not be retained.
-   */
-  public ByteString getBytes() throws IOException {
-    ByteString.Output out = ByteString.newOutput();
-    writeTo(out);
-    return out.toByteString();
-  }
-
   /**
    * Returns the metadata for this input if available. Null otherwise.
    *

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -1663,6 +1663,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/concurrent",
+        "//src/main/java/com/google/devtools/build/lib/unsafe:string",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:detailed_exit_code",
         "//src/main/java/com/google/devtools/build/lib/util:string",

src/main/java/com/google/devtools/build/lib/analysis/actions/DeterministicWriter.java

@@ -22,6 +22,9 @@
  * invocation of writeOutputFile().
  */
 public interface DeterministicWriter {
+  /**
+   * @throws IOException only if out throws an IOException
+   */
   void writeOutputFile(OutputStream out) throws IOException;
 
   /**

src/main/java/com/google/devtools/build/lib/analysis/actions/FileWriteActionContext.java

@@ -56,4 +56,12 @@ default ImmutableList<SpawnResult> writeOutputToFile(
         isRemotable,
         Iterables.getOnlyElement(action.getOutputs()));
   }
+
+  /**
+   * Returns whether the {@link DeterministicWriter} may be retained after {@link
+   * #writeOutputToFile} returns.
+   */
+  default boolean mayRetainWriter() {
+    return false;
+  }
 }

src/main/java/com/google/devtools/build/lib/analysis/actions/LocalTemplateExpansionStrategy.java

@@ -14,7 +14,6 @@
 
 package com.google.devtools.build.lib.analysis.actions;
 
-import static java.nio.charset.StandardCharsets.ISO_8859_1;
 
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.actions.AbstractAction;
@@ -25,6 +24,7 @@
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.server.FailureDetails.Execution;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
+import com.google.devtools.build.lib.unsafe.StringUnsafe;
 import com.google.devtools.build.lib.util.StringUtilities;
 import java.io.IOException;
 import java.util.List;
@@ -44,18 +44,43 @@ public ImmutableList<SpawnResult> expandTemplate(
       TemplateExpansionContext.TemplateMetadata templateMetadata)
       throws InterruptedException, ExecException {
     try {
+      // If writeOutputToFile may retain the writer, make sure that it doesn't capture the
+      // expanded template string. Since expansion may throw and the writer must not, expand the
+      // template once before calling writeOutputToFile. It is assumed to be deterministic, so if
+      // it doesn't throw once, it won't throw again.
       final String expandedTemplate =
           getExpandedTemplateUnsafe(
               templateMetadata.template(), templateMetadata.substitutions(), ctx.getPathResolver());
-      DeterministicWriter deterministicWriter =
-          out -> out.write(expandedTemplate.getBytes(ISO_8859_1));
-      return ctx.getContext(FileWriteActionContext.class)
-          .writeOutputToFile(
-              action,
-              ctx,
-              deterministicWriter,
-              templateMetadata.makeExecutable(),
-              /* isRemotable= */ true);
+      FileWriteActionContext fileWriteActionContext = ctx.getContext(FileWriteActionContext.class);
+      DeterministicWriter deterministicWriter;
+      if (fileWriteActionContext.mayRetainWriter()) {
+        ArtifactPathResolver pathResolver = ctx.getPathResolver();
+        deterministicWriter =
+            out -> {
+              try {
+                out.write(
+                    StringUnsafe.getInternalStringBytes(
+                        getExpandedTemplateUnsafe(
+                            templateMetadata.template(),
+                            templateMetadata.substitutions(),
+                            pathResolver)));
+              } catch (EvalException e) {
+                throw new IllegalStateException(
+                    "Template expansion is not deterministic, first succeeded and then failed with: "
+                        + e.getMessage(),
+                    e);
+              }
+            };
+      } else {
+        deterministicWriter =
+            out -> out.write(StringUnsafe.getInternalStringBytes(expandedTemplate));
+      }
+      return fileWriteActionContext.writeOutputToFile(
+          action,
+          ctx,
+          deterministicWriter,
+          templateMetadata.makeExecutable(),
+          /* isRemotable= */ true);
     } catch (IOException | EvalException e) {
       throw new EnvironmentalExecException(
           e,

src/main/java/com/google/devtools/build/lib/exec/BUILD

@@ -141,10 +141,14 @@ java_library(
 
 java_library(
     name = "file_write_strategy",
-    srcs = ["FileWriteStrategy.java"],
+    srcs = [
+        "EagerFileWriteStrategy.java",
+        "LazyFileWriteStrategy.java",
+    ],
     deps = [
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/deterministic_writer",
         "//src/main/java/com/google/devtools/build/lib/analysis:actions/file_write_action_context",
         "//src/main/java/com/google/devtools/build/lib/profiler",

src/main/java/com/google/devtools/build/lib/exec/FileWriteStrategy.java → src/main/java/com/google/devtools/build/lib/exec/EagerFileWriteStrategy.java

@@ -35,9 +35,10 @@
 
 /**
  * A strategy for executing an {@link
- * com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction}.
+ * com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction} that writes the file to
+ * disk eagerly.
  */
-public final class FileWriteStrategy implements FileWriteActionContext {
+public class EagerFileWriteStrategy implements FileWriteActionContext {
   private static final Duration MIN_LOGGING = Duration.ofMillis(100);
 
   @Override

src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java

@@ -562,6 +562,29 @@ public boolean usingLocalTestJobs() {
               + " code 39.")
   public int remoteRetryOnTransientCacheError;
 
+  /** An enum for specifying different file write strategies. */
+  public enum FileWriteStrategy {
+    EAGER,
+    LAZY;
+
+    /** Converts to {@link FileWriteStrategy}. */
+    public static class Converter extends EnumConverter<FileWriteStrategy> {
+      public Converter() {
+        super(FileWriteStrategy.class, "file write strategy");
+      }
+    }
+  }
+
+  @Option(
+      name = "file_write_strategy",
+      defaultValue = "eager",
+      converter = FileWriteStrategy.Converter.class,
+      documentationCategory = OptionDocumentationCategory.REMOTE,
+      effectTags = {OptionEffectTag.EXECUTION},
+      help =
+          "If set to 'eager', Bazel will write the output of file write actions to disk. If set to 'lazy', their output will be kept in-memory if possible and is only written out if needed or requested, depending on the value of --remote_download_outputs.")
+  public FileWriteStrategy fileWriteStrategy;
+
   /** An enum for specifying different formats of test output. */
   public enum TestOutputFormat {
     SUMMARY, // Provide summary output only.