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

Example 6 with ByteStreamStub

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

the class ByteStreamServiceTest method uploadsCanResetInLine.

@Test
public void uploadsCanResetInLine() 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 flush() throws IOException {
            output.flush();
        }

        @Override
        public void write(byte[] b) throws IOException {
            output.write(b);
        }

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

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

        @Override
        public boolean isReady() {
            return true;
        }
    };
    Write write = mock(Write.class);
    doAnswer((Answer<Void>) invocation -> {
        output.reset();
        return null;
    }).when(write).reset();
    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.onNext(WriteRequest.newBuilder().setWriteOffset(0).setData(content).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(1)).reset();
    verify(write, times(1)).getFuture();
}

Example 7 with ByteStreamStub

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

the class Extract method downloadActionContents.

@SuppressWarnings("ResultOfMethodCallIgnored")
static void downloadActionContents(Path root, String instanceName, Set<Digest> actionDigests, Channel channel) throws IOException, InterruptedException {
    ByteStreamStub bsStub = ByteStreamGrpc.newStub(channel);
    ExecutorService service = newSingleThreadExecutor();
    ListeningScheduledExecutorService retryService = listeningDecorator(newSingleThreadScheduledExecutor());
    Set<Digest> visitedDigests = Sets.newHashSet();
    Set<Digest> visitedDirectories = Sets.newHashSet();
    AtomicLong outstandingOperations = new AtomicLong(0);
    for (Digest actionDigest : actionDigests) {
        ByteString content = getBlobIntoFile("action", instanceName, actionDigest, bsStub, root);
        Action action = Action.parseFrom(content);
        Digest commandDigest = action.getCommandDigest();
        if (!visitedDigests.contains(commandDigest)) {
            visitedDigests.add(commandDigest);
            outstandingOperations.getAndIncrement();
            service.execute(blobGetter(root, instanceName, commandDigest, bsStub, outstandingOperations, retryService));
        }
        Digest inputRootDigest = action.getInputRootDigest();
        if (!visitedDigests.contains(inputRootDigest)) {
            visitedDirectories.add(inputRootDigest);
            visitedDigests.add(inputRootDigest);
            outstandingOperations.getAndIncrement();
            service.execute(directoryGetter(root, instanceName, inputRootDigest, visitedDirectories, visitedDigests, bsStub, service, outstandingOperations, retryService));
        }
    }
    while (outstandingOperations.get() > 0) {
        System.out.println("Waiting on " + outstandingOperations.get() + " operations");
        TimeUnit.SECONDS.sleep(5);
    }
    service.shutdown();
    service.awaitTermination(1, TimeUnit.MINUTES);
    retryService.shutdown();
    retryService.awaitTermination(1, TimeUnit.MINUTES);
}