Search in sources :

Example 1 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest 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");
        }
    }
}
Also used : ExecuteOperationMetadata(build.bazel.remote.execution.v2.ExecuteOperationMetadata) ScheduledFuture(java.util.concurrent.ScheduledFuture) ManagedChannel(io.grpc.ManagedChannel) NegotiationType(io.grpc.netty.NegotiationType) Scanner(java.util.Scanner) SettableFuture(com.google.common.util.concurrent.SettableFuture) Channel(io.grpc.Channel) ByteStreamGrpc(com.google.bytestream.ByteStreamGrpc) DirectoryStream(java.nio.file.DirectoryStream) Executors.newSingleThreadScheduledExecutor(java.util.concurrent.Executors.newSingleThreadScheduledExecutor) StreamObserver(io.grpc.stub.StreamObserver) Digest(build.bazel.remote.execution.v2.Digest) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) ByteStreamUploader.uploadResourceName(build.buildfarm.instance.stub.ByteStreamUploader.uploadResourceName) EXECUTING(build.bazel.remote.execution.v2.ExecutionStage.Value.EXECUTING) Path(java.nio.file.Path) InvalidProtocolBufferException(com.google.protobuf.InvalidProtocolBufferException) ExecutionGrpc(build.bazel.remote.execution.v2.ExecutionGrpc) ExecutionStub(build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionStub) BatchUpdateBlobsRequest(build.bazel.remote.execution.v2.BatchUpdateBlobsRequest) ExecuteRequest(build.bazel.remote.execution.v2.ExecuteRequest) FileStatus(build.buildfarm.common.io.FileStatus) UUID(java.util.UUID) Preconditions.checkState(com.google.common.base.Preconditions.checkState) ByteString(com.google.protobuf.ByteString) List(java.util.List) Size(build.buildfarm.common.Size) MICROSECONDS(java.util.concurrent.TimeUnit.MICROSECONDS) FindMissingBlobsRequest(build.bazel.remote.execution.v2.FindMissingBlobsRequest) Stopwatch(com.google.common.base.Stopwatch) Operation(com.google.longrunning.Operation) DigestUtil(build.buildfarm.common.DigestUtil) ContentAddressableStorageGrpc(build.bazel.remote.execution.v2.ContentAddressableStorageGrpc) ImmutableList(com.google.common.collect.ImmutableList) ByteStreamStub(com.google.bytestream.ByteStreamGrpc.ByteStreamStub) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) FindMissingBlobsResponse(build.bazel.remote.execution.v2.FindMissingBlobsResponse) Code(com.google.rpc.Code) MoreExecutors.shutdownAndAwaitTermination(com.google.common.util.concurrent.MoreExecutors.shutdownAndAwaitTermination) Utils.stat(build.buildfarm.common.io.Utils.stat) FileStore(java.nio.file.FileStore) Files(java.nio.file.Files) HashCode(com.google.common.hash.HashCode) ContentAddressableStorageBlockingStub(build.bazel.remote.execution.v2.ContentAddressableStorageGrpc.ContentAddressableStorageBlockingStub) IOException(java.io.IOException) BatchUpdateBlobsResponse(build.bazel.remote.execution.v2.BatchUpdateBlobsResponse) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) NettyChannelBuilder(io.grpc.netty.NettyChannelBuilder) AtomicLong(java.util.concurrent.atomic.AtomicLong) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) Paths(java.nio.file.Paths) ExecuteResponse(build.bazel.remote.execution.v2.ExecuteResponse) Request(build.bazel.remote.execution.v2.BatchUpdateBlobsRequest.Request) SECONDS(java.util.concurrent.TimeUnit.SECONDS) InputStream(java.io.InputStream) ContentAddressableStorageBlockingStub(build.bazel.remote.execution.v2.ContentAddressableStorageGrpc.ContentAddressableStorageBlockingStub) NettyChannelBuilder(io.grpc.netty.NettyChannelBuilder) Stopwatch(com.google.common.base.Stopwatch) ByteString(com.google.protobuf.ByteString) HashCode(com.google.common.hash.HashCode) UUID(java.util.UUID) Path(java.nio.file.Path) Digest(build.bazel.remote.execution.v2.Digest) BatchUpdateBlobsResponse(build.bazel.remote.execution.v2.BatchUpdateBlobsResponse) ByteStreamStub(com.google.bytestream.ByteStreamGrpc.ByteStreamStub) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) InputStream(java.io.InputStream) BatchUpdateBlobsRequest(build.bazel.remote.execution.v2.BatchUpdateBlobsRequest) ExecuteRequest(build.bazel.remote.execution.v2.ExecuteRequest) FindMissingBlobsRequest(build.bazel.remote.execution.v2.FindMissingBlobsRequest) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) Request(build.bazel.remote.execution.v2.BatchUpdateBlobsRequest.Request) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) BatchUpdateBlobsRequest(build.bazel.remote.execution.v2.BatchUpdateBlobsRequest)

Example 2 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest 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 3 with WriteRequest

use of com.google.bytestream.ByteStreamProto.WriteRequest 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();
}
Also used : Write(build.buildfarm.common.Write) Digest(build.bazel.remote.execution.v2.Digest) ByteString(com.google.protobuf.ByteString) ByteStreamStub(com.google.bytestream.ByteStreamGrpc.ByteStreamStub) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) Channel(io.grpc.Channel) ByteString(com.google.protobuf.ByteString) FeedbackOutputStream(build.buildfarm.common.io.FeedbackOutputStream) HashCode(com.google.common.hash.HashCode) TimeUnit(java.util.concurrent.TimeUnit) UUID(java.util.UUID) Test(org.junit.Test)

Example 4 with WriteRequest

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

the class ByteStreamServiceTest method writePutsIntoBlobStore.

@Test
public void writePutsIntoBlobStore() 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> call = channel.newCall(ByteStreamGrpc.getWriteMethod(), CallOptions.DEFAULT);
    ClientCall.Listener<WriteResponse> callListener = new ClientCall.Listener<WriteResponse>() {

        boolean complete = false;

        boolean callHalfClosed = false;

        @Override
        public void onReady() {
            while (call.isReady()) {
                if (complete) {
                    if (!callHalfClosed) {
                        call.halfClose();
                        callHalfClosed = true;
                    }
                    return;
                }
                call.sendMessage(WriteRequest.newBuilder().setResourceName(resourceName).setData(helloWorld).setFinishWrite(true).build());
                complete = true;
            }
        }
    };
    call.start(callListener, new Metadata());
    call.request(1);
    verify(simpleBlobStore, times(1)).put(eq(digest.getHash()), eq(digest.getSizeBytes()), any(InputStream.class));
}
Also used : Digest(build.bazel.remote.execution.v2.Digest) ByteString(com.google.protobuf.ByteString) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) ClientCall(io.grpc.ClientCall) InputStream(java.io.InputStream) Channel(io.grpc.Channel) WriteResponse(com.google.bytestream.ByteStreamProto.WriteResponse) Metadata(io.grpc.Metadata) ByteString(com.google.protobuf.ByteString) Test(org.junit.Test)

Example 5 with WriteRequest

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

the class BuildFarmServerTest method progressiveUploadCompletes.

@Test
public void progressiveUploadCompletes() throws Exception {
    DigestUtil digestUtil = new DigestUtil(HashFunction.SHA256);
    ByteString content = ByteString.copyFromUtf8("Hello, World!");
    Digest digest = digestUtil.compute(content);
    HashCode hash = HashCode.fromString(digest.getHash());
    UUID uuid = UUID.randomUUID();
    String resourceName = ByteStreamUploader.uploadResourceName(INSTANCE_NAME, uuid, hash, content.size());
    assertThat(getBlob(digest)).isNull();
    FutureWriteResponseObserver futureResponder = new FutureWriteResponseObserver();
    StreamObserver<WriteRequest> requestObserver = ByteStreamGrpc.newStub(inProcessChannel).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 = ByteStreamGrpc.newStub(inProcessChannel).write(futureResponder);
    requestObserver.onNext(WriteRequest.newBuilder().setWriteOffset(6).setResourceName(resourceName).setData(content.substring(6)).setFinishWrite(true).build());
    requestObserver.onCompleted();
    assertThat(futureResponder.get()).isEqualTo(WriteResponse.newBuilder().setCommittedSize(content.size()).build());
    assertThat(getBlob(digest)).isEqualTo(content);
}
Also used : HashCode(com.google.common.hash.HashCode) Digest(build.bazel.remote.execution.v2.Digest) ByteString(com.google.protobuf.ByteString) WriteRequest(com.google.bytestream.ByteStreamProto.WriteRequest) ByteString(com.google.protobuf.ByteString) DigestUtil(build.buildfarm.common.DigestUtil) UUID(java.util.UUID) Test(org.junit.Test)

Aggregations

WriteRequest (com.google.bytestream.ByteStreamProto.WriteRequest)10 ByteString (com.google.protobuf.ByteString)9 Test (org.junit.Test)9 WriteResponse (com.google.bytestream.ByteStreamProto.WriteResponse)7 Digest (build.bazel.remote.execution.v2.Digest)6 StreamObserver (io.grpc.stub.StreamObserver)6 Channel (io.grpc.Channel)5 Write (build.buildfarm.common.Write)4 HashCode (com.google.common.hash.HashCode)4 DigestUtil (build.buildfarm.common.DigestUtil)3 Instance (build.buildfarm.instance.Instance)3 ByteStreamImplBase (com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase)3 ByteStreamStub (com.google.bytestream.ByteStreamGrpc.ByteStreamStub)3 QueryWriteStatusResponse (com.google.bytestream.ByteStreamProto.QueryWriteStatusResponse)3 IOException (java.io.IOException)3 InputStream (java.io.InputStream)3 UUID (java.util.UUID)3 FeedbackOutputStream (build.buildfarm.common.io.FeedbackOutputStream)2 ByteStreamGrpc (com.google.bytestream.ByteStreamGrpc)2 QueryWriteStatusRequest (com.google.bytestream.ByteStreamProto.QueryWriteStatusRequest)2