aboutsummaryrefslogtreecommitdiffhomepage
path: root/src/main/java/com/google/devtools/build/lib/remote/GrpcActionCache.java
blob: 5004748b00aea2746d3816da312e83ad10a60724 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
// 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.analysis.config.InvalidConfigurationException;
import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
import com.google.devtools.build.lib.remote.CasServiceGrpc.CasServiceBlockingStub;
import com.google.devtools.build.lib.remote.CasServiceGrpc.CasServiceStub;
import com.google.devtools.build.lib.remote.ContentDigests.ActionKey;
import com.google.devtools.build.lib.remote.ExecutionCacheServiceGrpc.ExecutionCacheServiceBlockingStub;
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.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.Iterator;
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 final class GrpcActionCache implements RemoteActionCache {

  /** Channel over which to send gRPC CAS queries. */
  private final ManagedChannel channel;

  // TODO(olaola): proper profiling to determine the best values for these.
  private final int grpcTimeoutSeconds;
  private final int maxBatchInputs;
  private final int maxChunkSizeBytes;
  private final int maxBatchSizeBytes;

  private static final int MAX_MEMORY_KBYTES = 512 * 1024;

  /** 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, maxChunkSizeBytes);
      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, (long) maxChunkSizeBytes)];
        currentStream.read(blob);
        chunk.setData(ByteString.copyFrom(blob));
        bytesLeft -= blob.length;
      }
      if (bytesLeft == 0) {
        currentStream.close();
        advanceInput();
      }
      return chunk.build();
    }
  }

  @VisibleForTesting
  public GrpcActionCache(ManagedChannel channel, RemoteOptions options) {
    this.channel = channel;
    maxBatchInputs = options.grpcMaxBatchInputs;
    maxChunkSizeBytes = options.grpcMaxChunkSizeBytes;
    maxBatchSizeBytes = options.grpcMaxBatchSizeBytes;
    grpcTimeoutSeconds = options.grpcTimeoutSeconds;
  }

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

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

  private CasServiceBlockingStub getBlockingStub() {
    return CasServiceGrpc.newBlockingStub(channel)
        .withDeadlineAfter(grpcTimeoutSeconds, TimeUnit.SECONDS);
  }

  private CasServiceStub getStub() {
    return CasServiceGrpc.newStub(channel).withDeadlineAfter(grpcTimeoutSeconds, TimeUnit.SECONDS);
  }

  private ImmutableSet<ContentDigest> getMissingDigests(Iterable<ContentDigest> digests) {
    CasLookupRequest.Builder request = CasLookupRequest.newBuilder().addAllDigest(digests);
    if (request.getDigestCount() == 0) {
      return ImmutableSet.of();
    }
    CasStatus status = getBlockingStub().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 = getBlockingStub().uploadTreeMetadata(metaRequest.build());
      if (!reply.getStatus().getSucceeded()) {
        throw new RuntimeException(reply.getStatus().getErrorDetail());
      }
    }
    if (!actionInputs.isEmpty()) {
      ArrayList<Path> paths = new ArrayList<>();
      for (ActionInput actionInput : actionInputs) {
        paths.add(execRoot.getRelative(actionInput.getExecPathString()));
      }
      uploadChunks(paths.size(), new BlobChunkFileIterator(missingDigests, paths.iterator()));
    }
  }

  /**
   * 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 {
    // 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();
    ArrayList<Output> fileOutputs = new ArrayList<>();
    for (Output output : result.getOutputList()) {
      Path path = execRoot.getRelative(output.getPath());
      if (output.getContentCase() == ContentCase.FILE_METADATA) {
        ContentDigest digest = output.getFileMetadata().getDigest();
        if (digest.getSizeBytes() > 0) {
          request.addDigest(digest);
          fileOutputs.add(output);
        } else {
          // Handle empty file locally.
          FileSystemUtils.createDirectoryAndParents(path.getParentDirectory());
          FileSystemUtils.writeContent(path, new byte[0]);
        }
      } else {
        downloadTree(output.getDigest(), path);
      }
    }
    Iterator<CasDownloadReply> replies = getBlockingStub().downloadBlob(request.build());
    for (Output output : fileOutputs) {
      createFileFromStream(
          execRoot.getRelative(output.getPath()), output.getFileMetadata(), replies);
    }
  }

  private void createFileFromStream(
      Path path, FileMetadata fileMetadata, 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(digest.equals(fileMetadata.getDigest()));
    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(fileMetadata.getExecutable());
    }
  }

  private byte[] getBlobFromStream(ContentDigest blobDigest, Iterator<CasDownloadReply> replies)
      throws 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(digest.equals(blobDigest));
    // 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();
    }
    return result;
  }

  /** 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. Note: this is horribly inefficient, need to
   * patch through an overload that uses an ActionInputFile cache to compute the digests!
   *
   * @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;
  }

  /**
   * Download a blob keyed by the given digest and write it to the specified path. Set the
   * executable parameter to the specified value.
   */
  @Override
  public void downloadFileContents(ContentDigest digest, Path dest, boolean executable)
      throws IOException, 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().addDigest(digest);
    Iterator<CasDownloadReply> replies = getBlockingStub().downloadBlob(request.build());
    FileMetadata fileMetadata =
        FileMetadata.newBuilder().setDigest(digest).setExecutable(executable).build();
    createFileFromStream(dest, fileMetadata, replies);
  }

  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.
          if (batchedInputs % maxBatchInputs == 0
              || chunk.getDigest().getSizeBytes() + currentBatchBytes > maxBatchSizeBytes) {
            // The batches execute simultaneously.
            if (requestObserver != null) {
              batchedInputs = 0;
              currentBatchBytes = 0;
              requestObserver.onCompleted();
            }
            batches++;
            responseObserver = new UploadBlobReplyStreamObserver(finishLatch, exception);
            requestObserver = getStub().uploadBlob(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(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;
    if (request.getDigestCount() > 0) {
      replies = getBlockingStub().downloadBlob(request.build());
    }
    ArrayList<byte[]> result = new ArrayList<>();
    for (ContentDigest digest : digests) {
      result.add(digest.getSizeBytes() == 0 ? new byte[0] : getBlobFromStream(digest, replies));
    }
    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) {
    ExecutionCacheServiceBlockingStub stub =
        ExecutionCacheServiceGrpc.newBlockingStub(channel)
            .withDeadlineAfter(grpcTimeoutSeconds, TimeUnit.SECONDS);
    ExecutionCacheRequest request =
        ExecutionCacheRequest.newBuilder().setActionDigest(actionKey.getDigest()).build();
    ExecutionCacheReply reply = stub.getCachedResult(request);
    ExecutionCacheStatus status = reply.getStatus();
    if (!status.getSucceeded()
        && status.getError() != ExecutionCacheStatus.ErrorCode.UNSUPPORTED
        && 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 {
    ExecutionCacheServiceBlockingStub stub =
        ExecutionCacheServiceGrpc.newBlockingStub(channel)
            .withDeadlineAfter(grpcTimeoutSeconds, TimeUnit.SECONDS);
    ExecutionCacheSetRequest request =
        ExecutionCacheSetRequest.newBuilder()
            .setActionDigest(actionKey.getDigest())
            .setResult(result)
            .build();
    ExecutionCacheSetReply reply = stub.setCachedResult(request);
    ExecutionCacheStatus status = reply.getStatus();
    if (!status.getSucceeded()) {
      throw new RuntimeException(status.getErrorDetail());
    }
  }
}