Support bounded parallel chunk transfers

Author: tyler-french

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

@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.remote;
 
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture;
 
 import build.bazel.remote.execution.v2.Digest;
@@ -23,10 +24,19 @@
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.LinkedBlockingQueue;
 
-/** Downloads blobs by sequentially fetching chunks via the SplitBlob API. */
+/** Downloads blobs by fetching chunks through a per-blob sliding window via the SplitBlob API. */
 public class ChunkedBlobDownloader {
+  // Guard against pathological fanout from a single large chunked blob. This is only a per-blob
+  // cap; chunk requests still flow through CombinedCache and the shared remote cache transport
+  // stack below it, which is what bounds active remote RPC concurrency across blobs.
+  private static final int MAX_IN_FLIGHT_CHUNK_DOWNLOADS = 32;
+
   private final GrpcCacheClient grpcCacheClient;
   private final CombinedCache combinedCache;
 
@@ -37,8 +47,8 @@ public ChunkedBlobDownloader(GrpcCacheClient grpcCacheClient, CombinedCache comb
 
   /**
    * Downloads a blob using chunked download via the SplitBlob API. This should be called with
-   * virtual threads, as it blocks on futures via {@link
-   * com.google.devtools.build.lib.remote.util.Utils#getFromFuture}.
+   * virtual threads, as it may block while waiting for chunk metadata and completed chunk
+   * downloads.
    */
   public void downloadChunked(
       RemoteActionExecutionContext context, Digest blobDigest, OutputStream out)
@@ -64,11 +74,123 @@ private List<Digest> getChunkDigests(RemoteActionExecutionContext context, Diges
     return chunkDigests;
   }
 
+  private static final class PendingDownload {
+    private final Digest digest;
+    private final ListenableFuture<byte[]> future;
+    private final List<Integer> chunkIndices = new ArrayList<>(1);
+
+    PendingDownload(Digest digest, ListenableFuture<byte[]> future, int firstChunkIndex) {
+      this.digest = digest;
+      this.future = future;
+      chunkIndices.add(firstChunkIndex);
+    }
+
+    void addChunkIndex(int chunkIndex) {
+      chunkIndices.add(chunkIndex);
+    }
+
+    Digest digest() {
+      return digest;
+    }
+
+    ListenableFuture<byte[]> future() {
+      return future;
+    }
+
+    List<Integer> chunkIndices() {
+      return chunkIndices;
+    }
+  }
+
   private void downloadAndReassembleChunks(
       RemoteActionExecutionContext context, List<Digest> chunkDigests, OutputStream out)
       throws IOException, InterruptedException {
-    for (Digest chunkDigest : chunkDigests) {
-      getFromFuture(combinedCache.downloadBlob(context, chunkDigest, out));
+    new DownloadSession(context, chunkDigests, out).run();
+  }
+
+  private final class DownloadSession {
+    private final LinkedBlockingQueue<PendingDownload> completedDownloads =
+        new LinkedBlockingQueue<>();
+    private final Map<Digest, PendingDownload> activeDownloads =
+        new HashMap<>(MAX_IN_FLIGHT_CHUNK_DOWNLOADS);
+    private final Map<Integer, byte[]> readyChunks =
+        new HashMap<>(MAX_IN_FLIGHT_CHUNK_DOWNLOADS);
+    private final RemoteActionExecutionContext context;
+    private final List<Digest> chunkDigests;
+    private final OutputStream out;
+    private int nextToStart = 0;
+    private int nextToWrite = 0;
+
+    DownloadSession(
+        RemoteActionExecutionContext context, List<Digest> chunkDigests, OutputStream out) {
+      this.context = context;
+      this.chunkDigests = chunkDigests;
+      this.out = out;
+    }
+
+    void run() throws IOException, InterruptedException {
+      try {
+        fillWindow();
+        while (nextToWrite < chunkDigests.size()) {
+          awaitCompletedDownload();
+          fillWindow();
+          drainReadyChunks();
+        }
+      } catch (IOException | InterruptedException | RuntimeException | Error e) {
+        cancelAllDownloads();
+        throw e;
+      }
+    }
+
+    private void fillWindow() {
+      while (nextToStart < chunkDigests.size()) {
+        Digest chunkDigest = chunkDigests.get(nextToStart);
+        PendingDownload existing = activeDownloads.get(chunkDigest);
+        if (existing != null) {
+          existing.addChunkIndex(nextToStart);
+          nextToStart++;
+          continue;
+        }
+        if (activeDownloads.size() >= MAX_IN_FLIGHT_CHUNK_DOWNLOADS) {
+          return;
+        }
+        startDownload(chunkDigest, nextToStart);
+        nextToStart++;
+      }
+    }
+
+    private void startDownload(Digest chunkDigest, int chunkIndex) {
+      PendingDownload download =
+          new PendingDownload(
+              chunkDigest, combinedCache.downloadBlob(context, chunkDigest), chunkIndex);
+      activeDownloads.put(chunkDigest, download);
+      download.future().addListener(() -> completedDownloads.add(download), directExecutor());
+    }
+
+    private void awaitCompletedDownload() throws IOException, InterruptedException {
+      PendingDownload download = completedDownloads.take();
+      activeDownloads.remove(download.digest());
+      byte[] chunkData = getFromFuture(download.future());
+      for (int chunkIndex : download.chunkIndices()) {
+        readyChunks.put(chunkIndex, chunkData);
+      }
+    }
+
+    private void drainReadyChunks() throws IOException {
+      while (true) {
+        byte[] chunk = readyChunks.remove(nextToWrite);
+        if (chunk == null) {
+          return;
+        }
+        out.write(chunk);
+        nextToWrite++;
+      }
+    }
+
+    private void cancelAllDownloads() {
+      for (PendingDownload download : activeDownloads.values()) {
+        download.future().cancel(/* mayInterruptIfRunning= */ true);
+      }
     }
   }
 }

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

@@ -14,11 +14,13 @@
 
 package com.google.devtools.build.lib.remote;
 
+import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture;
 
 import build.bazel.remote.execution.v2.Digest;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.ByteStreams;
+import com.google.common.util.concurrent.ListenableFuture;
 import com.google.devtools.build.lib.remote.chunking.ChunkingConfig;
 import com.google.devtools.build.lib.remote.chunking.FastCdcChunker;
 import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
@@ -30,6 +32,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.LinkedBlockingQueue;
 
 /**
  * Uploads blobs in chunks using Content-Defined Chunking with FastCDC 2020.
@@ -44,6 +47,10 @@
  * </ol>
  */
 public class ChunkedBlobUploader {
+  // Guard against pathological fanout from a single large chunked blob. This is only a per-blob
+  // cap; chunk uploads still flow through CombinedCache and the shared remote cache transport
+  // stack below it, which is what bounds active remote RPC concurrency across blobs.
+  private static final int MAX_IN_FLIGHT_CHUNK_UPLOADS = 32;
 
   private final GrpcCacheClient grpcCacheClient;
   private final CombinedCache combinedCache;
@@ -104,18 +111,95 @@ private void uploadMissingChunks(
     if (missingDigests.isEmpty()) {
       return;
     }
+    new UploadSession(context, missingDigests, chunkDigests).run(file);
+  }
 
-    Set<Digest> uploaded = new HashSet<>();
-    try (InputStream input = file.getInputStream()) {
-      for (Digest chunkDigest : chunkDigests) {
-        if (missingDigests.contains(chunkDigest) && uploaded.add(chunkDigest)) {
-          ByteString.Output out = ByteString.newOutput((int) chunkDigest.getSizeBytes());
-          ByteStreams.limit(input, chunkDigest.getSizeBytes()).transferTo(out);
-          getFromFuture(combinedCache.uploadBlob(context, chunkDigest, out.toByteString()));
-        } else {
-          input.skipNBytes(chunkDigest.getSizeBytes());
+  private final class UploadSession {
+    private final LinkedBlockingQueue<ListenableFuture<Void>> completedUploads =
+        new LinkedBlockingQueue<>();
+    private final Set<ListenableFuture<Void>> inFlightUploads =
+        new HashSet<>(MAX_IN_FLIGHT_CHUNK_UPLOADS);
+    private final Set<Digest> scheduledDigests = new HashSet<>();
+    private final RemoteActionExecutionContext context;
+    private final ImmutableSet<Digest> missingDigests;
+    private final List<Digest> chunkDigests;
+
+    UploadSession(
+        RemoteActionExecutionContext context,
+        ImmutableSet<Digest> missingDigests,
+        List<Digest> chunkDigests) {
+      this.context = context;
+      this.missingDigests = missingDigests;
+      this.chunkDigests = chunkDigests;
+    }
+
+    void run(Path file) throws IOException, InterruptedException {
+      try (InputStream input = file.getInputStream()) {
+        for (Digest chunkDigest : chunkDigests) {
+          drainCompletedUploads();
+          if (!shouldScheduleUpload(chunkDigest)) {
+            skipChunk(input, chunkDigest);
+            continue;
+          }
+          if (inFlightUploads.size() >= MAX_IN_FLIGHT_CHUNK_UPLOADS) {
+            awaitCompletedUpload();
+          }
+          startUpload(chunkDigest, readChunk(input, chunkDigest));
+        }
+        while (!inFlightUploads.isEmpty()) {
+          awaitCompletedUpload();
         }
+      } catch (IOException | InterruptedException | RuntimeException | Error e) {
+        cancelAllUploads();
+        throw e;
       }
     }
+
+    private boolean shouldScheduleUpload(Digest chunkDigest) {
+      return missingDigests.contains(chunkDigest) && scheduledDigests.add(chunkDigest);
+    }
+
+    private void startUpload(Digest chunkDigest, ByteString chunkData) {
+      ListenableFuture<Void> upload = combinedCache.uploadBlob(context, chunkDigest, chunkData);
+      inFlightUploads.add(upload);
+      upload.addListener(() -> completedUploads.add(upload), directExecutor());
+    }
+
+    private void drainCompletedUploads() throws IOException, InterruptedException {
+      while (true) {
+        ListenableFuture<Void> upload = completedUploads.poll();
+        if (upload == null) {
+          return;
+        }
+        finishUpload(upload);
+      }
+    }
+
+    private void awaitCompletedUpload() throws IOException, InterruptedException {
+      finishUpload(completedUploads.take());
+      drainCompletedUploads();
+    }
+
+    private void finishUpload(ListenableFuture<Void> upload)
+        throws IOException, InterruptedException {
+      inFlightUploads.remove(upload);
+      getFromFuture(upload);
+    }
+
+    private void cancelAllUploads() {
+      for (ListenableFuture<Void> upload : inFlightUploads) {
+        upload.cancel(/* mayInterruptIfRunning= */ true);
+      }
+    }
+  }
+
+  private static void skipChunk(InputStream input, Digest chunkDigest) throws IOException {
+    input.skipNBytes(chunkDigest.getSizeBytes());
+  }
+
+  private static ByteString readChunk(InputStream input, Digest chunkDigest) throws IOException {
+    ByteString.Output out = ByteString.newOutput((int) chunkDigest.getSizeBytes());
+    ByteStreams.limit(input, chunkDigest.getSizeBytes()).transferTo(out);
+    return out.toByteString();
   }
 }

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

@@ -1,4 +1,5 @@
 load("@rules_java//java:defs.bzl", "java_library", "java_test")
+load("//src:java_opt_binary.bzl", "java_opt_binary")
 
 package(
     default_applicable_licenses = ["//:license"],
@@ -65,6 +66,7 @@ java_library(
             "RemoteActionFileSystemTestBase.java",
             "BuildWithoutTheBytesIntegrationTest.java",
             "BuildWithoutTheBytesIntegrationTestBase.java",
+            "ChunkedTransferBenchmark.java",
             "ChunkedCacheIntegrationTest.java",
             "ChunkedDiskCacheIntegrationTest.java",
             "DiskCacheIntegrationTest.java",
@@ -251,6 +253,27 @@ java_test(
     ],
 )
 
+java_opt_binary(
+    name = "ChunkedTransferBenchmark",
+    srcs = ["ChunkedTransferBenchmark.java"],
+    main_class = "org.openjdk.jmh.Main",
+    deps = [
+        "@com_google_protobuf//java/core:lite_runtime_only",
+        "//src/main/java/com/google/devtools/build/lib/clock",
+        "//src/main/java/com/google/devtools/build/lib/remote:combined_cache",
+        "//src/main/java/com/google/devtools/build/lib/remote:grpc_cache_client",
+        "//src/main/java/com/google/devtools/build/lib/remote/chunking",
+        "//src/main/java/com/google/devtools/build/lib/remote/common",
+        "//src/main/java/com/google/devtools/build/lib/remote/util:digest_utils",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
+        "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
+        "//third_party:guava",
+        "//third_party:jmh",
+        "//third_party:mockito",
+        "@remoteapis//:build_bazel_remote_execution_v2_remote_execution_java_proto",
+    ],
+)
+
 java_library(
     name = "build_without_the_bytes_integration_test_base",
     srcs = [