Search in sources :

Example 1 with QueryWriteStatusResponse

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

the class StubInstanceTest method completedWriteBeforeCloseThrowsOnNextInteraction.

@Test
public void completedWriteBeforeCloseThrowsOnNextInteraction() throws IOException, InterruptedException {
    String resourceName = "early-completed-output-stream-test";
    serviceRegistry.addService(new ByteStreamImplBase() {

        boolean completed = false;

        int writtenBytes = 0;

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

                @Override
                public void onNext(WriteRequest request) {
                    if (!completed) {
                        writtenBytes = request.getData().size();
                        responseObserver.onNext(WriteResponse.newBuilder().setCommittedSize(writtenBytes).build());
                        responseObserver.onCompleted();
                        completed = true;
                    }
                }

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

                @Override
                public void onCompleted() {
                }
            };
        }

        @Override
        public void queryWriteStatus(QueryWriteStatusRequest request, StreamObserver<QueryWriteStatusResponse> responseObserver) {
            if (request.getResourceName().equals(resourceName)) {
                responseObserver.onNext(QueryWriteStatusResponse.newBuilder().setCommittedSize(writtenBytes).setComplete(completed).build());
                responseObserver.onCompleted();
            } else {
                responseObserver.onError(Status.NOT_FOUND.asException());
            }
        }
    });
    Instance instance = newStubInstance("early-completed-outputStream-test");
    ByteString content = ByteString.copyFromUtf8("test-content");
    boolean writeThrewException = false;
    Write operationStreamWrite = instance.getOperationStreamWrite(resourceName);
    try (OutputStream out = operationStreamWrite.getOutput(1, SECONDS, () -> {
    })) {
        content.writeTo(out);
        try {
            content.writeTo(out);
        } catch (Exception e) {
            writeThrewException = true;
        }
    }
    assertThat(writeThrewException).isTrue();
    instance.stop();
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) ServerCallStreamObserver(io.grpc.stub.ServerCallStreamObserver) ClientCallStreamObserver(io.grpc.stub.ClientCallStreamObserver) Write(build.buildfarm.common.Write) Instance(build.buildfarm.instance.Instance) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) ByteString(com.google.protobuf.ByteString) ByteArrayOutputStream(java.io.ByteArrayOutputStream) OutputStream(java.io.OutputStream) ByteStreamImplBase(com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) ByteString(com.google.protobuf.ByteString) IOException(java.io.IOException) ExecutionException(java.util.concurrent.ExecutionException) QueryWriteStatusResponse(com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse) QueryWriteStatusRequest(com.google.bytestream.ByteStreamProto.QueryWriteStatusRequest) Test(org.junit.Test)

Example 2 with QueryWriteStatusResponse

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

the class ByteStreamServiceTest method completedWriteQueryIsFound.

@Test
public void completedWriteQueryIsFound() 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);
    when(simpleBlobStore.containsKey(digest.getHash())).thenReturn(true);
    Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
    ByteStreamBlockingStub service = ByteStreamGrpc.newBlockingStub(channel);
    QueryWriteStatusResponse response = service.queryWriteStatus(QueryWriteStatusRequest.newBuilder().setResourceName(resourceName).build());
    assertThat(response).isEqualTo(QueryWriteStatusResponse.newBuilder().setCommittedSize(digest.getSizeBytes()).setComplete(true).build());
    verify(simpleBlobStore, times(1)).containsKey(eq(digest.getHash()));
}
Also used : Digest(build.bazel.remote.execution.v2.Digest) ByteString(com.google.protobuf.ByteString) Channel(io.grpc.Channel) ByteStreamBlockingStub(com.google.bytestream.ByteStreamGrpc.ByteStreamBlockingStub) QueryWriteStatusResponse(com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse) ByteString(com.google.protobuf.ByteString) Test(org.junit.Test)

Example 3 with QueryWriteStatusResponse

use of com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse 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());
}
Also used : Digest(build.bazel.remote.execution.v2.Digest) ByteString(com.google.protobuf.ByteString) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) InputStream(java.io.InputStream) Channel(io.grpc.Channel) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) Metadata(io.grpc.Metadata) ByteString(com.google.protobuf.ByteString) ClientCall(io.grpc.ClientCall) ByteStreamBlockingStub(com.google.bytestream.ByteStreamGrpc.ByteStreamBlockingStub) QueryWriteStatusResponse(com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse) Test(org.junit.Test)

Example 4 with QueryWriteStatusResponse

use of com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse 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();
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) ServerCallStreamObserver(io.grpc.stub.ServerCallStreamObserver) ClientCallStreamObserver(io.grpc.stub.ClientCallStreamObserver) Write(build.buildfarm.common.Write) Instance(build.buildfarm.instance.Instance) ByteString(com.google.protobuf.ByteString) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) ByteArrayOutputStream(java.io.ByteArrayOutputStream) OutputStream(java.io.OutputStream) ByteStreamImplBase(com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) AtomicReference(java.util.concurrent.atomic.AtomicReference) ByteString(com.google.protobuf.ByteString) QueryWriteStatusResponse(com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse) QueryWriteStatusRequest(com.google.bytestream.ByteStreamProto.QueryWriteStatusRequest) Test(org.junit.Test)

Example 5 with QueryWriteStatusResponse

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

the class StubWriteOutputStreamTest method resetExceptionsAreInterpreted.

@SuppressWarnings("unchecked")
@Test
public void resetExceptionsAreInterpreted() {
    String unimplementedResourceName = "unimplemented-resource";
    QueryWriteStatusRequest unimplementedRequest = QueryWriteStatusRequest.newBuilder().setResourceName(unimplementedResourceName).build();
    doAnswer(invocation -> {
        StreamObserver<QueryWriteStatusResponse> observer = invocation.getArgument(1);
        observer.onError(Status.UNIMPLEMENTED.asException());
        return null;
    }).when(serviceImpl).queryWriteStatus(eq(unimplementedRequest), any(StreamObserver.class));
    String notFoundResourceName = "not-found-resource";
    QueryWriteStatusRequest notFoundRequest = QueryWriteStatusRequest.newBuilder().setResourceName(notFoundResourceName).build();
    doAnswer(invocation -> {
        StreamObserver<QueryWriteStatusResponse> observer = invocation.getArgument(1);
        observer.onError(Status.NOT_FOUND.asException());
        return null;
    }).when(serviceImpl).queryWriteStatus(eq(notFoundRequest), any(StreamObserver.class));
    StubWriteOutputStream write = new StubWriteOutputStream(Suppliers.ofInstance(ByteStreamGrpc.newBlockingStub(channel)), Suppliers.ofInstance(ByteStreamGrpc.newStub(channel)), unimplementedResourceName, Functions.identity(), /* expectedSize=*/
    StubWriteOutputStream.UNLIMITED_EXPECTED_SIZE, /* autoflush=*/
    true);
    assertThat(write.getCommittedSize()).isEqualTo(0);
    verify(serviceImpl, times(1)).queryWriteStatus(eq(unimplementedRequest), any(StreamObserver.class));
    write = new StubWriteOutputStream(Suppliers.ofInstance(ByteStreamGrpc.newBlockingStub(channel)), Suppliers.ofInstance(ByteStreamGrpc.newStub(channel)), notFoundResourceName, Functions.identity(), /* expectedSize=*/
    StubWriteOutputStream.UNLIMITED_EXPECTED_SIZE, /* autoflush=*/
    true);
    assertThat(write.getCommittedSize()).isEqualTo(0);
    verify(serviceImpl, times(1)).queryWriteStatus(eq(notFoundRequest), any(StreamObserver.class));
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) QueryWriteStatusResponse(com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse) ByteString(com.google.protobuf.ByteString) QueryWriteStatusRequest(com.google.bytestream.ByteStreamProto.QueryWriteStatusRequest) Test(org.junit.Test)

Aggregations

QueryWriteStatusResponse (com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse)5 ByteString (com.google.protobuf.ByteString)5 Test (org.junit.Test)5 QueryWriteStatusRequest (com.google.bytestream.ByteStreamProto.QueryWriteStatusRequest)3 WriteRequest (com.google.bytestream.ByteStreamProto.WriteRequest)3 WriteResponse (com.google.bytestream.ByteStreamProto.WriteResponse)3 StreamObserver (io.grpc.stub.StreamObserver)3 Digest (build.bazel.remote.execution.v2.Digest)2 Write (build.buildfarm.common.Write)2 Instance (build.buildfarm.instance.Instance)2 ByteStreamBlockingStub (com.google.bytestream.ByteStreamGrpc.ByteStreamBlockingStub)2 ByteStreamImplBase (com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase)2 Channel (io.grpc.Channel)2 ClientCallStreamObserver (io.grpc.stub.ClientCallStreamObserver)2 ServerCallStreamObserver (io.grpc.stub.ServerCallStreamObserver)2 ByteArrayOutputStream (java.io.ByteArrayOutputStream)2 OutputStream (java.io.OutputStream)2 ClientCall (io.grpc.ClientCall)1 Metadata (io.grpc.Metadata)1 IOException (java.io.IOException)1