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");
}
}
}
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();
}
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();
}
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));
}
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);
}
Aggregations