Search in sources :

Example 1 with ByteStringQueueInputStream

use of build.buildfarm.common.io.ByteStringQueueInputStream in project bazel-buildfarm by bazelbuild.

the class ByteStreamHelper method newInput.

@SuppressWarnings("Guava")
public static InputStream newInput(String resourceName, long offset, Supplier<ByteStreamStub> bsStubSupplier, Supplier<Backoff> backoffSupplier, Predicate<Status> isRetriable, @Nullable ListeningScheduledExecutorService retryService) throws IOException {
    ReadRequest request = ReadRequest.newBuilder().setResourceName(resourceName).setReadOffset(offset).build();
    BlockingQueue<ByteString> queue = new ArrayBlockingQueue<>(1);
    ByteStringQueueInputStream inputStream = new ByteStringQueueInputStream(queue);
    // this interface needs to operate similar to open, where it
    // throws an exception on creation. We will need to wait around
    // for the response to come back in order to supply the stream or
    // throw the exception it receives
    SettableFuture<InputStream> streamReadyFuture = SettableFuture.create();
    StreamObserver<ReadResponse> responseObserver = new StreamObserver<ReadResponse>() {

        long requestOffset = offset;

        long currentOffset = offset;

        Backoff backoff = backoffSupplier.get();

        @Override
        public void onNext(ReadResponse response) {
            streamReadyFuture.set(inputStream);
            ByteString data = response.getData();
            try {
                queue.put(data);
                currentOffset += data.size();
            } catch (InterruptedException e) {
                // cancel context?
                inputStream.setException(e);
            }
        }

        private void retryRequest() {
            requestOffset = currentOffset;
            bsStubSupplier.get().read(request.toBuilder().setReadOffset(requestOffset).build(), this);
        }

        @Override
        public void onError(Throwable t) {
            Status status = Status.fromThrowable(t);
            long nextDelayMillis = backoff.nextDelayMillis();
            if (status.getCode() == Status.Code.DEADLINE_EXCEEDED && currentOffset != requestOffset) {
                backoff = backoffSupplier.get();
                retryRequest();
            } else if (retryService == null || nextDelayMillis < 0 || !isRetriable.test(status)) {
                streamReadyFuture.setException(t);
                inputStream.setException(t);
            } else {
                try {
                    ListenableFuture<?> schedulingResult = retryService.schedule(this::retryRequest, nextDelayMillis, TimeUnit.MILLISECONDS);
                    schedulingResult.addListener(() -> {
                        try {
                            schedulingResult.get();
                        } catch (ExecutionException e) {
                            inputStream.setException(e.getCause());
                        } catch (InterruptedException e) {
                            inputStream.setException(e);
                        }
                    }, MoreExecutors.directExecutor());
                } catch (RejectedExecutionException e) {
                    inputStream.setException(e);
                }
            }
        }

        @Override
        public void onCompleted() {
            inputStream.setCompleted();
        }
    };
    bsStubSupplier.get().read(request, responseObserver);
    // perfectly reasonable to be used as a wait point
    try {
        return streamReadyFuture.get();
    } catch (InterruptedException e) {
        try {
            inputStream.close();
        } catch (RuntimeException closeEx) {
            e.addSuppressed(e);
        }
        IOException ioEx = new ClosedByInterruptException();
        ioEx.addSuppressed(e);
        throw ioEx;
    } catch (ExecutionException e) {
        Throwable cause = e.getCause();
        Status status = Status.fromThrowable(cause);
        if (status.getCode() == Status.Code.NOT_FOUND) {
            IOException ioEx = new NoSuchFileException(resourceName);
            ioEx.addSuppressed(cause);
            throw ioEx;
        }
        Throwables.throwIfInstanceOf(cause, IOException.class);
        throw new IOException(cause);
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) Status(io.grpc.Status) ByteString(com.google.protobuf.ByteString) ByteStringQueueInputStream(build.buildfarm.common.io.ByteStringQueueInputStream) InputStream(java.io.InputStream) NoSuchFileException(java.nio.file.NoSuchFileException) IOException(java.io.IOException) ByteStringQueueInputStream(build.buildfarm.common.io.ByteStringQueueInputStream) Backoff(build.buildfarm.common.grpc.Retrier.Backoff) RejectedExecutionException(java.util.concurrent.RejectedExecutionException) ClosedByInterruptException(java.nio.channels.ClosedByInterruptException) ArrayBlockingQueue(java.util.concurrent.ArrayBlockingQueue) ReadResponse(com.google.bytestream.ByteStreamProto.ReadResponse) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) RejectedExecutionException(java.util.concurrent.RejectedExecutionException) ExecutionException(java.util.concurrent.ExecutionException) ReadRequest(com.google.bytestream.ByteStreamProto.ReadRequest)

Aggregations

Backoff (build.buildfarm.common.grpc.Retrier.Backoff)1 ByteStringQueueInputStream (build.buildfarm.common.io.ByteStringQueueInputStream)1 ReadRequest (com.google.bytestream.ByteStreamProto.ReadRequest)1 ReadResponse (com.google.bytestream.ByteStreamProto.ReadResponse)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 ByteString (com.google.protobuf.ByteString)1 Status (io.grpc.Status)1 StreamObserver (io.grpc.stub.StreamObserver)1 IOException (java.io.IOException)1 InputStream (java.io.InputStream)1 ClosedByInterruptException (java.nio.channels.ClosedByInterruptException)1 NoSuchFileException (java.nio.file.NoSuchFileException)1 ArrayBlockingQueue (java.util.concurrent.ArrayBlockingQueue)1 ExecutionException (java.util.concurrent.ExecutionException)1 RejectedExecutionException (java.util.concurrent.RejectedExecutionException)1