src/main/java/com/google/devtools/build/lib/remote/ChunkedBlobUploader.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.base.Preconditions.checkNotNull;
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 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.ContentDefinedChunker;
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.common.RemoteCacheClient.Blob;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.vfs.Path;
import java.io.EOFException;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.channels.FileChannel;
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.
 *
 * <p>Upload flow for blobs above threshold:
 *
 * <ol>
 *   <li>Chunk file with the configured chunking function
 *   <li>Call findMissingDigests on chunk digests
 *   <li>Upload only missing chunks
 *   <li>Call SpliceBlob to register the blob as the concatenation of chunks
 * </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 = 16;

  private final RemoteCacheClient remoteCacheClient;
  private final CombinedCache combinedCache;
  private final ContentDefinedChunker chunker;
  private final ChunkingFunction.Value chunkingFunction;
  private final long chunkingThreshold;

  /**
   * Creates a new uploader with the given chunking configuration.
   *
   * @param remoteCacheClient client used for {@code FindMissingDigests} and {@code SpliceBlob} RPCs
   * @param combinedCache cache used to upload individual chunks
   * @param config chunking parameters negotiated from server capabilities
   */
  public ChunkedBlobUploader(
      RemoteCacheClient remoteCacheClient, CombinedCache combinedCache, ChunkingConfig config) {
    this.remoteCacheClient = remoteCacheClient;
    this.combinedCache = combinedCache;
    this.chunker = config.newChunker(combinedCache.digestUtil());
    this.chunkingFunction = config.chunkingFunction();
    this.chunkingThreshold = config.chunkingThreshold();
  }

  /** Returns the minimum blob size for chunked upload. */
  public long getChunkingThreshold() {
    return chunkingThreshold;
  }

  public void uploadChunked(RemoteActionExecutionContext context, Digest blobDigest, Path file)
      throws IOException, InterruptedException {
    uploadChunked(context, blobDigest, file, /* force= */ false);
  }

  /**
   * Uploads a blob in content-defined chunks. The file is chunked with the configured chunking
   * function, missing chunks are uploaded, and {@code SpliceBlob} is called to register the blob as
   * the concatenation of its chunks.
   *
   * @param force whether to upload missing chunks even if the cache client has already completed
   *     uploads of them
   */
  public void uploadChunked(
      RemoteActionExecutionContext context, Digest blobDigest, Path file, boolean force)
      throws IOException, InterruptedException {
    List<Digest> chunkDigests;
    try (InputStream input = file.getInputStream()) {
      chunkDigests = chunker.chunkToDigests(input);
    }
    if (chunkDigests.isEmpty()) {
      return;
    }

    ImmutableSet<Digest> missingDigests =
        getFromFuture(remoteCacheClient.findMissingDigests(context, chunkDigests));
    uploadMissingChunks(context, missingDigests, chunkDigests, file, force);
    getFromFuture(
        remoteCacheClient.spliceBlob(context, blobDigest, chunkDigests, chunkingFunction));
  }

  private void uploadMissingChunks(
      RemoteActionExecutionContext context,
      ImmutableSet<Digest> missingDigests,
      List<Digest> chunkDigests,
      Path file,
      boolean force)
      throws IOException, InterruptedException {
    if (missingDigests.isEmpty()) {
      return;
    }
    new UploadSession(context, missingDigests, chunkDigests, force).run(file);
  }

  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;
    private final boolean force;

    UploadSession(
        RemoteActionExecutionContext context,
        ImmutableSet<Digest> missingDigests,
        List<Digest> chunkDigests,
        boolean force) {
      this.context = context;
      this.missingDigests = missingDigests;
      this.chunkDigests = chunkDigests;
      this.force = force;
    }

    void run(Path file) throws IOException, InterruptedException {
      try {
        long offset = 0;
        for (Digest chunkDigest : chunkDigests) {
          drainCompletedUploads();
          long chunkOffset = offset;
          offset += chunkDigest.getSizeBytes();
          if (!shouldScheduleUpload(chunkDigest)) {
            continue;
          }
          if (inFlightUploads.size() >= MAX_IN_FLIGHT_CHUNK_UPLOADS) {
            awaitCompletedUpload();
          }
          startUpload(file, chunkOffset, chunkDigest, force);
        }
        while (!inFlightUploads.isEmpty()) {
          awaitCompletedUpload();
        }
      } finally {
        cancelAllUploads();
      }
    }

    private boolean shouldScheduleUpload(Digest chunkDigest) {
      return missingDigests.contains(chunkDigest) && scheduledDigests.add(chunkDigest);
    }

    private void startUpload(Path file, long chunkOffset, Digest chunkDigest, boolean force) {
      ListenableFuture<Void> upload =
          combinedCache.uploadBlob(
              context, chunkDigest, new ChunkBlob(file, chunkOffset, chunkDigest), force);
      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 final class ChunkBlob implements Blob {
    private final Path file;
    private final long offset;
    private final Digest digest;

    private ChunkBlob(Path file, long offset, Digest digest) {
      this.file = file;
      this.offset = offset;
      this.digest = digest;
    }

    @Override
    public InputStream get() throws IOException {
      InputStream input = file.getInputStream();
      boolean success = false;
      try {
        seekOrSkip(input, offset);
        InputStream limitedInput = ByteStreams.limit(input, digest.getSizeBytes());
        success = true;
        return limitedInput;
      } catch (EOFException e) {
        throw new IOException("file was concurrently modified during upload: " + file, e);
      } finally {
        if (!success) {
          input.close();
        }
      }
    }

    @Override
    public String description() {
      return "chunk %s at offset %d of file %s"
          .formatted(DigestUtil.toString(digest), offset, file);
    }
  }

  private static void seekOrSkip(InputStream input, long offset) throws IOException {
    if (offset == 0) {
      return;
    }
    if (input instanceof FileInputStream fileInputStream) {
      FileChannel channel = fileInputStream.getChannel();
      if (channel.size() < offset) {
        throw new EOFException();
      }
      channel.position(offset);
      return;
    }
    input.skipNBytes(offset);
  }
}