Examples with ByteStreamStub com.google.bytestream.ByteStreamGrpc.ByteStreamStub used on opensource projects

Example 1 with ByteStreamStub

use of com.google.bytestream.ByteStreamGrpc.ByteStreamStub in project bazel-buildfarm by bazelbuild.

the class Executor method loadFilesIntoCAS.

private static void loadFilesIntoCAS(String instanceName, Channel channel, Path blobsDir) throws Exception {
    ContentAddressableStorageBlockingStub casStub = ContentAddressableStorageGrpc.newBlockingStub(channel);
    List<Digest> missingDigests = findMissingBlobs(instanceName, blobsDir, casStub);
    UUID uploadId = UUID.randomUUID();
    int[] bucketSizes = new int[128];
    BatchUpdateBlobsRequest.Builder[] buckets = new BatchUpdateBlobsRequest.Builder[128];
    for (int i = 0; i < 128; i++) {
        bucketSizes[i] = 0;
        buckets[i] = BatchUpdateBlobsRequest.newBuilder().setInstanceName(instanceName);
    }
    ByteStreamStub bsStub = ByteStreamGrpc.newStub(channel);
    for (Digest missingDigest : missingDigests) {
        Path path = blobsDir.resolve(missingDigest.getHash() + "_" + missingDigest.getSizeBytes());
        if (missingDigest.getSizeBytes() < Size.mbToBytes(1)) {
            Request request = Request.newBuilder().setDigest(missingDigest).setData(ByteString.copyFrom(Files.readAllBytes(path))).build();
            int maxBucketSize = 0;
            long minBucketSize = Size.mbToBytes(2) + 1;
            int maxBucketIndex = 0;
            int minBucketIndex = -1;
            int size = (int) missingDigest.getSizeBytes() + 48;
            for (int i = 0; i < 128; i++) {
                int newBucketSize = bucketSizes[i] + size;
                if (newBucketSize < Size.mbToBytes(2) && bucketSizes[i] < minBucketSize) {
                    minBucketSize = bucketSizes[i];
                    minBucketIndex = i;
                }
                if (bucketSizes[i] > maxBucketSize) {
                    maxBucketSize = bucketSizes[i];
                    maxBucketIndex = i;
                }
            }
            if (minBucketIndex < 0) {
                bucketSizes[maxBucketIndex] = size;
                BatchUpdateBlobsRequest batchRequest = buckets[maxBucketIndex].build();
                Stopwatch stopwatch = Stopwatch.createStarted();
                BatchUpdateBlobsResponse batchResponse = casStub.batchUpdateBlobs(batchRequest);
                long usecs = stopwatch.elapsed(MICROSECONDS);
                checkState(batchResponse.getResponsesList().stream().allMatch(response -> Code.forNumber(response.getStatus().getCode()) == Code.OK));
                System.out.println("Updated " + batchRequest.getRequestsCount() + " blobs in " + (usecs / 1000.0) + "ms");
                buckets[maxBucketIndex] = BatchUpdateBlobsRequest.newBuilder().setInstanceName(instanceName).addRequests(request);
            } else {
                bucketSizes[minBucketIndex] += size;
                buckets[minBucketIndex].addRequests(request);
            }
        } else {
            Stopwatch stopwatch = Stopwatch.createStarted();
            SettableFuture<WriteResponse> writtenFuture = SettableFuture.create();
            StreamObserver<WriteRequest> requestObserver = bsStub.write(new StreamObserver<WriteResponse>() {

                @Override
                public void onNext(WriteResponse response) {
                    writtenFuture.set(response);
                }

                @Override
                public void onCompleted() {
                }

                @Override
                public void onError(Throwable t) {
                    writtenFuture.setException(t);
                }
            });
            HashCode hash = HashCode.fromString(missingDigest.getHash());
            String resourceName = uploadResourceName(instanceName, uploadId, hash, missingDigest.getSizeBytes());
            try (InputStream in = Files.newInputStream(path)) {
                boolean first = true;
                long writtenBytes = 0;
                byte[] buf = new byte[64 * 1024];
                while (writtenBytes != missingDigest.getSizeBytes()) {
                    int len = in.read(buf);
                    WriteRequest.Builder request = WriteRequest.newBuilder();
                    if (first) {
                        request.setResourceName(resourceName);
                    }
                    request.setData(ByteString.copyFrom(buf, 0, len)).setWriteOffset(writtenBytes);
                    if (writtenBytes + len == missingDigest.getSizeBytes()) {
                        request.setFinishWrite(true);
                    }
                    requestObserver.onNext(request.build());
                    writtenBytes += len;
                    first = false;
                }
                writtenFuture.get();
                System.out.println("Wrote long " + DigestUtil.toString(missingDigest) + " in " + (stopwatch.elapsed(MICROSECONDS) / 1000.0) + "ms");
            }
        }
    }
    for (int i = 0; i < 128; i++) {
        if (bucketSizes[i] > 0) {
            BatchUpdateBlobsRequest batchRequest = buckets[i].build();
            Stopwatch stopwatch = Stopwatch.createStarted();
            BatchUpdateBlobsResponse batchResponse = casStub.batchUpdateBlobs(batchRequest);
            long usecs = stopwatch.elapsed(MICROSECONDS);
            checkState(batchResponse.getResponsesList().stream().allMatch(response -> Code.forNumber(response.getStatus().getCode()) == Code.OK));
            System.out.println("Updated " + batchRequest.getRequestsCount() + " blobs in " + (usecs / 1000.0) + "ms");
        }
    }
}

Example 2 with ByteStreamStub

use of com.google.bytestream.ByteStreamGrpc.ByteStreamStub in project bazel-buildfarm by bazelbuild.

the class ByteStreamServiceTest method readSlicesLargeChunksFromInstance.

@SuppressWarnings("unchecked")
@Test
public void readSlicesLargeChunksFromInstance() throws Exception {
    // pick a large chunk size
    long size = CHUNK_SIZE * 10 + CHUNK_SIZE - 47;
    ByteString content;
    try (ByteString.Output out = ByteString.newOutput(ByteStreamService.CHUNK_SIZE * 10 + ByteStreamService.CHUNK_SIZE - 47)) {
        for (long i = 0; i < size; i++) {
            out.write((int) (i & 0xff));
        }
        content = out.toByteString();
    }
    Digest digest = DIGEST_UTIL.compute(content);
    String resourceName = "blobs/" + DigestUtil.toString(digest);
    ReadRequest request = ReadRequest.newBuilder().setResourceName(resourceName).build();
    doAnswer(answerVoid((blobDigest, offset, limit, chunkObserver, metadata) -> {
    })).when(instance).getBlob(eq(digest), eq(request.getReadOffset()), eq((long) content.size()), any(ServerCallStreamObserver.class), eq(RequestMetadata.getDefaultInstance()));
    Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
    ByteStreamStub service = ByteStreamGrpc.newStub(channel);
    CountingReadObserver readObserver = new CountingReadObserver();
    service.read(request, readObserver);
    ArgumentCaptor<ServerCallStreamObserver<ByteString>> observerCaptor = ArgumentCaptor.forClass(ServerCallStreamObserver.class);
    verify(instance, times(1)).getBlob(eq(digest), eq(request.getReadOffset()), eq((long) content.size()), observerCaptor.capture(), eq(RequestMetadata.getDefaultInstance()));
    StreamObserver<ByteString> responseObserver = observerCaptor.getValue();
    // supply entire content
    responseObserver.onNext(content);
    responseObserver.onCompleted();
    assertThat(readObserver.isCompleted()).isTrue();
    assertThat(readObserver.getData()).isEqualTo(content);
    List<Integer> sizes = readObserver.getSizesList();
    // 10 + 1 incomplete chunk
    assertThat(sizes.size()).isEqualTo(11);
    assertThat(sizes.stream().filter((responseSize) -> responseSize > CHUNK_SIZE).collect(Collectors.toList())).isEmpty();
}

Example 3 with ByteStreamStub

use of com.google.bytestream.ByteStreamGrpc.ByteStreamStub in project bazel-buildfarm by bazelbuild.

the class ByteStreamServiceTest method uploadsCanProgressAfterCancellation.

@Test
public void uploadsCanProgressAfterCancellation() throws Exception {
    ByteString content = ByteString.copyFromUtf8("Hello, World!");
    Digest digest = DIGEST_UTIL.compute(content);
    UUID uuid = UUID.randomUUID();
    SettableFuture<Long> writtenFuture = SettableFuture.create();
    ByteString.Output output = ByteString.newOutput((int) digest.getSizeBytes());
    FeedbackOutputStream out = new FeedbackOutputStream() {

        @Override
        public void close() {
            if (output.size() == digest.getSizeBytes()) {
                writtenFuture.set(digest.getSizeBytes());
            }
        }

        @Override
        public void write(byte[] b, int off, int len) {
            output.write(b, off, len);
        }

        @Override
        public void write(int b) {
            output.write(b);
        }

        @Override
        public boolean isReady() {
            return true;
        }
    };
    Write write = mock(Write.class);
    when(write.getOutput(any(Long.class), any(TimeUnit.class), any(Runnable.class))).thenReturn(out);
    doAnswer(invocation -> (long) output.size()).when(write).getCommittedSize();
    when(write.getFuture()).thenReturn(writtenFuture);
    when(instance.getBlobWrite(digest, uuid, RequestMetadata.getDefaultInstance())).thenReturn(write);
    HashCode hash = HashCode.fromString(digest.getHash());
    String resourceName = ByteStreamUploader.uploadResourceName(/* instanceName=*/
    null, uuid, hash, digest.getSizeBytes());
    Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
    ByteStreamStub service = ByteStreamGrpc.newStub(channel);
    FutureWriteResponseObserver futureResponder = new FutureWriteResponseObserver();
    StreamObserver<WriteRequest> requestObserver = service.write(futureResponder);
    ByteString shortContent = content.substring(0, 6);
    requestObserver.onNext(WriteRequest.newBuilder().setWriteOffset(0).setResourceName(resourceName).setData(shortContent).build());
    requestObserver.onError(Status.CANCELLED.asException());
    // should be done
    assertThat(futureResponder.isDone()).isTrue();
    futureResponder = new FutureWriteResponseObserver();
    requestObserver = service.write(futureResponder);
    requestObserver.onNext(WriteRequest.newBuilder().setWriteOffset(6).setResourceName(resourceName).setData(content.substring(6)).setFinishWrite(true).build());
    assertThat(futureResponder.get()).isEqualTo(WriteResponse.newBuilder().setCommittedSize(content.size()).build());
    requestObserver.onCompleted();
    verify(write, atLeastOnce()).getCommittedSize();
    verify(write, atLeastOnce()).getOutput(any(Long.class), any(TimeUnit.class), any(Runnable.class));
    verify(write, times(2)).getFuture();
}

Example 4 with ByteStreamStub

use of com.google.bytestream.ByteStreamGrpc.ByteStreamStub in project bazel-buildfarm by bazelbuild.

the class ByteStreamServiceTest method skippedInputIsNotInResponse.

@Test
public void skippedInputIsNotInResponse() throws ExecutionException, IOException, InterruptedException {
    ByteString helloWorld = ByteString.copyFromUtf8("Hello, World!");
    Digest digest = DIGEST_UTIL.compute(helloWorld);
    Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
    ByteStreamStub service = ByteStreamGrpc.newStub(channel);
    SettableFuture<Boolean> getComplete = SettableFuture.create();
    when(simpleBlobStore.get(eq(digest.getHash()), any(OutputStream.class))).thenReturn(getComplete);
    ArgumentCaptor<OutputStream> outputStreamCaptor = ArgumentCaptor.forClass(OutputStream.class);
    ReadRequest request = ReadRequest.newBuilder().setResourceName(createBlobDownloadResourceName(digest)).setReadOffset(6).build();
    SettableFuture<ByteString> readComplete = SettableFuture.create();
    service.read(request, new StreamObserver<ReadResponse>() {

        ByteString content = ByteString.EMPTY;

        @Override
        public void onNext(ReadResponse response) {
            content = content.concat(response.getData());
        }

        @Override
        public void onError(Throwable t) {
            readComplete.setException(t);
        }

        @Override
        public void onCompleted() {
            readComplete.set(content);
        }
    });
    verify(simpleBlobStore, times(1)).get(eq(digest.getHash()), outputStreamCaptor.capture());
    try (OutputStream outputStream = outputStreamCaptor.getValue()) {
        outputStream.write(helloWorld.toByteArray());
        getComplete.set(true);
    }
    assertThat(readComplete.get()).isEqualTo(helloWorld.substring(6));
}

Example 5 with ByteStreamStub

use of com.google.bytestream.ByteStreamGrpc.ByteStreamStub in project bazel-buildfarm by bazelbuild.

the class GrpcCAS method newWrite.

@SuppressWarnings("Guava")
public static Write newWrite(Channel channel, String instanceName, Digest digest, UUID uuid, RequestMetadata requestMetadata) {
    HashCode hash = HashCode.fromString(digest.getHash());
    String resourceName = ByteStreamUploader.uploadResourceName(instanceName, uuid, hash, digest.getSizeBytes());
    Supplier<ByteStreamBlockingStub> bsBlockingStub = Suppliers.memoize(() -> ByteStreamGrpc.newBlockingStub(channel).withInterceptors(attachMetadataInterceptor(requestMetadata)));
    Supplier<ByteStreamStub> bsStub = Suppliers.memoize(() -> ByteStreamGrpc.newStub(channel).withInterceptors(attachMetadataInterceptor(requestMetadata)));
    return new StubWriteOutputStream(bsBlockingStub, bsStub, resourceName, Functions.identity(), digest.getSizeBytes(), /* autoflush=*/
    false);
}