src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java - bazel - Git at Google

 // 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.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputFileCache;
 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.Channel;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 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.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;

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

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

   public GrpcActionCache(RemoteOptions options, ChannelOptions channelOptions) {
     this(RemoteUtils.createChannel(options.remoteCache, channelOptions), options, channelOptions);
   }

   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 Chunker.Builder()
               .chunkSize(options.grpcMaxChunkSizeBytes)
               .addAllInputs(actionInputs, repository.getInputFileCache(), execRoot)
               .onlyUseDigests(missingDigests)
               .build());
     }
   }

   /**
    * 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<>();
     Chunker.Builder b = new Chunker.Builder().chunkSize(options.grpcMaxChunkSizeBytes);
     for (Path file : files) {
       digests.add(ContentDigests.computeDigest(file));
       b.addInput(file);
     }
     ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
     if (!missing.isEmpty()) {
       uploadChunks(missing.size(), b.onlyUseDigests(missing).build());
     }
     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, Chunker.from(file, options.grpcMaxChunkSizeBytes));
     }
     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, Chunker.from(input, options.grpcMaxChunkSizeBytes, inputCache, execRoot));
     }
     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, Chunker 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<>();
     Chunker.Builder b = new Chunker.Builder().chunkSize(options.grpcMaxChunkSizeBytes);
     for (byte[] blob : blobs) {
       digests.add(ContentDigests.computeDigest(blob));
       b.addInput(blob);
     }
     ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
     try {
       if (!missing.isEmpty()) {
         uploadChunks(missing.size(), b.onlyUseDigests(missing).build());
       }
       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, Chunker.from(blob, options.grpcMaxChunkSizeBytes));
       }
       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());
     }
   }
 }
	// 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.devtools.build.lib.actions.ActionInput;
	import com.google.devtools.build.lib.actions.ActionInputFileCache;
	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.Channel;
	import io.grpc.Status;
	import io.grpc.StatusRuntimeException;
	import io.grpc.stub.StreamObserver;
	import java.io.IOException;
	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.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;

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

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

	public GrpcActionCache(RemoteOptions options, ChannelOptions channelOptions) {
	this(RemoteUtils.createChannel(options.remoteCache, channelOptions), options, channelOptions);
	}

	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 Chunker.Builder()
	.chunkSize(options.grpcMaxChunkSizeBytes)
	.addAllInputs(actionInputs, repository.getInputFileCache(), execRoot)
	.onlyUseDigests(missingDigests)
	.build());
	}
	}

	/**
	* 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<>();
	Chunker.Builder b = new Chunker.Builder().chunkSize(options.grpcMaxChunkSizeBytes);
	for (Path file : files) {
	digests.add(ContentDigests.computeDigest(file));
	b.addInput(file);
	}
	ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
	if (!missing.isEmpty()) {
	uploadChunks(missing.size(), b.onlyUseDigests(missing).build());
	}
	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, Chunker.from(file, options.grpcMaxChunkSizeBytes));
	}
	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, Chunker.from(input, options.grpcMaxChunkSizeBytes, inputCache, execRoot));
	}
	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, Chunker 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<>();
	Chunker.Builder b = new Chunker.Builder().chunkSize(options.grpcMaxChunkSizeBytes);
	for (byte[] blob : blobs) {
	digests.add(ContentDigests.computeDigest(blob));
	b.addInput(blob);
	}
	ImmutableSet<ContentDigest> missing = getMissingDigests(digests);
	try {
	if (!missing.isEmpty()) {
	uploadChunks(missing.size(), b.onlyUseDigests(missing).build());
	}
	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, Chunker.from(blob, options.grpcMaxChunkSizeBytes));
	}
	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());
	}
	}
	}