Search in sources :

Example 1 with CasServiceStub

use of com.google.devtools.build.lib.remote.CasServiceGrpc.CasServiceStub in project bazel by bazelbuild.

the class GrpcActionCache method uploadChunks.

private void uploadChunks(int numItems, BlobChunkIterator blobs) throws InterruptedException, IOException {
    // Maximal number of batches.
    CountDownLatch finishLatch = new CountDownLatch(numItems);
    AtomicReference<RuntimeException> exception = new AtomicReference<>(null);
    UploadBlobReplyStreamObserver responseObserver = null;
    StreamObserver<CasUploadBlobRequest> requestObserver = null;
    int currentBatchBytes = 0;
    int batchedInputs = 0;
    int batches = 0;
    CasServiceStub stub = getStub();
    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 = stub.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) {
        // Finish last batch.
        requestObserver.onCompleted();
    }
    while (batches++ < numItems) {
        // Non-sent batches.
        finishLatch.countDown();
    }
    finishLatch.await(options.grpcTimeoutSeconds, TimeUnit.SECONDS);
    if (exception.get() != null) {
        // Re-throw the first encountered exception.
        throw exception.get();
    }
}
Also used : StatusRuntimeException(io.grpc.StatusRuntimeException) BlobChunk(com.google.devtools.build.lib.remote.RemoteProtocol.BlobChunk) CasUploadBlobRequest(com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadBlobRequest) CasServiceStub(com.google.devtools.build.lib.remote.CasServiceGrpc.CasServiceStub) AtomicReference(java.util.concurrent.atomic.AtomicReference) CountDownLatch(java.util.concurrent.CountDownLatch)

Aggregations

CasServiceStub (com.google.devtools.build.lib.remote.CasServiceGrpc.CasServiceStub)1 BlobChunk (com.google.devtools.build.lib.remote.RemoteProtocol.BlobChunk)1 CasUploadBlobRequest (com.google.devtools.build.lib.remote.RemoteProtocol.CasUploadBlobRequest)1 StatusRuntimeException (io.grpc.StatusRuntimeException)1 CountDownLatch (java.util.concurrent.CountDownLatch)1 AtomicReference (java.util.concurrent.atomic.AtomicReference)1