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

// Copyright 2017 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.MoreObjects.firstNonNull;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
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.common.ProgressStatusListener.NO_ACTION;
import static com.google.devtools.build.lib.remote.util.Futures.getFromFuture;
import static com.google.devtools.build.lib.util.StringUtilities.bytesCountToDisplayString;

import build.bazel.remote.execution.v2.ActionResult;
import build.bazel.remote.execution.v2.CacheCapabilities;
import build.bazel.remote.execution.v2.ChunkingFunction;
import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.ServerCapabilities;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.flogger.GoogleLogger;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.devtools.build.lib.concurrent.ThreadSafety;
import com.google.devtools.build.lib.exec.SpawnCheckingCacheEvent;
import com.google.devtools.build.lib.exec.SpawnProgressEvent;
import com.google.devtools.build.lib.remote.chunking.ChunkingConfig;
import com.google.devtools.build.lib.remote.common.ActionKey;
import com.google.devtools.build.lib.remote.common.BlobNotSplittableException;
import com.google.devtools.build.lib.remote.common.CacheNotFoundException;
import com.google.devtools.build.lib.remote.common.LazyFileOutputStream;
import com.google.devtools.build.lib.remote.common.MaybePathBacked;
import com.google.devtools.build.lib.remote.common.OutputDigestMismatchException;
import com.google.devtools.build.lib.remote.common.ProgressStatusListener;
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.disk.DiskCacheClient;
import com.google.devtools.build.lib.remote.options.RemoteOptions.ChunkingFunctionValue;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
import com.google.devtools.build.lib.server.FailureDetails.RemoteExecution;
import com.google.devtools.build.lib.server.FailureDetails.RemoteExecution.Code;
import com.google.devtools.build.lib.util.io.OutErr;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.PathFragment;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.protobuf.ByteString;
import io.netty.util.AbstractReferenceCounted;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import javax.annotation.Nullable;

/**
 * Provides unified access to a disk cache, remote cache, or both.
 *
 * <p>The cache is reference counted. Initially, the reference count is 1. Use {@link #retain()} to
 * increase and {@link #release()} to decrease the reference count respectively. Once the reference
 * count is reached to 0, the underlying resources will be released (after pending I/O is finished).
 *
 * <p>Use {@link #awaitTermination()} to wait for pending I/O to finish. Use {@link #shutdownNow()}
 * to cancel all pending I/O and reject new requests.
 */
@ThreadSafety.ThreadSafe
public class CombinedCache extends AbstractReferenceCounted {
  private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

  private static final ListenableFuture<Void> COMPLETED_SUCCESS = immediateFuture(null);
  private static final ListenableFuture<byte[]> EMPTY_BYTES = immediateFuture(new byte[0]);
  private static final SpawnCheckingCacheEvent SPAWN_CHECKING_DISK_CACHE_EVENT =
      SpawnCheckingCacheEvent.create("disk-cache");
  private static final SpawnCheckingCacheEvent SPAWN_CHECKING_REMOTE_CACHE_EVENT =
      SpawnCheckingCacheEvent.create("remote-cache");

  private final CountDownLatch closeCountDownLatch = new CountDownLatch(1);

  private final ListeningExecutorService virtualThreadExecutor =
      MoreExecutors.listeningDecorator(
          Executors.newThreadPerTaskExecutor(
              Thread.ofVirtual().name("combined-cache-", 0).factory()));

  @Nullable protected final RemoteCacheClient remoteCacheClient;
  @Nullable protected final DiskCacheClient diskCacheClient;
  @Nullable protected final String symlinkTemplate;
  private final DigestUtil digestUtil;

  static class Chunking {
    @Nullable private final ChunkingConfig config;
    @Nullable private final ChunkedBlobDownloader downloader;
    @Nullable private final ChunkedBlobUploader uploader;

    Chunking() {
      this.config = null;
      this.downloader = null;
      this.uploader = null;
    }

    Chunking(
        ServerCapabilities serverCapabilities,
        RemoteCacheClient remoteCacheClient,
        ChunkingFunctionValue chunkingFunction,
        ChunkLocationMap chunkLocationMap,
        CombinedCache combinedCache) {
      config =
          switch (chunkingFunction) {
            case AUTO -> ChunkingConfig.fromServerCapabilities(serverCapabilities);
            case FAST_CDC_2020 ->
                ChunkingConfig.fromServerCapabilities(
                    serverCapabilities, ChunkingFunction.Value.FAST_CDC_2020);
            case REP_MAX_CDC ->
                ChunkingConfig.fromServerCapabilities(
                    serverCapabilities, ChunkingFunction.Value.REP_MAX_CDC);
          };
      if (config != null) {
        downloader =
            new ChunkedBlobDownloader(remoteCacheClient, combinedCache, config, chunkLocationMap);
        uploader = new ChunkedBlobUploader(remoteCacheClient, combinedCache, config);
      } else {
        downloader = null;
        uploader = null;
      }
    }

    boolean supported() {
      return config != null;
    }

    ChunkingConfig config() {
      return checkNotNull(config, "must not call config() unless supported() is true");
    }

    ChunkedBlobDownloader downloader() {
      return checkNotNull(downloader, "must not call downloader() unless supported() is true");
    }

    ChunkedBlobUploader uploader() {
      return checkNotNull(uploader, "must not call uploader() unless supported() is true");
    }
  }

  private final ListenableFuture<Chunking> chunking;

  private static ListenableFuture<Chunking> createChunking(
      @Nullable RemoteCacheClient remoteCacheClient,
      @Nullable ChunkingFunctionValue chunkingFunction,
      ChunkLocationMap chunkLocationMap,
      CombinedCache combinedCache) {
    if (remoteCacheClient == null || chunkingFunction == null) {
      return immediateFuture(new Chunking());
    }
    return Futures.transform(
        remoteCacheClient.serverCapabilities(),
        serverCapabilities ->
            new Chunking(
                serverCapabilities,
                remoteCacheClient,
                chunkingFunction,
                chunkLocationMap,
                combinedCache),
        directExecutor());
  }

  public CombinedCache(
      @Nullable RemoteCacheClient remoteCacheClient,
      @Nullable DiskCacheClient diskCacheClient,
      @Nullable String symlinkTemplate,
      DigestUtil digestUtil,
      @Nullable ChunkingFunctionValue chunkingFunction,
      ChunkLocationMap chunkLocationMap) {
    checkArgument(
        remoteCacheClient != null || diskCacheClient != null,
        "remoteCacheClient and diskCacheClient cannot be null at the same time");
    this.remoteCacheClient = remoteCacheClient;
    this.diskCacheClient = diskCacheClient;
    this.symlinkTemplate = symlinkTemplate;
    this.digestUtil = digestUtil;
    this.chunking = createChunking(remoteCacheClient, chunkingFunction, chunkLocationMap, this);
  }

  public CombinedCache(
      @Nullable RemoteCacheClient remoteCacheClient,
      @Nullable DiskCacheClient diskCacheClient,
      @Nullable String symlinkTemplate,
      DigestUtil digestUtil,
      ListenableFuture<Chunking> chunking) {
    checkArgument(
        remoteCacheClient != null || diskCacheClient != null,
        "remoteCacheClient and diskCacheClient cannot be null at the same time");
    this.remoteCacheClient = remoteCacheClient;
    this.diskCacheClient = diskCacheClient;
    this.symlinkTemplate = symlinkTemplate;
    this.digestUtil = digestUtil;
    this.chunking = chunking;
  }

  public DigestUtil digestUtil() {
    return digestUtil;
  }

  public CacheCapabilities getRemoteCacheCapabilities() throws IOException {
    return getRemoteServerCapabilities().getCacheCapabilities();
  }

  public ListenableFuture<String> getRemoteAuthority() {
    if (remoteCacheClient == null) {
      return immediateFuture("");
    }
    return remoteCacheClient.getAuthority();
  }

  public ServerCapabilities getRemoteServerCapabilities() throws IOException {
    if (remoteCacheClient == null) {
      return ServerCapabilities.getDefaultInstance();
    }
    return remoteCacheClient.getServerCapabilities();
  }

  /**
   * Class to keep track of which cache (disk or remote) a given [cached] ActionResult comes from.
   */
  public record CachedActionResult(ActionResult actionResult, String cacheName) {
    @Nullable
    public static CachedActionResult remote(ActionResult actionResult) {
      if (actionResult == null) {
        return null;
      }
      return new CachedActionResult(actionResult, "remote");
    }

    @Nullable
    public static CachedActionResult disk(ActionResult actionResult) {
      if (actionResult == null) {
        return null;
      }
      return new CachedActionResult(actionResult, "disk");
    }
  }

  public CachedActionResult downloadActionResult(
      RemoteActionExecutionContext context,
      ActionKey actionKey,
      boolean inlineOutErr,
      Set<String> inlineOutputFiles)
      throws IOException, InterruptedException {
    return getFromFuture(
        downloadActionResultAsync(context, actionKey, inlineOutErr, inlineOutputFiles));
  }

  public ListenableFuture<CachedActionResult> downloadActionResultAsync(
      RemoteActionExecutionContext context,
      ActionKey actionKey,
      boolean inlineOutErr,
      Set<String> inlineOutputFiles) {
    var spawnExecutionContext = context.getSpawnExecutionContext();

    ListenableFuture<CachedActionResult> future = immediateFuture(null);

    if (diskCacheClient != null && context.getReadCachePolicy().allowDiskCache()) {
      if (spawnExecutionContext != null) {
        spawnExecutionContext.report(SPAWN_CHECKING_DISK_CACHE_EVENT);
      }
      future =
          Futures.transform(
              diskCacheClient.downloadActionResult(actionKey),
              CachedActionResult::disk,
              directExecutor());
    }

    if (remoteCacheClient != null && context.getReadCachePolicy().allowRemoteCache()) {
      future =
          Futures.transformAsync(
              future,
              (result) -> {
                if (result == null) {
                  if (spawnExecutionContext != null) {
                    spawnExecutionContext.report(SPAWN_CHECKING_REMOTE_CACHE_EVENT);
                  }
                  return Futures.transform(
                      downloadActionResultFromRemote(
                          context, actionKey, inlineOutErr, inlineOutputFiles),
                      CachedActionResult::remote,
                      directExecutor());
                } else {
                  return immediateFuture(result);
                }
              },
              directExecutor());
    }

    return future;
  }

  private ListenableFuture<ActionResult> downloadActionResultFromRemote(
      RemoteActionExecutionContext context,
      ActionKey actionKey,
      boolean inlineOutErr,
      Set<String> inlineOutputFiles) {
    checkState(remoteCacheClient != null && context.getReadCachePolicy().allowRemoteCache());
    return Futures.transformAsync(
        remoteCacheClient.downloadActionResult(context, actionKey, inlineOutErr, inlineOutputFiles),
        (actionResult) -> {
          if (actionResult == null) {
            return immediateFuture(null);
          }

          if (diskCacheClient != null && context.getWriteCachePolicy().allowDiskCache()) {
            return Futures.transform(
                diskCacheClient.uploadActionResult(actionKey, actionResult),
                v -> actionResult,
                directExecutor());
          }

          return immediateFuture(actionResult);
        },
        directExecutor());
  }

  /**
   * Returns a set of digests that the remote cache does not know about. The returned set is
   * guaranteed to be a subset of {@code digests}.
   */
  public ListenableFuture<ImmutableSet<Digest>> findMissingDigests(
      RemoteActionExecutionContext context, Iterable<Digest> digests) {
    if (Iterables.isEmpty(digests)) {
      return immediateFuture(ImmutableSet.of());
    }

    if (remoteCacheClient != null && context.getWriteCachePolicy().allowRemoteCache()) {
      return remoteCacheClient.findMissingDigests(context, digests);
    }

    return immediateFuture(ImmutableSet.of());
  }

  /** Returns whether the remote action cache supports updating action results. */
  public boolean remoteActionCacheSupportsUpdate() {
    try {
      return getRemoteCacheCapabilities().getActionCacheUpdateCapabilities().getUpdateEnabled();
    } catch (IOException ignored) {
      return false;
    }
  }

  /** Upload the action result to the remote cache. */
  public ListenableFuture<Void> uploadActionResult(
      RemoteActionExecutionContext context, ActionKey actionKey, ActionResult actionResult) {
    ListenableFuture<Void> diskCacheFuture = Futures.immediateVoidFuture();
    if (diskCacheClient != null && context.getWriteCachePolicy().allowDiskCache()) {
      diskCacheFuture = diskCacheClient.uploadActionResult(actionKey, actionResult);
    }

    ListenableFuture<Void> remoteCacheFuture = Futures.immediateVoidFuture();
    if (remoteCacheClient != null && context.getWriteCachePolicy().allowRemoteCache()) {
      remoteCacheFuture = remoteCacheClient.uploadActionResult(context, actionKey, actionResult);
    }

    return Futures.whenAllSucceed(diskCacheFuture, remoteCacheFuture)
        .call(() -> null, directExecutor());
  }

  private ListenableFuture<Void> uploadFileToDisk(Digest digest, Path file) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }
    if (diskCacheClient == null) {
      return Futures.immediateVoidFuture();
    }
    return diskCacheClient.uploadFile(digest, file);
  }

  private ListenableFuture<Void> uploadBlobToDisk(Digest digest, Blob blob) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }
    if (diskCacheClient == null) {
      return Futures.immediateVoidFuture();
    }
    return diskCacheClient.uploadBlob(digest, blob);
  }

  /* must not interact with casUploadCache */
  protected ListenableFuture<Void> uploadOrChunkFile(
      RemoteActionExecutionContext context, Digest digest, Path file, boolean force) {
    return Futures.transformAsync(
        chunking,
        chunking -> {
          if (chunking.supported()
              && digest.getSizeBytes() > chunking.config().chunkingThreshold()) {
            return uploadChunked(chunking, context, digest, file, force);
          }
          return remoteCacheClient.uploadFile(context, digest, file, force);
        },
        directExecutor());
  }

  private ListenableFuture<Void> uploadFileToRemote(
      RemoteActionExecutionContext context, Digest digest, Path file, boolean force) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }
    if (remoteCacheClient == null) {
      return Futures.immediateVoidFuture();
    }

    return Futures.transformAsync(
        chunking,
        chunking -> {
          if (chunking.supported()
              && digest.getSizeBytes() > chunking.config().chunkingThreshold()) {
            return remoteCacheClient.dedupUpload(
                digest, () -> uploadChunked(chunking, context, digest, file, force), force);
          }
          return remoteCacheClient.uploadFile(context, digest, file, force);
        },
        directExecutor());
  }

  private ListenableFuture<Void> uploadBlobToRemote(
      RemoteActionExecutionContext context, Digest digest, Blob blob, boolean force) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }
    if (remoteCacheClient == null) {
      return Futures.immediateVoidFuture();
    }
    return remoteCacheClient.uploadBlob(context, digest, blob, force);
  }

  /** Upload a local file to the remote and/or disk cache. */
  public ListenableFuture<Void> uploadFile(
      RemoteActionExecutionContext context, Digest digest, Path file) {
    return uploadFile(context, digest, file, /* force= */ false);
  }

  /**
   * Upload a local file to the remote and/or disk cache.
   *
   * <p>Trying to upload the same file multiple times concurrently, results in only one upload being
   * performed. An upload to the remote cache is also skipped if this instance has already completed
   * an upload of the same blob, unless {@code force} is set.
   *
   * @param context the context for the action.
   * @param digest the digest of the file.
   * @param file the file to upload.
   * @param force whether to upload to the remote cache even if this instance has already completed
   *     an upload of the same blob, e.g. because it may have been evicted since. Concurrent uploads
   *     are still deduplicated.
   */
  public ListenableFuture<Void> uploadFile(
      RemoteActionExecutionContext context, Digest digest, Path file, boolean force) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }

    ListenableFuture<Void> diskCacheFuture = Futures.immediateVoidFuture();
    if (diskCacheClient != null && context.getWriteCachePolicy().allowDiskCache()) {
      diskCacheFuture = uploadFileToDisk(digest, file);
    }

    ListenableFuture<Void> remoteCacheFuture = Futures.immediateVoidFuture();
    if (remoteCacheClient != null && context.getWriteCachePolicy().allowRemoteCache()) {
      remoteCacheFuture = uploadFileToRemote(context, digest, file, force);
    }

    return Futures.whenAllSucceed(diskCacheFuture, remoteCacheFuture)
        .call(() -> null, directExecutor());
  }

  private ListenableFuture<Void> uploadChunked(
      Chunking chunking,
      RemoteActionExecutionContext context,
      Digest digest,
      Path file,
      boolean force) {
    return virtualThreadExecutor.submit(
        () -> {
          chunking.uploader().uploadChunked(context, digest, file, force);
          return null;
        });
  }

  /**
   * Uploads a blob to the cache from a repeatable stream supplier.
   *
   * <p>The supplier may be opened more than once, including concurrently when both disk and remote
   * cache writes are enabled.
   *
   * @param force whether to upload to the remote cache even if this instance has already completed
   *     an upload of the same blob, e.g. because it may have been evicted since. Concurrent uploads
   *     are still deduplicated.
   */
  public ListenableFuture<Void> uploadBlob(
      RemoteActionExecutionContext context, Digest digest, Blob blob, boolean force) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }

    ListenableFuture<Void> diskCacheFuture = Futures.immediateVoidFuture();
    if (diskCacheClient != null && context.getWriteCachePolicy().allowDiskCache()) {
      diskCacheFuture = uploadBlobToDisk(digest, blob);
    }

    ListenableFuture<Void> remoteCacheFuture = Futures.immediateVoidFuture();
    if (remoteCacheClient != null && context.getWriteCachePolicy().allowRemoteCache()) {
      remoteCacheFuture = uploadBlobToRemote(context, digest, blob, force);
    }

    return Futures.whenAllSucceed(diskCacheFuture, remoteCacheFuture)
        .call(() -> null, directExecutor());
  }

  // Only for use by tests and the remote executor implementation.
  public ListenableFuture<byte[]> downloadBlob(
      RemoteActionExecutionContext context, Digest digest) {
    return downloadBlob(context, /* blobName= */ "", /* execPath= */ null, digest);
  }

  /**
   * Downloads a blob with content hash {@code digest} and stores its content in memory.
   *
   * @return a future that completes after the download completes (succeeds / fails). If successful,
   *     the content is stored in the future's {@code byte[]}.
   */
  public ListenableFuture<byte[]> downloadBlob(
      RemoteActionExecutionContext context,
      String blobName,
      @Nullable PathFragment execPath,
      Digest digest) {
    if (digest.getSizeBytes() == 0) {
      return EMPTY_BYTES;
    }
    ByteArrayOutputStream bOut = new ByteArrayOutputStream((int) digest.getSizeBytes());
    var download = downloadBlob(context, blobName, execPath, digest, bOut);
    return Futures.transform(download, (v) -> bOut.toByteArray(), directExecutor());
  }

  /**
   * Downloads a blob with content hash {@code digest} and stores its content in memory.
   *
   * <p>Unlike {@link #downloadBlob(RemoteActionExecutionContext, String, PathFragment, Digest)},
   * the content is returned as a {@link ByteString} without an extra copy of the downloaded bytes.
   *
   * @return a future that completes after the download completes (succeeds / fails). If successful,
   *     the content is stored in the future's {@link ByteString}.
   */
  public ListenableFuture<ByteString> downloadBlobAsByteString(
      RemoteActionExecutionContext context,
      String blobName,
      @Nullable PathFragment execPath,
      Digest digest) {
    if (digest.getSizeBytes() == 0) {
      return immediateFuture(ByteString.empty());
    }
    ByteString.Output bOut = ByteString.newOutput((int) digest.getSizeBytes());
    var download = downloadBlob(context, blobName, execPath, digest, bOut);
    return Futures.transform(download, (v) -> bOut.toByteString(), directExecutor());
  }

  private ListenableFuture<Void> downloadBlob(
      RemoteActionExecutionContext context,
      String blobName,
      @Nullable PathFragment execPath,
      Digest digest,
      OutputStream out) {
    if (digest.getSizeBytes() == 0) {
      return COMPLETED_SUCCESS;
    }
    var future = downloadBlob(context, digest, out);
    return Futures.catchingAsync(
        future,
        CacheNotFoundException.class,
        (cacheNotFoundException) -> {
          cacheNotFoundException.setExecPath(execPath);
          cacheNotFoundException.setFilename(blobName);
          return immediateFailedFuture(cacheNotFoundException);
        },
        directExecutor());
  }

  ListenableFuture<Void> downloadBlob(
      RemoteActionExecutionContext context, Digest digest, OutputStream out) {
    ListenableFuture<Void> future = immediateFailedFuture(new CacheNotFoundException(digest));

    if (diskCacheClient != null && context.getReadCachePolicy().allowDiskCache()) {
      future = diskCacheClient.downloadBlob(digest, out);
    }

    if (remoteCacheClient != null && context.getReadCachePolicy().allowRemoteCache()) {
      future =
          Futures.catchingAsync(
              future,
              CacheNotFoundException.class,
              (unused) -> downloadBlobFromRemote(context, digest, out),
              directExecutor());
    }

    return future;
  }

  private ListenableFuture<Void> downloadBlobFromRemote(
      RemoteActionExecutionContext context, Digest digest, OutputStream out) {
    return Futures.transformAsync(
        chunking,
        chunking -> {
          if (chunking.supported()
              && digest.getSizeBytes() > chunking.config().chunkingThreshold()) {
            ListenableFuture<Void> chunkedDownloadFuture =
                virtualThreadExecutor.submit(
                    () -> {
                      chunking.downloader().downloadChunked(context, digest, out);
                      return null;
                    });
            // Only a failure to split the blob is recoverable by downloading it as a whole. Once
            // the
            // server has described the blob as a sequence of chunks, the download is committed to
            // that
            // path: chunks are written to `out` as they arrive, so restarting would append the blob
            // to
            // the chunks already written. A missing chunk is reported as a missing blob instead,
            // letting
            // the usual lost input handling regenerate it.
            return Futures.catchingAsync(
                chunkedDownloadFuture,
                BlobNotSplittableException.class,
                (e) -> regularDownloadBlobFromRemote(context, digest, out),
                directExecutor());
          }

          return regularDownloadBlobFromRemote(context, digest, out);
        },
        directExecutor());
  }

  private ListenableFuture<Void> regularDownloadBlobFromRemote(
      RemoteActionExecutionContext context, Digest digest, OutputStream out) {
    checkState(remoteCacheClient != null && context.getReadCachePolicy().allowRemoteCache());

    if (diskCacheClient != null && context.getWriteCachePolicy().allowDiskCache()) {
      Path tempPath = diskCacheClient.getTempPath();
      LazyFileOutputStream tempOut = new LazyFileOutputStream(tempPath);
      ListenableFuture<Void> download = remoteCacheClient.downloadBlob(context, digest, tempOut);
      return cleanupTempFileOnError(
          Futures.transformAsync(
              download,
              (unused) -> {
                try {
                  // Fsync temp before we rename it to avoid data loss in the case of machine
                  // crashes (the OS may reorder the writes and the rename).
                  tempOut.syncIfPossible();
                  tempOut.close();
                  diskCacheClient.captureFile(tempPath, digest, Store.CAS);
                } catch (IOException e) {
                  return immediateFailedFuture(e);
                }
                // The garbage collection of another server sharing the disk cache may have removed
                // the blob again, in which case it is served from the remote cache directly.
                return Futures.catchingAsync(
                    diskCacheClient.downloadBlob(digest, out),
                    CacheNotFoundException.class,
                    e -> remoteCacheClient.downloadBlob(context, digest, out),
                    directExecutor());
              },
              directExecutor()),
          tempPath,
          tempOut);
    }

    return remoteCacheClient.downloadBlob(context, digest, out);
  }

  private static ListenableFuture<Void> cleanupTempFileOnError(
      ListenableFuture<Void> f, Path tempPath, OutputStream tempOut) {
    return Futures.catchingAsync(
        f,
        Exception.class,
        (rootCause) -> {
          try {
            tempOut.close();
          } catch (IOException e) {
            rootCause.addSuppressed(e);
          }
          try {
            tempPath.delete();
          } catch (IOException e) {
            rootCause.addSuppressed(e);
          }
          return immediateFailedFuture(rootCause);
        },
        directExecutor());
  }

  /** A reporter that reports download progresses. */
  public static class DownloadProgressReporter {
    private static final Pattern PATTERN = Pattern.compile("^bazel-out/[^/]+/[^/]+/");
    private final boolean includeFile;
    private final ProgressStatusListener listener;
    private final String id;
    private final String file;
    private final String totalSize;
    private final AtomicLong downloadedBytes = new AtomicLong(0);

    public DownloadProgressReporter(ProgressStatusListener listener, String file, long totalSize) {
      this(/* includeFile= */ true, listener, file, totalSize);
    }

    public DownloadProgressReporter(
        boolean includeFile, ProgressStatusListener listener, String file, long totalSize) {
      this.includeFile = includeFile;
      this.listener = listener;
      this.id = file;
      this.totalSize = bytesCountToDisplayString(totalSize);

      Matcher matcher = PATTERN.matcher(file);
      this.file = matcher.replaceFirst("");
    }

    void started() {
      reportProgress(false, false);
    }

    void downloadedBytes(int count) {
      downloadedBytes.addAndGet(count);
      reportProgress(true, false);
    }

    void finished() {
      reportProgress(true, true);
    }

    private void reportProgress(boolean includeBytes, boolean finished) {
      String progress;
      if (includeBytes) {
        if (includeFile) {
          progress =
              String.format(
                  "Downloading %s, %s / %s",
                  file, bytesCountToDisplayString(downloadedBytes.get()), totalSize);
        } else {
          progress =
              String.format("%s / %s", bytesCountToDisplayString(downloadedBytes.get()), totalSize);
        }
      } else {
        if (includeFile) {
          progress = String.format("Downloading %s", file);
        } else {
          progress = "";
        }
      }
      listener.onProgressStatus(SpawnProgressEvent.create(id, progress, finished));
    }
  }

  /**
   * Downloads a file (that is not a directory) into {@code localPath}.
   *
   * @param localPath the path to write the file to
   * @param finalPath the path the content will remain at after the command, e.g. when {@code
   *     localPath} is a temporary staging path, or {@code null} if that is {@code localPath}
   *     itself. Used to remember where chunks of the file can be found so that later chunked
   *     downloads can read them from local disk instead of the remote cache.
   */
  public ListenableFuture<Void> downloadFile(
      RemoteActionExecutionContext context,
      String outputPath,
      @Nullable PathFragment execPath,
      Path localPath,
      @Nullable Path finalPath,
      Digest digest,
      DownloadProgressReporter reporter)
      throws IOException {
    ListenableFuture<Void> f = downloadFile(context, localPath, finalPath, digest, reporter);
    return Futures.catchingAsync(
        f,
        Throwable.class,
        (throwable) -> {
          if (throwable instanceof CacheNotFoundException cacheNotFoundException) {
            cacheNotFoundException.setExecPath(execPath);
            cacheNotFoundException.setFilename(outputPath);
          } else if (throwable instanceof OutputDigestMismatchException e) {
            e.setOutputPath(outputPath);
            e.setLocalPath(localPath);
          }
          return immediateFailedFuture(throwable);
        },
        directExecutor());
  }

  /**
   * Downloads a file (that is not a directory). The content is fetched from the digest.
   *
   * <p>Use {@link #downloadFile(RemoteActionExecutionContext, String, PathFragment, Path, Path,
   * Digest, DownloadProgressReporter)} instead for build outputs as it provides progress
   * information and correctly handles unexpected cache misses.
   */
  public ListenableFuture<Void> downloadFile(
      RemoteActionExecutionContext context, Path path, Digest digest) throws IOException {
    return downloadFile(
        context,
        path.getPathString(),
        /* execPath= */ null,
        path,
        /* finalPath= */ null,
        digest,
        new DownloadProgressReporter(NO_ACTION, "", 0));
  }

  /** Downloads a file (that is not a directory). The content is fetched from the digest. */
  private ListenableFuture<Void> downloadFile(
      RemoteActionExecutionContext context,
      Path path,
      @Nullable Path finalPath,
      Digest digest,
      DownloadProgressReporter reporter)
      throws IOException {
    checkNotNull(path.getParentDirectory()).createDirectoryAndParents();
    if (digest.getSizeBytes() == 0) {
      // Handle empty file locally.
      FileSystemUtils.writeContent(path, new byte[0]);
      return COMPLETED_SUCCESS;
    }

    if (symlinkTemplate != null) {
      // Don't actually download files from the CAS. Instead, create a
      // symbolic link that points to a location where CAS objects may
      // be found. This could, for example, be a FUSE file system.
      path.createSymbolicLink(
          path.getRelative(
              symlinkTemplate
                  .replace("{hash}", digest.getHash())
                  .replace("{size_bytes}", String.valueOf(digest.getSizeBytes()))));
      return COMPLETED_SUCCESS;
    }

    reporter.started();
    OutputStream out =
        new ReportingOutputStream(
            new LazyFileOutputStream(path), reporter, firstNonNull(finalPath, path));

    ListenableFuture<Void> f = downloadBlob(context, digest, out);
    f.addListener(
        () -> {
          try {
            out.close();
          } catch (IOException e) {
            logger.atWarning().withCause(e).log(
                "Unexpected exception closing output stream after downloading %s/%d to %s",
                digest.getHash(), digest.getSizeBytes(), path);
          } finally {
            reporter.finished();
          }
        },
        directExecutor());
    return f;
  }

  /**
   * Download the stdout and stderr of an executed action.
   *
   * @param context the context for the action.
   * @param result the result of the action.
   * @param outErr the {@link OutErr} that the stdout and stderr will be downloaded to.
   */
  public final List<ListenableFuture<Void>> downloadOutErr(
      RemoteActionExecutionContext context, ActionResult result, OutErr outErr) {
    return downloadOutErr(context, result, outErr, /* downloadStdout= */ true);
  }

  /**
   * Download the stdout and stderr of an executed action.
   *
   * @param context the context for the action.
   * @param result the result of the action.
   * @param outErr the {@link OutErr} that the stdout and stderr will be downloaded to.
   * @param downloadStdout whether to download stdout.
   */
  public final List<ListenableFuture<Void>> downloadOutErr(
      RemoteActionExecutionContext context,
      ActionResult result,
      OutErr outErr,
      boolean downloadStdout) {
    List<ListenableFuture<Void>> downloads = new ArrayList<>();
    if (downloadStdout) {
      if (!result.getStdoutRaw().isEmpty()) {
        try {
          result.getStdoutRaw().writeTo(outErr.getOutputStream());
          outErr.getOutputStream().flush();
        } catch (IOException e) {
          downloads.add(Futures.immediateFailedFuture(e));
        }
      } else if (result.hasStdoutDigest()) {
        downloads.add(
            downloadBlob(
                context,
                /* blobName= */ "<stdout>",
                /* execPath= */ null,
                result.getStdoutDigest(),
                outErr.getOutputStream()));
      }
    }
    if (!result.getStderrRaw().isEmpty()) {
      try {
        result.getStderrRaw().writeTo(outErr.getErrorStream());
        outErr.getErrorStream().flush();
      } catch (IOException e) {
        downloads.add(Futures.immediateFailedFuture(e));
      }
    } else if (result.hasStderrDigest()) {
      downloads.add(
          downloadBlob(
              context,
              /* blobName= */ "<stderr>",
              /* execPath= */ null,
              result.getStderrDigest(),
              outErr.getErrorStream()));
    }
    return downloads;
  }

  public boolean hasRemoteCache() {
    return remoteCacheClient != null;
  }

  public boolean hasDiskCache() {
    return diskCacheClient != null;
  }

  @Override
  protected void deallocate() {
    if (diskCacheClient != null) {
      diskCacheClient.close();
    }
    virtualThreadExecutor.shutdown();
    if (remoteCacheClient != null) {
      remoteCacheClient.shutdownUploads();
      remoteCacheClient.close();
    }

    closeCountDownLatch.countDown();
  }

  @Override
  public CombinedCache touch(Object o) {
    return this;
  }

  @CanIgnoreReturnValue
  @Override
  public CombinedCache retain() {
    super.retain();
    return this;
  }

  /** Waits for active network I/Os to finish. */
  public void awaitTermination() throws InterruptedException {
    if (remoteCacheClient != null) {
      remoteCacheClient.awaitUploadTermination();
    }
    closeCountDownLatch.await();
    virtualThreadExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
  }

  /** Shuts the cache down and cancels active network I/Os. */
  public void shutdownNow() {
    if (remoteCacheClient != null) {
      remoteCacheClient.shutdownUploadsNow();
    }
    virtualThreadExecutor.shutdownNow();
  }

  public static FailureDetail createFailureDetail(String message, Code detailedCode) {
    return FailureDetail.newBuilder()
        .setMessage(message)
        .setRemoteExecution(RemoteExecution.newBuilder().setCode(detailedCode))
        .build();
  }

  /**
   * An {@link OutputStream} that reports all the write operations with {@link
   * DownloadProgressReporter}.
   */
  private static class ReportingOutputStream extends OutputStream implements MaybePathBacked {

    private final OutputStream out;
    private final DownloadProgressReporter reporter;
    private final Path finalPath;

    ReportingOutputStream(OutputStream out, DownloadProgressReporter reporter, Path finalPath) {
      this.out = out;
      this.reporter = reporter;
      this.finalPath = finalPath;
    }

    @Override
    public void write(byte[] b) throws IOException {
      out.write(b);
      reporter.downloadedBytes(b.length);
    }

    @Override
    public void write(byte[] b, int off, int len) throws IOException {
      out.write(b, off, len);
      reporter.downloadedBytes(len);
    }

    @Override
    public void write(int b) throws IOException {
      out.write(b);
      reporter.downloadedBytes(1);
    }

    @Override
    public void flush() throws IOException {
      out.flush();
    }

    @Override
    public void close() throws IOException {
      out.close();
    }

    @Nullable
    @Override
    public Path maybeGetPath() {
      return out instanceof MaybePathBacked maybePathBacked ? maybePathBacked.maybeGetPath() : null;
    }

    @Override
    public Path maybeGetFinalPath() {
      return finalPath;
    }
  }
}