com.google.bytestream.ByteStreamProto.ReadRequest Java Examples

@Test
public void getHandlesNotFound() {
  Digest digest = DIGEST_UTIL.compute(ByteString.copyFromUtf8("nonexistent"));
  String instanceName = "test";
  final AtomicReference<Boolean> readCalled = new AtomicReference<>(false);
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName())
              .isEqualTo(String.format("%s/blobs/%s", instanceName, DigestUtil.toString(digest)));
          readCalled.compareAndSet(false, true);
          responseObserver.onError(Status.NOT_FOUND.asException());
        }
      });

  GrpcCAS cas =
      new GrpcCAS(
          instanceName,
          InProcessChannelBuilder.forName(fakeServerName).directExecutor().build(),
          mock(ByteStreamUploader.class),
          onExpirations);
  assertThat(cas.get(digest)).isNull();
  assertThat(readCalled.get()).isTrue();
}
 

@Test
public void inputStreamThrowsOnDeadlineExceededWithoutProgress()
    throws IOException, InterruptedException {
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          responseObserver.onError(Status.DEADLINE_EXCEEDED.asException());
        }
      });
  OutputStream out = mock(OutputStream.class);
  IOException ioException = null;
  Instance instance = newStubInstance("input-stream-deadline-exceeded");
  Digest timeoutDigest = Digest.newBuilder().setHash("timeout-blob-name").setSizeBytes(1).build();
  try (InputStream in =
      instance.newBlobInput(timeoutDigest, 0, 1, SECONDS, RequestMetadata.getDefaultInstance())) {
    ByteStreams.copy(in, out);
  } catch (IOException e) {
    ioException = e;
  }
  assertThat(ioException).isNotNull();
  Status status = Status.fromThrowable(ioException);
  assertThat(status.getCode()).isEqualTo(Code.DEADLINE_EXCEEDED);
  verifyZeroInteractions(out);
  instance.stop();
}
 

@Test
public void missingBlobReadIsNotFound() {
  ByteString helloWorld = ByteString.copyFromUtf8("Hello, World!");
  Digest digest = DIGEST_UTIL.compute(helloWorld);

  Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
  ByteStreamBlockingStub service = ByteStreamGrpc.newBlockingStub(channel);

  when(simpleBlobStore.get(eq(digest.getHash()), any(OutputStream.class)))
      .thenReturn(immediateFuture(false));
  ReadRequest request =
      ReadRequest.newBuilder().setResourceName(createBlobDownloadResourceName(digest)).build();
  StatusRuntimeException notFoundException = null;
  try {
    if (service.read(request).hasNext()) {
      fail("no responses should be available");
    }
  } catch (StatusRuntimeException e) {
    assertThat(Status.fromThrowable(e).getCode()).isEqualTo(Code.NOT_FOUND);
    notFoundException = e;
  }
  assertThat(notFoundException).isNotNull();
}
 

@Test
public void testDisablingDigestVerification() throws Exception {
  // Test that when digest verification is disabled a corrupted download works.

  RemoteOptions remoteOptions = Options.getDefaults(RemoteOptions.class);
  remoteOptions.remoteVerifyDownloads = false;

  GrpcCacheClient client = newClient(remoteOptions);
  Digest digest = DIGEST_UTIL.computeAsUtf8("foo");
  ByteString downloadContents = ByteString.copyFromUtf8("bar");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          responseObserver.onNext(ReadResponse.newBuilder().setData(downloadContents).build());
          responseObserver.onCompleted();
        }
      });

  assertThat(downloadBlob(client, digest)).isEqualTo(downloadContents.toByteArray());
}
 

@Test
public void testDownloadFailsOnDigestMismatch() throws Exception {
  // Test that the download fails when a blob/file has a different content hash than expected.

  GrpcCacheClient client = newClient();
  Digest digest = DIGEST_UTIL.computeAsUtf8("foo");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          ByteString data = ByteString.copyFromUtf8("bar");
          responseObserver.onNext(ReadResponse.newBuilder().setData(data).build());
          responseObserver.onCompleted();
        }
      });
  IOException e = assertThrows(IOException.class, () -> downloadBlob(client, digest));
  assertThat(e).hasMessageThat().contains(digest.getHash());
  assertThat(e).hasMessageThat().contains(DIGEST_UTIL.computeAsUtf8("bar").getHash());
}
 

@Test
public void downloadBlobPassesThroughDeadlineExceededWithoutProgress() throws IOException {
  Backoff mockBackoff = Mockito.mock(Backoff.class);
  Mockito.when(mockBackoff.nextDelayMillis()).thenReturn(-1L);
  final GrpcCacheClient client =
      newClient(Options.getDefaults(RemoteOptions.class), () -> mockBackoff);
  final Digest digest = DIGEST_UTIL.computeAsUtf8("abcdefg");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName().contains(digest.getHash())).isTrue();
          ByteString data = ByteString.copyFromUtf8("abcdefg");
          if (request.getReadOffset() == 0) {
            responseObserver.onNext(
                ReadResponse.newBuilder().setData(data.substring(0, 2)).build());
          }
          responseObserver.onError(Status.DEADLINE_EXCEEDED.asException());
        }
      });
  IOException e = assertThrows(IOException.class, () -> downloadBlob(client, digest));
  Status st = Status.fromThrowable(e);
  assertThat(st.getCode()).isEqualTo(Status.Code.DEADLINE_EXCEEDED);
  Mockito.verify(mockBackoff, Mockito.times(1)).nextDelayMillis();
}
 

@Test
public void testDownloadBlobSingleChunk() throws Exception {
  final GrpcCacheClient client = newClient();
  final Digest digest = DIGEST_UTIL.computeAsUtf8("abcdefg");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName().contains(digest.getHash())).isTrue();
          responseObserver.onNext(
              ReadResponse.newBuilder().setData(ByteString.copyFromUtf8("abcdefg")).build());
          responseObserver.onCompleted();
        }
      });
  assertThat(new String(downloadBlob(client, digest), UTF_8)).isEqualTo("abcdefg");
}
 

@Test
public void testDownloadBlobMultipleChunks() throws Exception {
  final GrpcCacheClient client = newClient();
  final Digest digest = DIGEST_UTIL.computeAsUtf8("abcdefg");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName().contains(digest.getHash())).isTrue();
          responseObserver.onNext(
              ReadResponse.newBuilder().setData(ByteString.copyFromUtf8("abc")).build());
          responseObserver.onNext(
              ReadResponse.newBuilder().setData(ByteString.copyFromUtf8("def")).build());
          responseObserver.onNext(
              ReadResponse.newBuilder().setData(ByteString.copyFromUtf8("g")).build());
          responseObserver.onCompleted();
        }
      });
  assertThat(new String(downloadBlob(client, digest), UTF_8)).isEqualTo("abcdefg");
}
 

@Override
public void get(
    Digest digest,
    long offset,
    long count,
    ServerCallStreamObserver<ByteString> blobObserver,
    RequestMetadata requestMetadata) {
  ReadRequest request =
      ReadRequest.newBuilder()
          .setResourceName(getBlobName(digest))
          .setReadOffset(offset)
          .setReadLimit(count)
          .build();
  ByteStreamGrpc.newStub(channel)
      .withInterceptors(attachMetadataInterceptor(requestMetadata))
      .read(
          request,
          new DelegateServerCallStreamObserver<ReadResponse, ByteString>(blobObserver) {
            @Override
            public void onNext(ReadResponse response) {
              blobObserver.onNext(response.getData());
            }

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

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

@Test
public void onExpirationCalledWhenNotFound() {
  Digest digest = DIGEST_UTIL.compute(ByteString.copyFromUtf8("nonexistent"));
  String instanceName = "test";
  final AtomicReference<Boolean> readCalled = new AtomicReference<>(false);
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName())
              .isEqualTo(String.format("%s/blobs/%s", instanceName, DigestUtil.toString(digest)));
          readCalled.compareAndSet(false, true);
          responseObserver.onError(Status.NOT_FOUND.asException());
        }
      });

  Runnable onExpiration = mock(Runnable.class);
  onExpirations.put(digest, onExpiration);

  GrpcCAS cas =
      new GrpcCAS(
          instanceName,
          InProcessChannelBuilder.forName(fakeServerName).directExecutor().build(),
          mock(ByteStreamUploader.class),
          onExpirations);
  assertThat(cas.get(digest)).isNull();
  assertThat(readCalled.get()).isTrue();
  verify(onExpiration, times(1)).run();
}
 

@Test
public void inputStreamThrowsNonDeadlineExceededCausal()
    throws IOException, InterruptedException {
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          responseObserver.onError(Status.UNAVAILABLE.asException());
        }
      });
  OutputStream out = mock(OutputStream.class);
  IOException ioException = null;
  Instance instance = newStubInstance("input-stream-non-deadline-exceeded");
  Digest unavailableDigest =
      Digest.newBuilder().setHash("unavailable-blob-name").setSizeBytes(1).build();
  try (InputStream in =
      instance.newBlobInput(
          unavailableDigest, 0, 1, SECONDS, RequestMetadata.getDefaultInstance())) {
    ByteStreams.copy(in, out);
  } catch (IOException e) {
    ioException = e;
  }
  assertThat(ioException).isNotNull();
  Status status = Status.fromThrowable(ioException);
  assertThat(status.getCode()).isEqualTo(Code.UNAVAILABLE);
  verifyZeroInteractions(out);
  instance.stop();
}
 

@Test
public void inputStreamRetriesOnDeadlineExceededWithProgress()
    throws IOException, InterruptedException {
  ByteString content = ByteString.copyFromUtf8("1");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        boolean first = true;

        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          if (first && request.getReadOffset() == 0) {
            first = false;
            responseObserver.onNext(ReadResponse.newBuilder().setData(content).build());
            responseObserver.onError(Status.DEADLINE_EXCEEDED.asException());
          } else if (request.getReadOffset() == 1) {
            responseObserver.onCompleted();
          } else {
            // all others fail with unimplemented
            responseObserver.onError(Status.UNIMPLEMENTED.asException());
          }
        }
      });
  Instance instance = newStubInstance("input-stream-stalled");
  ByteArrayOutputStream out = new ByteArrayOutputStream();
  Digest delayedDigest = Digest.newBuilder().setHash("delayed-blob-name").setSizeBytes(1).build();
  try (InputStream in =
      instance.newBlobInput(delayedDigest, 0, 1, SECONDS, RequestMetadata.getDefaultInstance())) {
    ByteStreams.copy(in, out);
  }
  assertThat(ByteString.copyFrom(out.toByteArray())).isEqualTo(content);
  instance.stop();
}
 

@Test
public void newInputThrowsOnNotFound() {
  String resourceName = "not/found/resource";
  ReadRequest readRequest = ReadRequest.newBuilder().setResourceName(resourceName).build();
  doAnswer(
          (Answer)
              invocation -> {
                StreamObserver<ReadResponse> observer = invocation.getArgument(1);
                observer.onError(Status.NOT_FOUND.asException());
                return null;
              })
      .when(serviceImpl)
      .read(eq(readRequest), any(StreamObserver.class));

  try (InputStream in =
      ByteStreamHelper.newInput(
          resourceName,
          /* offset=*/ 0,
          Suppliers.ofInstance(ByteStreamGrpc.newStub(channel)),
          NO_RETRIES::newBackoff,
          NO_RETRIES::isRetriable,
          /* retryService=*/ null)) {
    fail("should not get here");
  } catch (IOException e) {
    assertThat(e).isInstanceOf(NoSuchFileException.class);
  }

  verify(serviceImpl, times(1)).read(eq(readRequest), any(StreamObserver.class));
}
 

@Override
public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
  assertThat(cannedReplies.keySet()).contains(request);
  int errCount = numErrors.getOrDefault(request, 0);
  if (errCount < MAX_ERRORS) {
    numErrors.put(request, errCount + 1);
    responseObserver.onError(Status.UNAVAILABLE.asRuntimeException());  // Retriable error.
  } else {
    responseObserver.onNext(cannedReplies.get(request));
    responseObserver.onCompleted();
  }
}
 

@Test
public void downloadBlobIsRetriedWithProgress() throws IOException, InterruptedException {
  Backoff mockBackoff = Mockito.mock(Backoff.class);
  final GrpcCacheClient client =
      newClient(Options.getDefaults(RemoteOptions.class), () -> mockBackoff);
  final Digest digest = DIGEST_UTIL.computeAsUtf8("abcdefg");
  serviceRegistry.addService(
      new ByteStreamImplBase() {
        @Override
        public void read(ReadRequest request, StreamObserver<ReadResponse> responseObserver) {
          assertThat(request.getResourceName().contains(digest.getHash())).isTrue();
          ByteString data = ByteString.copyFromUtf8("abcdefg");
          int off = (int) request.getReadOffset();
          if (off == 0) {
            data = data.substring(0, 1);
          } else {
            data = data.substring(off);
          }
          responseObserver.onNext(ReadResponse.newBuilder().setData(data).build());
          if (off == 0) {
            responseObserver.onError(Status.DEADLINE_EXCEEDED.asException());
          } else {
            responseObserver.onCompleted();
          }
        }
      });
  assertThat(new String(downloadBlob(client, digest), UTF_8)).isEqualTo("abcdefg");
  Mockito.verify(mockBackoff, Mockito.never()).nextDelayMillis();
}
 

/** Reads a ByteStream onto the arg consumer. */
public static ListenableFuture<Unit> readByteStream(
    String instanceName,
    Protocol.Digest digest,
    ByteStreamStub byteStreamStub,
    ThrowingConsumer<ByteString, IOException> dataConsumer,
    int casDeadline) {
  String name = getResourceName(instanceName, digest);
  SettableFuture<Unit> future = SettableFuture.create();
  byteStreamStub
      .withDeadlineAfter(casDeadline, TimeUnit.SECONDS)
      .read(
          ReadRequest.newBuilder().setResourceName(name).setReadLimit(0).setReadOffset(0).build(),
          new StreamObserver<ReadResponse>() {
            int size = 0;
            MessageDigest messageDigest = PROTOCOL.getMessageDigest();

            @Override
            public void onNext(ReadResponse value) {
              try {
                ByteString data = value.getData();
                size += data.size();
                messageDigest.update(data.asReadOnlyByteBuffer());
                dataConsumer.accept(data);
              } catch (IOException e) {
                onError(e);
              }
            }

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

            @Override
            public void onCompleted() {
              String digestHash = HashCode.fromBytes(messageDigest.digest()).toString();
              if (size == digest.getSize() && digestHash.equals(digest.getHash())) {
                future.set(null);
              } else {
                future.setException(
                    new BuckUncheckedExecutionException(
                        "Digest of received bytes: "
                            + digestHash
                            + ":"
                            + size
                            + " doesn't match expected digest: "
                            + digest));
              }
            }
          });
  return future;
}
 

@Override
public void handleReq(ReadRequest message) {
  builder.setRequest(message);
}
 

@Test
public void readSlicesLargeChunksFromInstance() throws Exception {
  // pick a large chunk size
  long size = CHUNK_SIZE * 10 + CHUNK_SIZE - 47;
  ByteString content;
  try (ByteString.Output out =
      ByteString.newOutput(
          ByteStreamService.CHUNK_SIZE * 10 + ByteStreamService.CHUNK_SIZE - 47)) {
    for (long i = 0; i < size; i++) {
      out.write((int) (i & 0xff));
    }
    content = out.toByteString();
  }
  Digest digest = DIGEST_UTIL.compute(content);
  String resourceName = "blobs/" + DigestUtil.toString(digest);
  ReadRequest request = ReadRequest.newBuilder().setResourceName(resourceName).build();

  Instance instance = mock(Instance.class);
  when(instances.getFromBlob(eq(resourceName))).thenReturn(instance);
  doAnswer(answerVoid((blobDigest, offset, limit, chunkObserver, metadata) -> {}))
      .when(instance)
      .getBlob(
          eq(digest),
          eq(request.getReadOffset()),
          eq((long) content.size()),
          any(ServerCallStreamObserver.class),
          eq(RequestMetadata.getDefaultInstance()));
  Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
  ByteStreamStub service = ByteStreamGrpc.newStub(channel);
  CountingReadObserver readObserver = new CountingReadObserver();
  service.read(request, readObserver);
  ArgumentCaptor<ServerCallStreamObserver<ByteString>> observerCaptor =
      ArgumentCaptor.forClass(ServerCallStreamObserver.class);
  verify(instance, times(1))
      .getBlob(
          eq(digest),
          eq(request.getReadOffset()),
          eq((long) content.size()),
          observerCaptor.capture(),
          eq(RequestMetadata.getDefaultInstance()));

  StreamObserver<ByteString> responseObserver = observerCaptor.getValue();
  // supply entire content
  responseObserver.onNext(content);
  responseObserver.onCompleted();

  assertThat(readObserver.isCompleted()).isTrue();
  assertThat(readObserver.getData()).isEqualTo(content);
  List<Integer> sizes = readObserver.getSizesList();
  assertThat(sizes.size()).isEqualTo(11); // 10 + 1 incomplete chunk
  assertThat(Iterables.filter(sizes, (responseSize) -> responseSize > CHUNK_SIZE)).isEmpty();
}
 

@Test
public void skippedInputIsNotInResponse()
    throws ExecutionException, IOException, InterruptedException {
  ByteString helloWorld = ByteString.copyFromUtf8("Hello, World!");
  Digest digest = DIGEST_UTIL.compute(helloWorld);

  Channel channel = InProcessChannelBuilder.forName(fakeServerName).directExecutor().build();
  ByteStreamStub service = ByteStreamGrpc.newStub(channel);

  SettableFuture<Boolean> getComplete = SettableFuture.create();
  when(simpleBlobStore.get(eq(digest.getHash()), any(OutputStream.class)))
      .thenReturn(getComplete);
  ArgumentCaptor<OutputStream> outputStreamCaptor = ArgumentCaptor.forClass(OutputStream.class);

  ReadRequest request =
      ReadRequest.newBuilder()
          .setResourceName(createBlobDownloadResourceName(digest))
          .setReadOffset(6)
          .build();
  SettableFuture<ByteString> readComplete = SettableFuture.create();
  service.read(
      request,
      new StreamObserver<ReadResponse>() {
        ByteString content = ByteString.EMPTY;

        @Override
        public void onNext(ReadResponse response) {
          content = content.concat(response.getData());
        }

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

        @Override
        public void onCompleted() {
          readComplete.set(content);
        }
      });

  verify(simpleBlobStore, times(1)).get(eq(digest.getHash()), outputStreamCaptor.capture());
  try (OutputStream outputStream = outputStreamCaptor.getValue()) {
    outputStream.write(helloWorld.toByteArray());
    getComplete.set(true);
  }
  assertThat(readComplete.get()).isEqualTo(helloWorld.substring(6));
}
 

private void readBlob(ReadRequest request, StreamObserver<ReadResponse> responseObserver)
    throws IOException, InterruptedException, InvalidResourceNameException {
  String resourceName = request.getResourceName();

  Digest digest = parseBlobDigest(resourceName);

  OutputStream responseOut =
      new ChunkOutputStream(DEFAULT_CHUNK_SIZE) {
        @Override
        public void onChunk(byte[] b, int off, int len) {
          responseObserver.onNext(
              ReadResponse.newBuilder().setData(ByteString.copyFrom(b, off, len)).build());
        }
      };

  addCallback(
      getBlob(digest, request.getReadOffset(), request.getReadLimit(), responseOut),
      new FutureCallback<Boolean>() {
        private void onError(Status status) {
          responseObserver.onError(status.asException());
        }

        @Override
        public void onSuccess(Boolean success) {
          if (success) {
            try {
              responseOut.close();
              responseObserver.onCompleted();
            } catch (IOException e) {
              onError(Status.fromThrowable(e));
            }
          } else {
            onError(Status.NOT_FOUND);
          }
        }

        @Override
        public void onFailure(Throwable t) {
          try {
            responseOut.close();
          } catch (IOException ioException) {
            ioException.printStackTrace();
          }
          onError(Status.fromThrowable(t));
        }
      },
      directExecutor());
}