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

// Copyright 2016 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 com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionInputFileCache;
import com.google.devtools.build.lib.analysis.config.InvalidConfigurationException;
import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
import com.google.devtools.build.lib.remote.ContentDigests.ActionKey;
import com.google.devtools.build.lib.remote.RemoteProtocol.ActionResult;
import com.google.devtools.build.lib.remote.RemoteProtocol.BlobChunk;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasDownloadBlobRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasDownloadReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasLookupRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasStatus;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadBlobReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadBlobRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadTreeMetadataReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadTreeMetadataRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ContentDigest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionCacheReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionCacheRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionCacheSetReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionCacheSetRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionCacheStatus;
import com.google.devtools.build.lib.remote.RemoteProtocol.FileMetadata;
import com.google.devtools.build.lib.remote.RemoteProtocol.FileNode;
import com.google.devtools.build.lib.remote.RemoteProtocol.Output;
import com.google.devtools.build.lib.remote.RemoteProtocol.Output.ContentCase;
import com.google.devtools.build.lib.remote.TreeNodeRepository.TreeNode;
import com.google.devtools.build.lib.util.Pair;
import com.google.devtools.build.lib.util.Preconditions;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.protobuf.ByteString;
import io.grpc.ManagedChannel;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/** A RemoteActionCache implementation that uses gRPC calls to a remote cache server. */
@ThreadSafe
public class GrpcActionCache implements RemoteActionCache {
  private static final int MAX_MEMORY_KBYTES = 512 * 1024;

  /** Channel over which to send gRPC CAS queries. */
  private final GrpcCasInterface casIface;
  private final GrpcExecutionCacheInterface iface;
  private final RemoteOptions options;

  /** Reads from multiple sequential inputs and chunks the data into BlobChunks. */
  static interface BlobChunkIterator {
    boolean hasNext();

    BlobChunk next() throws IOException; // IOException can be a result of file read.
  }

  final class BlobChunkInlineIterator implements BlobChunkIterator {
    private final Iterator<byte[]> blobIterator;
    private final Set<ContentDigest> digests;
    private int offset;
    private ContentDigest digest;
    private byte[] currentBlob;

    public BlobChunkInlineIterator(Set<ContentDigest> digests, Iterator<byte[]> blobIterator) {
      this.digests = digests;
      this.blobIterator = blobIterator;
      advanceInput();
    }

    public BlobChunkInlineIterator(byte[] blob) {
      blobIterator = null;
      offset = 0;
      currentBlob = blob;
      digest = ContentDigests.computeDigest(currentBlob);
      digests = null;
    }

    private void advanceInput() {
      offset = 0;
      do {
        if (blobIterator != null && blobIterator.hasNext()) {
          currentBlob = blobIterator.next();
          digest = ContentDigests.computeDigest(currentBlob);
        } else {
          currentBlob = null;
          digest = null;
        }
      } while (digest != null && !digests.contains(digest));
    }

    @Override
    public boolean hasNext() {
      return currentBlob != null;
    }

    @Override
    public BlobChunk next() throws IOException {
      if (!hasNext()) {
        throw new NoSuchElementException();
      }
      BlobChunk.Builder chunk = BlobChunk.newBuilder();
      if (offset == 0) {
        chunk.setDigest(digest);
      } else {
        chunk.setOffset(offset);
      }
      int size = Math.min(currentBlob.length - offset, options.grpcMaxChunkSizeBytes);
      if (size > 0) {
        chunk.setData(ByteString.copyFrom(currentBlob, offset, size));
        offset += size;
      }
      if (offset >= currentBlob.length) {
        advanceInput();
      }
      return chunk.build();
    }
  }

  final class BlobChunkFileIterator implements BlobChunkIterator {
    private final Iterator<Path> fileIterator;
    private InputStream currentStream;
    private final Set<ContentDigest> digests;
    private ContentDigest digest;
    private long bytesLeft;

    public BlobChunkFileIterator(Set<ContentDigest> digests, Iterator<Path> fileIterator)
        throws IOException {
      this.digests = digests;
      this.fileIterator = fileIterator;
      advanceInput();
    }

    public BlobChunkFileIterator(Path file) throws IOException {
      fileIterator = Iterators.singletonIterator(file);
      digests = ImmutableSet.of(ContentDigests.computeDigest(file));
      advanceInput();
    }

    private void advanceInput() throws IOException {
      do {
        if (fileIterator != null && fileIterator.hasNext()) {
          Path file = fileIterator.next();
          digest = ContentDigests.computeDigest(file);
          currentStream = file.getInputStream();
          bytesLeft = digest.getSizeBytes();
        } else {
          digest = null;
          currentStream = null;
          bytesLeft = 0;
        }
      } while (digest != null && !digests.contains(digest));
    }

    @Override
    public boolean hasNext() {
      return currentStream != null;
    }

    @Override
    public BlobChunk next() throws IOException {
      if (!hasNext()) {
        throw new NoSuchElementException();
      }
      BlobChunk.Builder chunk = BlobChunk.newBuilder();
      long offset = digest.getSizeBytes() - bytesLeft;
      if (offset == 0) {
        chunk.setDigest(digest);
      } else {
        chunk.setOffset(offset);
      }
      if (bytesLeft > 0) {
        byte[] blob = new byte[(int) Math.min(bytesLeft, options.grpcMaxChunkSizeBytes)];
        currentStream.read(blob);
        chunk.setData(ByteString.copyFrom(blob));
        bytesLeft -= blob.length;
      }
      if (bytesLeft == 0) {
        currentStream.close();
        advanceInput();
      }
      return chunk.build();
    }
  }

  final class BlobChunkActionInputIterator implements BlobChunkIterator {
    private final Iterator<? extends ActionInput> inputIterator;
    private final ActionInputFileCache inputCache;
    private final Path execRoot;
    private InputStream currentStream;
    private final Set<ContentDigest> digests;
    private ContentDigest digest;
    private long bytesLeft;

    public BlobChunkActionInputIterator(
        Set<ContentDigest> digests,
        Path execRoot,
        Iterator<? extends ActionInput> inputIterator,
        ActionInputFileCache inputCache)
        throws IOException {
      this.digests = digests;
      this.inputIterator = inputIterator;
      this.inputCache = inputCache;
      this.execRoot = execRoot;
      advanceInput();
    }

    public BlobChunkActionInputIterator(
        ActionInput input, Path execRoot, ActionInputFileCache inputCache) throws IOException {
      inputIterator = Iterators.singletonIterator(input);
      digests = ImmutableSet.of(ContentDigests.getDigestFromInputCache(input, inputCache));
      this.inputCache = inputCache;
      this.execRoot = execRoot;
      advanceInput();
    }

    private void advanceInput() throws IOException {
      do {
        if (inputIterator != null && inputIterator.hasNext()) {
          ActionInput input = inputIterator.next();
          digest = ContentDigests.getDigestFromInputCache(input, inputCache);
          currentStream = execRoot.getRelative(input.getExecPathString()).getInputStream();
          bytesLeft = digest.getSizeBytes();
        } else {
          digest = null;
          currentStream = null;
          bytesLeft = 0;
        }
      } while (digest != null && !digests.contains(digest));
    }

    @Override
    public boolean hasNext() {
      return currentStream != null;
    }

    @Override
    public BlobChunk next() throws IOException {
      if (!hasNext()) {
        throw new NoSuchElementException();
      }
      BlobChunk.Builder chunk = BlobChunk.newBuilder();
      long offset = digest.getSizeBytes() - bytesLeft;
      if (offset == 0) {
        chunk.setDigest(digest);
      } else {
        chunk.setOffset(offset);
      }
      if (bytesLeft > 0) {
        byte[] blob = new byte[(int) Math.min(bytesLeft, (long) options.grpcMaxChunkSizeBytes)];
        currentStream.read(blob);
        chunk.setData(ByteString.copyFrom(blob));
        bytesLeft -= blob.length;
      }
      if (bytesLeft == 0) {
        currentStream.close();
        advanceInput();
      }
      return chunk.build();
    }
  }

  public GrpcActionCache(
      RemoteOptions options, GrpcCasInterface casIface, GrpcExecutionCacheInterface iface) {
    this.options = options;
    this.casIface = casIface;
    this.iface = iface;
  }

  @VisibleForTesting
  public GrpcActionCache(ManagedChannel channel, RemoteOptions options) {
    this.options = options;
    this.casIface = GrpcInterfaces.casInterface(options.grpcTimeoutSeconds, channel);
    this.iface = GrpcInterfaces.executionCacheInterface(options.grpcTimeoutSeconds, channel);
  }

  public GrpcActionCache(RemoteOptions options) throws InvalidConfigurationException {
    this(RemoteUtils.createChannelLegacy(options.remoteCache), options);
  }

  public static boolean isRemoteCacheOptions(RemoteOptions options) {
    return options.remoteCache != null;
  }

  private ImmutableSet<ContentDigest> getMissingDigests(Iterable<ContentDigest> digests) {
    CasLookupRequest.Builder request = CasLookupRequest.newBuilder().addAllDigest(digests);
    if (request.getDigestCount() == 0) {
      return ImmutableSet.of();
    }
    CasStatus status = casIface.lookup(request.build()).getStatus();
    if (!status.getSucceeded() && status.getError() != CasStatus.ErrorCode.MISSING_DIGEST) {
      // TODO(olaola): here and below, add basic retry logic on transient errors!
      throw new RuntimeException(status.getErrorDetail());
    }
    return ImmutableSet.copyOf(status.getMissingDigestList());
  }

  /**
   * Upload enough of the tree metadata and data into remote cache so that the entire tree can be
   * reassembled remotely using the root digest.
   */
  @Override
  public void uploadTree(TreeNodeRepository repository, Path execRoot, TreeNode root)
      throws IOException, InterruptedException {
    repository.computeMerkleDigests(root);
    // TODO(olaola): avoid querying all the digests, only ask for novel subtrees.
    ImmutableSet<ContentDigest> missingDigests = getMissingDigests(repository.getAllDigests(root));

    // Only upload data that was missing from the cache.
    ArrayList<ActionInput> actionInputs = new ArrayList<>();
    ArrayList<FileNode> treeNodes = new ArrayList<>();
    repository.getDataFromDigests(missingDigests, actionInputs, treeNodes);

    if (!treeNodes.isEmpty()) {
      CasUploadTreeMetadataRequest.Builder metaRequest =
          CasUploadTreeMetadataRequest.newBuilder().addAllTreeNode(treeNodes);
      CasUploadTreeMetadataReply reply = casIface.uploadTreeMetadata(metaRequest.build());
      if (!reply.getStatus().getSucceeded()) {
        throw new RuntimeException(reply.getStatus().getErrorDetail());
      }
    }
    if (!actionInputs.isEmpty()) {
      uploadChunks(
          actionInputs.size(),
          new BlobChunkActionInputIterator(
              missingDigests, execRoot, actionInputs.iterator(), repository.getInputFileCache()));
    }
  }

  /**
   * Download the entire tree data rooted by the given digest and write it into the given location.
   */
  @Override
  public void downloadTree(ContentDigest rootDigest, Path rootLocation)
      throws IOException, CacheNotFoundException {
    throw new UnsupportedOperationException();
  }

  private void handleDownloadStatus(CasStatus status) throws CacheNotFoundException {
    if (!status.getSucceeded()) {
      if (status.getError() == CasStatus.ErrorCode.MISSING_DIGEST) {
        throw new CacheNotFoundException(status.getMissingDigest(0));
      }
      // TODO(olaola): deal with other statuses better.
      throw new RuntimeException(status.getErrorDetail());
    }
  }

  /**
   * Download all results of a remotely executed action locally. TODO(olaola): will need to amend to
   * include the {@link com.google.devtools.build.lib.remote.TreeNodeRepository} for updating.
   */
  @Override
  public void downloadAllResults(ActionResult result, Path execRoot)
      throws IOException, CacheNotFoundException {
    if (result.getOutputList().isEmpty()) {
      return;
    }
    // Send all the file requests in a single synchronous batch.
    // TODO(olaola): profile to maybe replace with separate concurrent requests.
    CasDownloadBlobRequest.Builder request = CasDownloadBlobRequest.newBuilder();
    Map<ContentDigest, Pair<Path, FileMetadata>> metadataMap = new HashMap<>();
    for (Output output : result.getOutputList()) {
      Path path = execRoot.getRelative(output.getPath());
      if (output.getContentCase() == ContentCase.FILE_METADATA) {
        FileMetadata fileMetadata = output.getFileMetadata();
        ContentDigest digest = fileMetadata.getDigest();
        if (digest.getSizeBytes() > 0) {
          request.addDigest(digest);
          metadataMap.put(digest, Pair.of(path, fileMetadata));
        } else {
          // Handle empty file locally.
          FileSystemUtils.createDirectoryAndParents(path.getParentDirectory());
          FileSystemUtils.writeContent(path, new byte[0]);
        }
      } else {
        downloadTree(output.getDigest(), path);
      }
    }
    Iterator<CasDownloadReply> replies = casIface.downloadBlob(request.build());
    Set<ContentDigest> results = new HashSet<>();
    while (replies.hasNext()) {
      results.add(createFileFromStream(metadataMap, replies));
    }
    for (ContentDigest digest : metadataMap.keySet()) {
      if (!results.contains(digest)) {
        throw new CacheNotFoundException(digest);
      }
    }
  }

  private ContentDigest createFileFromStream(
      Map<ContentDigest, Pair<Path, FileMetadata>> metadataMap, Iterator<CasDownloadReply> replies)
      throws IOException, CacheNotFoundException {
    Preconditions.checkArgument(replies.hasNext());
    CasDownloadReply reply = replies.next();
    if (reply.hasStatus()) {
      handleDownloadStatus(reply.getStatus());
    }
    BlobChunk chunk = reply.getData();
    ContentDigest digest = chunk.getDigest();
    Preconditions.checkArgument(metadataMap.containsKey(digest));
    Pair<Path, FileMetadata> metadata = metadataMap.get(digest);
    Path path = metadata.first;
    FileSystemUtils.createDirectoryAndParents(path.getParentDirectory());
    try (OutputStream stream = path.getOutputStream()) {
      ByteString data = chunk.getData();
      data.writeTo(stream);
      long bytesLeft = digest.getSizeBytes() - data.size();
      while (bytesLeft > 0) {
        Preconditions.checkArgument(replies.hasNext());
        reply = replies.next();
        if (reply.hasStatus()) {
          handleDownloadStatus(reply.getStatus());
        }
        chunk = reply.getData();
        data = chunk.getData();
        Preconditions.checkArgument(!chunk.hasDigest());
        Preconditions.checkArgument(chunk.getOffset() == digest.getSizeBytes() - bytesLeft);
        data.writeTo(stream);
        bytesLeft -= data.size();
      }
      path.setExecutable(metadata.second.getExecutable());
    }
    return digest;
  }

  /** Upload all results of a locally executed action to the cache. */
  @Override
  public void uploadAllResults(Path execRoot, Collection<Path> files, ActionResult.Builder result)
      throws IOException, InterruptedException {
    ArrayList<ContentDigest> digests = new ArrayList<>();
    for (Path file : files) {
      digests.add(ContentDigests.computeDigest(file));
    }
    ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
    if (!missing.isEmpty()) {
      uploadChunks(missing.size(), new BlobChunkFileIterator(missing, files.iterator()));
    }
    int index = 0;
    for (Path file : files) {
      if (file.isDirectory()) {
        // TODO(olaola): to implement this for a directory, will need to create or pass a
        // TreeNodeRepository to call uploadTree.
        throw new UnsupportedOperationException("Storing a directory is not yet supported.");
      }
      // Add to protobuf.
      result
          .addOutputBuilder()
          .setPath(file.relativeTo(execRoot).getPathString())
          .getFileMetadataBuilder()
          .setDigest(digests.get(index++))
          .setExecutable(file.isExecutable());
    }
  }

  /**
   * Put the file contents cache if it is not already in it. No-op if the file is already stored in
   * cache. The given path must be a full absolute path.
   *
   * @return The key for fetching the file contents blob from cache.
   */
  @Override
  public ContentDigest uploadFileContents(Path file) throws IOException, InterruptedException {
    ContentDigest digest = ContentDigests.computeDigest(file);
    ImmutableSet<ContentDigest> missing = getMissingDigests(ImmutableList.of(digest));
    if (!missing.isEmpty()) {
      uploadChunks(1, new BlobChunkFileIterator(file));
    }
    return digest;
  }

  /**
   * Put the file contents cache if it is not already in it. No-op if the file is already stored in
   * cache. The given path must be a full absolute path.
   *
   * @return The key for fetching the file contents blob from cache.
   */
  @Override
  public ContentDigest uploadFileContents(
      ActionInput input, Path execRoot, ActionInputFileCache inputCache)
      throws IOException, InterruptedException {
    ContentDigest digest = ContentDigests.getDigestFromInputCache(input, inputCache);
    ImmutableSet<ContentDigest> missing = getMissingDigests(ImmutableList.of(digest));
    if (!missing.isEmpty()) {
      uploadChunks(1, new BlobChunkActionInputIterator(input, execRoot, inputCache));
    }
    return digest;
  }

  static class UploadBlobReplyStreamObserver implements StreamObserver<CasUploadBlobReply> {
    private final CountDownLatch finishLatch;
    private final AtomicReference<RuntimeException> exception;

    public UploadBlobReplyStreamObserver(
        CountDownLatch finishLatch, AtomicReference<RuntimeException> exception) {
      this.finishLatch = finishLatch;
      this.exception = exception;
    }

    @Override
    public void onNext(CasUploadBlobReply reply) {
      if (!reply.getStatus().getSucceeded()) {
        // TODO(olaola): add basic retry logic on transient errors!
        this.exception.compareAndSet(
            null, new RuntimeException(reply.getStatus().getErrorDetail()));
      }
    }

    @Override
    public void onError(Throwable t) {
      this.exception.compareAndSet(null, new StatusRuntimeException(Status.fromThrowable(t)));
      finishLatch.countDown();
    }

    @Override
    public void onCompleted() {
      finishLatch.countDown();
    }
  }

  private void uploadChunks(int numItems, BlobChunkIterator blobs)
      throws InterruptedException, IOException {
    CountDownLatch finishLatch = new CountDownLatch(numItems); // Maximal number of batches.
    AtomicReference<RuntimeException> exception = new AtomicReference<>(null);
    UploadBlobReplyStreamObserver responseObserver = null;
    StreamObserver<CasUploadBlobRequest> requestObserver = null;
    int currentBatchBytes = 0;
    int batchedInputs = 0;
    int batches = 0;
    try {
      while (blobs.hasNext()) {
        BlobChunk chunk = blobs.next();
        if (chunk.hasDigest()) {
          // Determine whether to start next batch.
          final long batchSize = chunk.getDigest().getSizeBytes() + currentBatchBytes;
          if (batchedInputs % options.grpcMaxBatchInputs == 0
              || batchSize > options.grpcMaxBatchSizeBytes) {
            // The batches execute simultaneously.
            if (requestObserver != null) {
              batchedInputs = 0;
              currentBatchBytes = 0;
              requestObserver.onCompleted();
            }
            batches++;
            responseObserver = new UploadBlobReplyStreamObserver(finishLatch, exception);
            requestObserver = casIface.uploadBlobAsync(responseObserver);
          }
          batchedInputs++;
        }
        currentBatchBytes += chunk.getData().size();
        requestObserver.onNext(CasUploadBlobRequest.newBuilder().setData(chunk).build());
        if (finishLatch.getCount() == 0) {
          // RPC completed or errored before we finished sending.
          throw new RuntimeException(
              "gRPC terminated prematurely: "
                  + (exception.get() != null ? exception.get() : "unknown cause"));
        }
      }
    } catch (RuntimeException e) {
      // Cancel RPC
      if (requestObserver != null) {
        requestObserver.onError(e);
      }
      throw e;
    }
    if (requestObserver != null) {
      requestObserver.onCompleted(); // Finish last batch.
    }
    while (batches++ < numItems) {
      finishLatch.countDown(); // Non-sent batches.
    }
    finishLatch.await(options.grpcTimeoutSeconds, TimeUnit.SECONDS);
    if (exception.get() != null) {
      throw exception.get(); // Re-throw the first encountered exception.
    }
  }

  @Override
  public ImmutableList<ContentDigest> uploadBlobs(Iterable<byte[]> blobs)
      throws InterruptedException {
    ArrayList<ContentDigest> digests = new ArrayList<>();
    for (byte[] blob : blobs) {
      digests.add(ContentDigests.computeDigest(blob));
    }
    ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
    try {
      if (!missing.isEmpty()) {
        uploadChunks(missing.size(), new BlobChunkInlineIterator(missing, blobs.iterator()));
      }
      return ImmutableList.copyOf(digests);
    } catch (IOException e) {
      // This will never happen.
      throw new RuntimeException(e);
    }
  }

  @Override
  public ContentDigest uploadBlob(byte[] blob) throws InterruptedException {
    ContentDigest digest = ContentDigests.computeDigest(blob);
    ImmutableSet<ContentDigest> missing = getMissingDigests(ImmutableList.of(digest));
    try {
      if (!missing.isEmpty()) {
        uploadChunks(1, new BlobChunkInlineIterator(blob));
      }
      return digest;
    } catch (IOException e) {
      // This will never happen.
      throw new RuntimeException();
    }
  }

  @Override
  public byte[] downloadBlob(ContentDigest digest) throws CacheNotFoundException {
    return downloadBlobs(ImmutableList.of(digest)).get(0);
  }

  @Override
  public ImmutableList<byte[]> downloadBlobs(Iterable<ContentDigest> digests)
      throws CacheNotFoundException {
    // Send all the file requests in a single synchronous batch.
    // TODO(olaola): profile to maybe replace with separate concurrent requests.
    CasDownloadBlobRequest.Builder request = CasDownloadBlobRequest.newBuilder();
    for (ContentDigest digest : digests) {
      if (digest.getSizeBytes() > 0) {
        request.addDigest(digest); // We handle empty blobs locally.
      }
    }
    Iterator<CasDownloadReply> replies = null;
    Map<ContentDigest, byte[]> results = new HashMap<>();
    int digestCount = request.getDigestCount();
    if (digestCount > 0) {
      replies = casIface.downloadBlob(request.build());
      while (digestCount-- > 0) {
        Preconditions.checkArgument(replies.hasNext());
        CasDownloadReply reply = replies.next();
        if (reply.hasStatus()) {
          handleDownloadStatus(reply.getStatus());
        }
        BlobChunk chunk = reply.getData();
        ContentDigest digest = chunk.getDigest();
        // This is not enough, but better than nothing.
        Preconditions.checkArgument(digest.getSizeBytes() / 1000.0 < MAX_MEMORY_KBYTES);
        byte[] result = new byte[(int) digest.getSizeBytes()];
        ByteString data = chunk.getData();
        data.copyTo(result, 0);
        int offset = data.size();
        while (offset < result.length) {
          Preconditions.checkArgument(replies.hasNext());
          reply = replies.next();
          if (reply.hasStatus()) {
            handleDownloadStatus(reply.getStatus());
          }
          chunk = reply.getData();
          Preconditions.checkArgument(!chunk.hasDigest());
          Preconditions.checkArgument(chunk.getOffset() == offset);
          data = chunk.getData();
          data.copyTo(result, offset);
          offset += data.size();
        }
        results.put(digest, result);
      }
    }

    ArrayList<byte[]> result = new ArrayList<>();
    for (ContentDigest digest : digests) {
      if (digest.getSizeBytes() == 0) {
        result.add(new byte[0]);
        continue;
      }
      if (!results.containsKey(digest)) {
        throw new CacheNotFoundException(digest);
      }
      result.add(results.get(digest));
    }
    return ImmutableList.copyOf(result);
  }

  // Execution Cache API

  /** Returns a cached result for a given Action digest, or null if not found in cache. */
  @Override
  public ActionResult getCachedActionResult(ActionKey actionKey) {
    ExecutionCacheRequest request =
        ExecutionCacheRequest.newBuilder().setActionDigest(actionKey.getDigest()).build();
    ExecutionCacheReply reply = iface.getCachedResult(request);
    ExecutionCacheStatus status = reply.getStatus();
    if (!status.getSucceeded()
        && status.getError() != ExecutionCacheStatus.ErrorCode.MISSING_RESULT) {
      throw new RuntimeException(status.getErrorDetail());
    }
    return reply.hasResult() ? reply.getResult() : null;
  }

  /** Sets the given result as result of the given Action. */
  @Override
  public void setCachedActionResult(ActionKey actionKey, ActionResult result)
      throws InterruptedException {
    ExecutionCacheSetRequest request =
        ExecutionCacheSetRequest.newBuilder()
            .setActionDigest(actionKey.getDigest())
            .setResult(result)
            .build();
    ExecutionCacheSetReply reply = iface.setCachedResult(request);
    ExecutionCacheStatus status = reply.getStatus();
    if (!status.getSucceeded() && status.getError() != ExecutionCacheStatus.ErrorCode.UNSUPPORTED) {
      throw new RuntimeException(status.getErrorDetail());
    }
  }
}