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

// Copyright 2026 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.google.devtools.build.lib.remote;

import static com.google.common.util.concurrent.Futures.immediateFuture;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static com.google.devtools.build.lib.remote.util.Futures.getFromFuture;

import build.bazel.remote.execution.v2.ChunkingFunction;
import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.SplitBlobResponse;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.devtools.build.lib.remote.chunking.ChunkingConfig;
import com.google.devtools.build.lib.remote.common.BlobNotSplittableException;
import com.google.devtools.build.lib.remote.common.MaybePathBacked;
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
import com.google.devtools.build.lib.remote.common.RemoteCacheClient;
import com.google.devtools.build.lib.remote.util.DigestOutputStream;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.remote.util.Utils;
import com.google.devtools.build.lib.vfs.Path;
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;
import javax.annotation.Nullable;

/** 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 = 16;

  private final RemoteCacheClient remoteCacheClient;
  private final CombinedCache combinedCache;
  private final ChunkLocationMap chunkLocationMap;
  private final ChunkingFunction.Value chunkingFunction;
  private final long maxChunkSize;

  public ChunkedBlobDownloader(
      RemoteCacheClient remoteCacheClient,
      CombinedCache combinedCache,
      ChunkingConfig chunkingConfig,
      ChunkLocationMap chunkLocationMap) {
    this.remoteCacheClient = remoteCacheClient;
    this.combinedCache = combinedCache;
    this.chunkLocationMap = chunkLocationMap;
    this.chunkingFunction = chunkingConfig.chunkingFunction();
    this.maxChunkSize = chunkingConfig.maxChunkSize();
  }

  /**
   * Downloads a blob using chunked download via the SplitBlob API. This should be called with
   * virtual threads, as it may block while waiting for chunk metadata and completed chunk
   * downloads.
   */
  public void downloadChunked(
      RemoteActionExecutionContext context, Digest blobDigest, OutputStream out)
      throws IOException, InterruptedException {
    // The file being written to must not be used as a chunk source, while the final path (which
    // differs when the caller stages the download in a temporary location) is recorded as one.
    @Nullable Path destination = null;
    @Nullable Path finalDestination = null;
    if (out instanceof MaybePathBacked pathBacked) {
      destination = pathBacked.maybeGetPath();
      finalDestination = pathBacked.maybeGetFinalPath();
    }
    @Nullable DigestOutputStream digestOut = null;
    if (remoteCacheClient.shouldVerifyDownloads()) {
      digestOut = combinedCache.digestUtil().newDigestOutputStream(out);
      out = digestOut;
    }

    List<Digest> chunkDigests = getChunkDigests(context, blobDigest);
    new DownloadSession(context, chunkDigests, destination, out).run();
    if (digestOut != null) {
      Utils.verifyBlobContents(blobDigest, digestOut.digest());
    }
    if (finalDestination != null) {
      // The content may not be at the final destination yet (out may not be flushed, and a staged
      // download has yet to be moved into place), but locations are only hints: an early read is
      // just a miss.
      chunkLocationMap.addFile(finalDestination, chunkDigests);
    }
  }

  /**
   * Returns the chunks the blob is composed of.
   *
   * @throws BlobNotSplittableException if the server cannot describe the blob as a sequence of
   *     chunks, in which case the caller may fall back to downloading the whole blob
   */
  private List<Digest> getChunkDigests(RemoteActionExecutionContext context, Digest blobDigest)
      throws IOException, InterruptedException {
    if (blobDigest.getSizeBytes() == 0) {
      return ImmutableList.of();
    }
    ListenableFuture<SplitBlobResponse> splitResponseFuture =
        remoteCacheClient.splitBlob(context, blobDigest, chunkingFunction);
    if (splitResponseFuture == null) {
      throw new BlobNotSplittableException(blobDigest);
    }
    List<Digest> chunkDigests = getFromFuture(splitResponseFuture).getChunkDigestsList();
    if (chunkDigests.isEmpty()) {
      throw new BlobNotSplittableException(blobDigest);
    }
    validateChunkDigests(blobDigest, chunkDigests);
    return chunkDigests;
  }

  private void validateChunkDigests(Digest blobDigest, List<Digest> chunkDigests)
      throws IOException {
    long remainingSize = blobDigest.getSizeBytes();
    if (remainingSize < 0) {
      throw new IOException(
          "Invalid SplitBlob response for %s: blob size is negative"
              .formatted(DigestUtil.toString(blobDigest)));
    }
    for (Digest chunkDigest : chunkDigests) {
      long chunkSize = chunkDigest.getSizeBytes();
      if (chunkSize <= 0) {
        throw new IOException(
            "Invalid SplitBlob response for %s: chunk %s has non-positive size"
                .formatted(DigestUtil.toString(blobDigest), DigestUtil.toString(chunkDigest)));
      }
      if (chunkSize > maxChunkSize) {
        throw new IOException(
            "Invalid SplitBlob response for %s: chunk %s exceeds max chunk size %d"
                .formatted(
                    DigestUtil.toString(blobDigest),
                    DigestUtil.toString(chunkDigest),
                    maxChunkSize));
      }
      if (chunkSize > remainingSize) {
        throw new IOException(
            "Invalid SplitBlob response for %s: chunk sizes exceed blob size"
                .formatted(DigestUtil.toString(blobDigest)));
      }
      remainingSize -= chunkSize;
    }
    if (remainingSize != 0) {
      throw new IOException(
          "Invalid SplitBlob response for %s: chunk sizes do not match blob size"
              .formatted(DigestUtil.toString(blobDigest)));
    }
  }

  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 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;
    @Nullable private final Path destination;
    private final OutputStream out;
    private int nextToStart = 0;
    private int nextToWrite = 0;

    DownloadSession(
        RemoteActionExecutionContext context,
        List<Digest> chunkDigests,
        @Nullable Path destination,
        OutputStream out) {
      this.context = context;
      this.chunkDigests = chunkDigests;
      this.destination = destination;
      this.out = out;
    }

    void run() throws IOException, InterruptedException {
      try {
        fillWindow();
        while (nextToWrite < chunkDigests.size()) {
          drainCompletedDownloads();
          drainReadyChunks();
          fillWindow();
        }
      } finally {
        cancelAllDownloads();
      }
    }

    private void fillWindow() {
      while (nextToStart < chunkDigests.size()) {
        if (nextToStart - nextToWrite >= MAX_IN_FLIGHT_CHUNK_DOWNLOADS) {
          return;
        }
        Digest chunkDigest = chunkDigests.get(nextToStart);
        PendingDownload existing = activeDownloads.get(chunkDigest);
        if (existing != null) {
          existing.addChunkIndex(nextToStart);
          nextToStart++;
          continue;
        }
        startDownload(chunkDigest, nextToStart);
        nextToStart++;
      }
    }

    private void startDownload(Digest chunkDigest, int chunkIndex) {
      PendingDownload download =
          new PendingDownload(chunkDigest, fetchChunk(chunkDigest), chunkIndex);
      activeDownloads.put(chunkDigest, download);
      download.future().addListener(() -> completedDownloads.add(download), directExecutor());
    }

    /** Returns the contents of a chunk, preferring a copy already on local disk over a download. */
    private ListenableFuture<byte[]> fetchChunk(Digest chunkDigest) {
      byte[] local = chunkLocationMap.read(chunkDigest, destination, combinedCache.digestUtil());
      return local != null
          ? immediateFuture(local)
          : combinedCache.downloadBlob(context, chunkDigest);
    }

    private void drainCompletedDownloads() throws IOException, InterruptedException {
      PendingDownload download = completedDownloads.take();
      do {
        processCompletedDownload(download);
        download = completedDownloads.poll();
      } while (download != null);
    }

    private void processCompletedDownload(PendingDownload download)
        throws IOException, InterruptedException {
      activeDownloads.remove(download.digest());
      byte[] chunkData = getFromFuture(download.future());
      for (int chunkIndex : download.chunkIndices()) {
        if (chunkIndex == nextToWrite) {
          out.write(chunkData);
          nextToWrite++;
        } else {
          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);
      }
    }
  }
}