Examples with WriteRequest com.google.bytestream.ByteStreamProto.WriteRequest used on opensource projects

Example 6 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest in project bazel-buildfarm by bazelbuild.

the class ByteStreamServiceTest method writeCanBeResumed.

@Test
public void writeCanBeResumed() throws IOException, InterruptedException {
    ByteString helloWorld = ByteString.copyFromUtf8("Hello, World!");
    Digest digest = DIGEST_UTIL.compute(helloWorld);
    String uuid = UUID.randomUUID().toString();
    String resourceName = createBlobUploadResourceName(uuid, digest);
    Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
    ClientCall<WriteRequest, WriteResponse> initialCall = channel.newCall(ByteStreamGrpc.getWriteMethod(), CallOptions.DEFAULT);
    ByteString initialData = helloWorld.substring(0, 6);
    ClientCall.Listener<WriteResponse> initialCallListener = new ClientCall.Listener<WriteResponse>() {

        boolean complete = false;

        boolean callHalfClosed = false;

        @Override
        public void onReady() {
            while (initialCall.isReady()) {
                if (complete) {
                    if (!callHalfClosed) {
                        initialCall.halfClose();
                        callHalfClosed = true;
                    }
                    return;
                }
                initialCall.sendMessage(WriteRequest.newBuilder().setResourceName(resourceName).setData(initialData).build());
                complete = true;
            }
        }
    };
    initialCall.start(initialCallListener, new Metadata());
    initialCall.request(1);
    ByteStreamBlockingStub service = ByteStreamGrpc.newBlockingStub(channel);
    QueryWriteStatusResponse response = service.queryWriteStatus(QueryWriteStatusRequest.newBuilder().setResourceName(resourceName).build());
    assertThat(response.getCommittedSize()).isEqualTo(initialData.size());
    assertThat(response.getComplete()).isFalse();
    ClientCall<WriteRequest, WriteResponse> finishCall = channel.newCall(ByteStreamGrpc.getWriteMethod(), CallOptions.DEFAULT);
    ClientCall.Listener<WriteResponse> finishCallListener = new ClientCall.Listener<WriteResponse>() {

        boolean complete = false;

        boolean callHalfClosed = false;

        @Override
        public void onReady() {
            while (finishCall.isReady()) {
                if (complete) {
                    if (!callHalfClosed) {
                        finishCall.halfClose();
                        callHalfClosed = true;
                    }
                    return;
                }
                finishCall.sendMessage(WriteRequest.newBuilder().setResourceName(resourceName).setWriteOffset(initialData.size()).setData(helloWorld.substring(initialData.size())).setFinishWrite(true).build());
                complete = true;
            }
        }
    };
    finishCall.start(finishCallListener, new Metadata());
    finishCall.request(1);
    ArgumentCaptor<InputStream> inputStreamCaptor = ArgumentCaptor.forClass(InputStream.class);
    verify(simpleBlobStore, times(1)).put(eq(digest.getHash()), eq(digest.getSizeBytes()), inputStreamCaptor.capture());
    InputStream inputStream = inputStreamCaptor.getValue();
    assertThat(inputStream.available()).isEqualTo(helloWorld.size());
    byte[] data = new byte[helloWorld.size()];
    assertThat(inputStream.read(data)).isEqualTo(helloWorld.size());
    assertThat(data).isEqualTo(helloWorld.toByteArray());
}

Example 7 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest in project bazel-buildfarm by bazelbuild.

the class ByteStreamUploaderTest method uploadBlobPropagatesIOExceptionFromIncompleteUpload.

@Test(expected = IOException.class)
public void uploadBlobPropagatesIOExceptionFromIncompleteUpload() throws IOException, InterruptedException {
    serviceRegistry.addService(new ByteStreamImplBase() {

        @Override
        public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> responseObserver) {
            responseObserver.onNext(WriteResponse.newBuilder().setCommittedSize(1).build());
            responseObserver.onCompleted();
            return new StreamObserver<WriteRequest>() {

                @Override
                public void onNext(WriteRequest request) {
                }

                @Override
                public void onError(Throwable t) {
                }

                @Override
                public void onCompleted() {
                }
            };
        }
    });
    ByteStreamUploader uploader = new ByteStreamUploader(/* instanceName=*/
    null, InProcessChannelBuilder.forName(fakeServerName).directExecutor().build(), /* callCredentials=*/
    null, /* callTimeoutSecs=*/
    1, NO_RETRIES);
    Chunker chunker = Chunker.builder().setInput(ByteString.copyFromUtf8("Hello, World!")).build();
    uploader.uploadBlob(HashCode.fromInt(42), chunker);
}

Example 8 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest in project bazel-buildfarm by bazelbuild.

the class StubInstanceTest method outputStreamWrites.

@Test
public void outputStreamWrites() throws IOException, InterruptedException {
    AtomicReference<ByteString> writtenContent = new AtomicReference<>();
    serviceRegistry.addService(new ByteStreamImplBase() {

        ByteString content = ByteString.EMPTY;

        boolean finished = false;

        public void queryWriteStatus(QueryWriteStatusRequest request, StreamObserver<QueryWriteStatusResponse> responseObserver) {
            responseObserver.onNext(QueryWriteStatusResponse.newBuilder().setCommittedSize(content.size()).setComplete(finished).build());
            responseObserver.onCompleted();
        }

        @Override
        public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> responseObserver) {
            return new StreamObserver<WriteRequest>() {

                @Override
                public void onNext(WriteRequest request) {
                    checkState(!finished);
                    if (request.getData().size() != 0) {
                        checkState(request.getWriteOffset() == content.size());
                        content = content.concat(request.getData());
                    }
                    finished = request.getFinishWrite();
                    if (finished) {
                        writtenContent.set(content);
                        responseObserver.onNext(WriteResponse.newBuilder().setCommittedSize(content.size()).build());
                    }
                }

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

                @Override
                public void onCompleted() {
                    responseObserver.onCompleted();
                }
            };
        }
    });
    Instance instance = newStubInstance("outputStream-test");
    String resourceName = "output-stream-test";
    ByteString content = ByteString.copyFromUtf8("test-content");
    Write operationStreamWrite = instance.getOperationStreamWrite(resourceName);
    try (OutputStream out = operationStreamWrite.getOutput(1, SECONDS, () -> {
    })) {
        content.writeTo(out);
    }
    assertThat(writtenContent.get()).isEqualTo(content);
    instance.stop();
}

Example 9 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest in project bazel-buildfarm by bazelbuild.

the class StubWriteOutputStreamTest method resetIsRespectedOnSubsequentWrite.

@SuppressWarnings("unchecked")
@Test
public void resetIsRespectedOnSubsequentWrite() throws IOException {
    String resourceName = "reset-resource";
    StubWriteOutputStream write = new StubWriteOutputStream(Suppliers.ofInstance(ByteStreamGrpc.newBlockingStub(channel)), Suppliers.ofInstance(ByteStreamGrpc.newStub(channel)), resourceName, Functions.identity(), /* expectedSize=*/
    StubWriteOutputStream.UNLIMITED_EXPECTED_SIZE, /* autoflush=*/
    true);
    ByteString content = ByteString.copyFromUtf8("Hello, World");
    try (OutputStream out = write.getOutput(1, SECONDS, () -> {
    })) {
        content.writeTo(out);
        write.reset();
        content.writeTo(out);
    }
    verify(serviceImpl, times(1)).write(any(StreamObserver.class));
    ArgumentCaptor<WriteRequest> writeRequestCaptor = ArgumentCaptor.forClass(WriteRequest.class);
    verify(writeObserver, times(3)).onNext(writeRequestCaptor.capture());
    List<WriteRequest> requests = writeRequestCaptor.getAllValues();
    assertThat(requests.get(0).getWriteOffset()).isEqualTo(requests.get(1).getWriteOffset());
    assertThat(requests.get(2).getFinishWrite()).isTrue();
}

Example 10 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest 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();
}