com.google.bytestream.ByteStreamProto.ReadResponse Java Examples

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

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

void readLimitedBlob(
    Instance instance,
    Digest digest,
    long offset,
    long limit,
    StreamObserver<ReadResponse> responseObserver) {
  ServerCallStreamObserver<ReadResponse> target =
      onErrorLogReadObserver(
          format("%s(%s)", DigestUtil.toString(digest), instance.getName()),
          offset,
          (ServerCallStreamObserver<ReadResponse>) responseObserver);
  try {
    instance.getBlob(
        digest,
        offset,
        limit,
        newChunkObserver(target),
        TracingMetadataUtils.fromCurrentContext());
  } catch (Exception e) {
    target.onError(e);
  }
}
 

void readBlob(
    Instance instance,
    Digest digest,
    long offset,
    long limit,
    StreamObserver<ReadResponse> responseObserver) {
  long available = digest.getSizeBytes() - offset;
  if (available == 0) {
    responseObserver.onCompleted();
  } else if (available < 0) {
    responseObserver.onError(OUT_OF_RANGE.asException());
  } else {
    if (limit == 0) {
      limit = available;
    } else {
      limit = Math.min(available, limit);
    }
    readLimitedBlob(instance, digest, offset, limit, responseObserver);
  }
}
 

void maybeInstanceRead(
    String resourceName, long offset, long limit, StreamObserver<ReadResponse> responseObserver)
    throws InstanceNotFoundException, InvalidResourceNameException {
  switch (detectResourceOperation(resourceName)) {
    case Blob:
      readBlob(
          instances.getFromBlob(resourceName),
          parseBlobDigest(resourceName),
          offset,
          limit,
          responseObserver);
      break;
    case OperationStream:
      readOperationStream(
          instances.getFromOperationStream(resourceName),
          resourceName,
          offset,
          limit,
          responseObserver);
      break;
    case UploadBlob:
    default:
      responseObserver.onError(INVALID_ARGUMENT.asException());
      break;
  }
}
 

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

@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 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 onNext(ReadResponse response) {
  ByteString data = response.getData();
  try {
    data.writeTo(sink);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  sizes.add(data.size());
}
 

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

ServerCallStreamObserver<ByteString> newChunkObserver(
    ServerCallStreamObserver<ReadResponse> responseObserver) {
  return new DelegateServerCallStreamObserver<ByteString, ReadResponse>(responseObserver) {
    @Override
    public void onNext(ByteString data) {
      while (!data.isEmpty()) {
        ByteString slice;
        if (data.size() > CHUNK_SIZE) {
          slice = data.substring(0, CHUNK_SIZE);
          data = data.substring(CHUNK_SIZE);
        } else {
          slice = data;
          data = ByteString.EMPTY;
        }
        responseObserver.onNext(ReadResponse.newBuilder().setData(slice).build());
      }
    }

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

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

ServerCallStreamObserver<ReadResponse> onErrorLogReadObserver(
    String name, long offset, ServerCallStreamObserver<ReadResponse> delegate) {
  return new UniformDelegateServerCallStreamObserver<ReadResponse>(delegate) {
    long responseCount = 0;
    long responseBytes = 0;

    @Override
    public void onNext(ReadResponse response) {
      delegate.onNext(response);
      responseCount++;
      responseBytes += response.getData().size();
    }

    @Override
    public void onError(Throwable t) {
      Status status = Status.fromThrowable(t);
      if (status.getCode() != Code.NOT_FOUND) {
        java.util.logging.Level level = Level.SEVERE;
        if (responseCount > 0 && status.getCode() == Code.DEADLINE_EXCEEDED
            || status.getCode() == Code.CANCELLED) {
          level = Level.WARNING;
        }
        String message = format("error reading %s at offset %d", name, offset);
        if (responseCount > 0) {
          message +=
              format(" after %d responses and %d bytes of content", responseCount, responseBytes);
        }
        logger.log(level, message, t);
      }
      delegate.onError(t);
    }

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

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

/** 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 handleResp(ReadResponse message) {
  numReads++;
  bytesRead += message.getData().size();
}
 

@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));
}
 

void readFrom(InputStream in, long limit, CallStreamObserver<ReadResponse> target) {
  final class ReadFromOnReadyHandler implements Runnable {
    private final byte[] buf = new byte[CHUNK_SIZE];
    private final boolean unlimited = limit == 0;
    private long remaining = limit;
    private boolean complete = false;

    ReadResponse next() throws IOException {
      int readBytes = in.read(buf, 0, (int) Math.min(remaining, buf.length));
      if (readBytes <= 0) {
        if (readBytes == -1) {
          if (!unlimited) {
            throw new UnexpectedEndOfStreamException(remaining, limit);
          }
          complete = true;
        }
        return ReadResponse.getDefaultInstance();
      }

      if (readBytes > remaining) {
        logger.log(Level.WARNING, format("read %d bytes, expected %d", readBytes, remaining));
        readBytes = (int) remaining;
      }
      remaining -= readBytes;
      complete = remaining == 0;
      return ReadResponse.newBuilder().setData(ByteString.copyFrom(buf, 0, readBytes)).build();
    }

    @Override
    public void run() {
      if (!complete) {
        copy();
      }
    }

    void copy() {
      try {
        while (target.isReady() && !complete) {
          ReadResponse response = next();
          if (response.getData().size() != 0) {
            target.onNext(response);
          }
        }

        if (complete) {
          in.close();
          target.onCompleted();
        }
      } catch (Exception e) {
        complete = true;
        try {
          in.close();
        } catch (IOException closeEx) {
          e.addSuppressed(closeEx);
        }
        if (e instanceof UnexpectedEndOfStreamException) {
          e = Status.UNAVAILABLE.withCause(e).asException();
        }
        target.onError(e);
      }
    }
  }
  target.setOnReadyHandler(new ReadFromOnReadyHandler());
}
 

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