Reduce retention and early materialization of VirtualActionInput

bazelbuild@b82dcdc introduced a way for parameter files to be materialized lazily and in a streaming fashion. For this to be effective, the contents of `VirtualActionInput`s should not be retained and, ideally, never or only briefly materialized in full.

This is achieved by replacing the digest computations with a streaming approach and avoiding retention of `VirtualActionInput` in `MerkleTree`s, which may be cached. All cache protocols are streamlined to accept a `Blob`, which is a closeable supplier of `InputStream`s that is automatically closed after the upload has completed (which may require retries).

Along the way, file nodes in `MerkleTree`s are stored unwrapped to reduce memory usage.

Closes bazelbuild#24891.

PiperOrigin-RevId: 718042524
Change-Id: I1567cf6ca36cbf5e1146c91df34157c9192bb822

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

@@ -25,7 +25,6 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.PathStrippable;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
-import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.ArrayList;
@@ -239,13 +238,6 @@ public byte[] atomicallyWriteTo(Path outputPath) throws IOException {
       return super.atomicallyWriteTo(outputPath);
     }
 
-    @Override
-    public ByteString getBytes() throws IOException {
-      ByteString.Output out = ByteString.newOutput();
-      writeTo(out);
-      return out.toByteString();
-    }
-
     @Override
     public String getExecPathString() {
       return paramFileExecPath.getPathString();

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

@@ -117,8 +117,15 @@ protected byte[] writeTo(Path target) throws IOException {
   /**
    * 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 abstract ByteString getBytes() throws IOException;
+  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.
@@ -174,7 +181,7 @@ public void writeTo(OutputStream out) throws IOException {
     }
 
     @Override
-    public ByteString getBytes() throws IOException {
+    public ByteString getBytes() {
       return ByteString.EMPTY;
     }
 

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

@@ -59,7 +59,6 @@ java_library(
         "//third_party:error_prone_annotations",
         "//third_party:guava",
         "//third_party:jsr305",
-        "@com_google_protobuf//:protobuf_java",
     ],
 )
 
@@ -287,6 +286,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/remote/options",
+        "//src/main/java/com/google/devtools/build/lib/remote/util:digest_utils",
         "//src/main/java/com/google/devtools/build/lib/util/io",
         "//src/main/java/com/google/devtools/build/lib/util/io:io-proto",
         "//src/main/java/com/google/devtools/build/lib/vfs",
@@ -298,6 +298,7 @@ java_library(
         "//third_party:jsr305",
         "@com_google_protobuf//:protobuf_java",
         "@com_google_protobuf//:protobuf_java_util",
+        "@remoteapis//:build_bazel_remote_execution_v2_remote_execution_java_proto",
         "@zstd-jni",
     ],
 )

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

@@ -30,7 +30,6 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Symlinks;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
-import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -201,13 +200,6 @@ protected byte[] atomicallyWriteTo(Path outputPath) throws IOException {
       return digest;
     }
 
-    @Override
-    public ByteString getBytes() throws IOException {
-      ByteString.Output out = ByteString.newOutput();
-      writeTo(out);
-      return out.toByteString();
-    }
-
     @Override
     public FileArtifactValue getMetadata() throws IOException {
       // We intentionally delay hashing until it is necessary.

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

@@ -35,6 +35,7 @@
 import com.google.devtools.build.lib.exec.Protos.EnvironmentVariable;
 import com.google.devtools.build.lib.exec.Protos.Platform;
 import com.google.devtools.build.lib.remote.options.RemoteOptions;
+import com.google.devtools.build.lib.remote.util.DigestUtil;
 import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.vfs.DigestUtils;
 import com.google.devtools.build.lib.vfs.FileStatus;
@@ -167,11 +168,9 @@ protected Digest computeDigest(
 
     if (input != null) {
       if (input instanceof VirtualActionInput virtualActionInput) {
-        byte[] blob = virtualActionInput.getBytes().toByteArray();
-        return builder
-            .setHash(digestHashFunction.getHashFunction().hashBytes(blob).toString())
-            .setSizeBytes(blob.length)
-            .build();
+        build.bazel.remote.execution.v2.Digest digest =
+            DigestUtil.compute(virtualActionInput, digestHashFunction.getHashFunction());
+        return builder.setHash(digest.getHash()).setSizeBytes(digest.getSizeBytes()).build();
       }
 
       // Try to obtain a digest from the input metadata.

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

@@ -108,6 +108,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/remote/merkletree",
         "//src/main/java/com/google/devtools/build/lib/remote/options",
         "//src/main/java/com/google/devtools/build/lib/remote/util",
+        "//src/main/java/com/google/devtools/build/lib/remote/util:digest_utils",
         "//src/main/java/com/google/devtools/build/lib/remote/zstd",
         "//src/main/java/com/google/devtools/build/lib/skyframe:action_execution_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",

src/main/java/com/google/devtools/build/lib/remote/ByteStreamUploader.java

@@ -16,8 +16,6 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
 import static com.google.devtools.build.lib.remote.util.DigestUtil.isOldStyleDigestFunction;
-import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture;
-import static com.google.devtools.build.lib.remote.util.Utils.waitForBulkTransfer;
 import static java.lang.String.format;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
@@ -50,9 +48,6 @@
 import io.grpc.stub.ClientResponseObserver;
 import io.netty.util.ReferenceCounted;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.Semaphore;
 import javax.annotation.Nullable;
@@ -104,54 +99,6 @@ final class ByteStreamUploader {
     this.digestFunction = digestFunction;
   }
 
-  @VisibleForTesting
-  ReferenceCountedChannel getChannel() {
-    return channel;
-  }
-
-  @VisibleForTesting
-  RemoteRetrier getRetrier() {
-    return retrier;
-  }
-
-  /**
-   * Uploads a BLOB, as provided by the {@link Chunker}, to the remote {@code ByteStream} service.
-   * The call blocks until the upload is complete, or throws an {@link Exception} in case of error.
-   *
-   * <p>Uploads are retried according to the specified {@link RemoteRetrier}. Retrying is
-   * transparent to the user of this API.
-   *
-   * @param digest the digest of the data to upload.
-   * @param chunker the data to upload.
-   * @throws IOException when reading of the {@link Chunker}s input source fails
-   */
-  public void uploadBlob(RemoteActionExecutionContext context, Digest digest, Chunker chunker)
-      throws IOException, InterruptedException {
-    getFromFuture(uploadBlobAsync(context, digest, chunker));
-  }
-
-  /**
-   * Uploads a list of BLOBs concurrently to the remote {@code ByteStream} service. The call blocks
-   * until the upload of all BLOBs is complete, or throws an {@link
-   * com.google.devtools.build.lib.remote.common.BulkTransferException} if there are errors.
-   *
-   * <p>Uploads are retried according to the specified {@link RemoteRetrier}. Retrying is
-   * transparent to the user of this API.
-   *
-   * @param chunkers the data to upload.
-   * @throws IOException when reading of the {@link Chunker}s input source or uploading fails
-   */
-  public void uploadBlobs(RemoteActionExecutionContext context, Map<Digest, Chunker> chunkers)
-      throws IOException, InterruptedException {
-    List<ListenableFuture<Void>> uploads = new ArrayList<>();
-
-    for (Map.Entry<Digest, Chunker> chunkerEntry : chunkers.entrySet()) {
-      uploads.add(uploadBlobAsync(context, chunkerEntry.getKey(), chunkerEntry.getValue()));
-    }
-
-    waitForBulkTransfer(uploads, /* cancelRemainingOnInterrupt= */ true);
-  }
-
   /**
    * Uploads a BLOB asynchronously to the remote {@code ByteStream} service. The call returns
    * immediately and one can listen to the returned future for the success/failure of the upload.
@@ -163,7 +110,7 @@ public void uploadBlobs(RemoteActionExecutionContext context, Map<Digest, Chunke
    * performed. This is transparent to the user of this API.
    *
    * @param digest the {@link Digest} of the data to upload.
-   * @param chunker the data to upload.
+   * @param chunker the data to upload. Callers are responsible for closing the {@link Chunker}.
    */
   public ListenableFuture<Void> uploadBlobAsync(
       RemoteActionExecutionContext context, Digest digest, Chunker chunker) {

src/main/java/com/google/devtools/build/lib/remote/Chunker.java

@@ -22,11 +22,8 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.io.ByteStreams;
-import com.google.devtools.build.lib.actions.ActionInput;
-import com.google.devtools.build.lib.actions.ActionInputHelper;
-import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
+import com.google.devtools.build.lib.remote.common.RemoteCacheClient.Blob;
 import com.google.devtools.build.lib.remote.zstd.ZstdCompressingInputStream;
-import com.google.devtools.build.lib.vfs.Path;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.protobuf.ByteString;
 import java.io.ByteArrayInputStream;
@@ -46,7 +43,7 @@
  *
  * <p>This class should not be extended - it's only non-final for testing.
  */
-public class Chunker {
+public class Chunker implements AutoCloseable {
 
   private static int defaultChunkSize = 1024 * 16;
 
@@ -100,12 +97,7 @@ public int hashCode() {
     }
   }
 
-  /** A supplier that provide data as {@link InputStream}. */
-  public interface ChunkDataSupplier {
-    InputStream get() throws IOException;
-  }
-
-  private final ChunkDataSupplier dataSupplier;
+  private final Blob blob;
   private final long uncompressedSize;
   private final int chunkSize;
   private final Chunk emptyChunk;
@@ -120,9 +112,8 @@ public interface ChunkDataSupplier {
   // lazily on the first call to next(), as opposed to opening it in the constructor or on reset().
   private boolean initialized;
 
-  Chunker(
-      ChunkDataSupplier dataSupplier, long uncompressedSize, int chunkSize, boolean compressed) {
-    this.dataSupplier = checkNotNull(dataSupplier);
+  Chunker(Blob blob, long uncompressedSize, int chunkSize, boolean compressed) {
+    this.blob = checkNotNull(blob);
     this.uncompressedSize = uncompressedSize;
     this.chunkSize = chunkSize;
     this.emptyChunk = new Chunk(ByteString.EMPTY, 0);
@@ -143,7 +134,7 @@ public long getUncompressedSize() {
    * <p>Closes any open resources (file handles, ...).
    */
   public void reset() throws IOException {
-    close();
+    closeInput();
     offset = 0;
     initialized = false;
   }
@@ -168,7 +159,7 @@ public void seek(long toOffset) throws IOException {
       initialize(toOffset);
     }
     if (uncompressedSize > 0 && data.finished()) {
-      close();
+      closeInput();
     }
   }
 
@@ -180,14 +171,20 @@ public boolean hasNext() {
   }
 
   /** Closes the input stream and reset chunk cache */
-  private void close() throws IOException {
+  private void closeInput() throws IOException {
     if (data != null) {
       data.close();
       data = null;
     }
     chunkCache = null;
   }
 
+  @Override
+  public void close() throws IOException {
+    reset();
+    blob.close();
+  }
+
   /** Attempts reading at most a full chunk and stores it in the chunkCache buffer */
   private int read() throws IOException {
     int count = 0;
@@ -217,7 +214,7 @@ public Chunk next() throws IOException {
     maybeInitialize();
 
     if (uncompressedSize == 0) {
-      close();
+      closeInput();
       return emptyChunk;
     }
 
@@ -249,7 +246,7 @@ public Chunk next() throws IOException {
     // or the guard in getActualSize won't work.
     offset += bytesRead;
     if (data.finished()) {
-      close();
+      closeInput();
     }
 
     return new Chunk(blob, offsetBefore);
@@ -268,7 +265,7 @@ private void initialize(long srcPos) throws IOException {
     checkState(offset == 0);
     checkState(chunkCache == null);
     try {
-      var src = dataSupplier.get();
+      var src = blob.get();
       ByteStreams.skipFully(src, srcPos);
       data =
           compressed
@@ -294,49 +291,23 @@ public static class Builder {
     private int chunkSize = getDefaultChunkSize();
     protected long size;
     private boolean compressed;
-    protected ChunkDataSupplier inputStream;
-
-    @CanIgnoreReturnValue
-    public Builder setInput(byte[] data) {
-      checkState(inputStream == null);
-      size = data.length;
-      setInputSupplier(() -> new ByteArrayInputStream(data));
-      return this;
-    }
+    protected Blob inputStream;
 
     @CanIgnoreReturnValue
-    public Builder setInput(long size, InputStream in) {
+    public Builder setInput(long size, Blob in) {
       checkState(inputStream == null);
       checkNotNull(in);
       this.size = size;
-      inputStream = () -> in;
-      return this;
-    }
-
-    @CanIgnoreReturnValue
-    public Builder setInput(long size, Path file) {
-      checkState(inputStream == null);
-      this.size = size;
-      inputStream = file::getInputStream;
-      return this;
-    }
-
-    @CanIgnoreReturnValue
-    public Builder setInput(long size, ActionInput actionInput, Path execRoot) {
-      checkState(inputStream == null);
-      this.size = size;
-      if (actionInput instanceof VirtualActionInput virtualActionInput) {
-        inputStream = () -> virtualActionInput.getBytes().newInput();
-      } else {
-        inputStream = () -> ActionInputHelper.toInputPath(actionInput, execRoot).getInputStream();
-      }
+      inputStream = in;
       return this;
     }
 
     @CanIgnoreReturnValue
     @VisibleForTesting
-    protected final Builder setInputSupplier(ChunkDataSupplier inputStream) {
-      this.inputStream = inputStream;
+    public Builder setInput(byte[] data) {
+      checkState(inputStream == null);
+      size = data.length;
+      this.inputStream = () -> new ByteArrayInputStream(data);
       return this;
     }