Examples with FeedbackOutputStream build.buildfarm.common.io.FeedbackOutputStream used on opensource projects

Example 1 with FeedbackOutputStream

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

the class CASFileCacheTest method interruptDeferredDuringExpirations.

@SuppressWarnings("unchecked")
@Test
public void interruptDeferredDuringExpirations() throws IOException, InterruptedException {
    Blob expiringBlob;
    try (ByteString.Output out = ByteString.newOutput(1024)) {
        for (int i = 0; i < 1024; i++) {
            out.write(0);
        }
        expiringBlob = new Blob(out.toByteString(), DIGEST_UTIL);
    }
    fileCache.put(expiringBlob);
    // state of CAS
    // 1024-byte key
    AtomicReference<Throwable> exRef = new AtomicReference(null);
    // 0 = not blocking
    // 1 = blocking
    // 2 = delegate write
    AtomicInteger writeState = new AtomicInteger(0);
    // this will ensure that the discharge task is blocked until we release it
    Future<Void> blockingExpiration = expireService.submit(() -> {
        writeState.getAndIncrement();
        while (writeState.get() != 0) {
            try {
                MICROSECONDS.sleep(1);
            } catch (InterruptedException e) {
            // ignore
            }
        }
        return null;
    });
    when(delegate.getWrite(eq(expiringBlob.getDigest()), any(UUID.class), any(RequestMetadata.class))).thenReturn(new NullWrite() {

        @Override
        public FeedbackOutputStream getOutput(long deadlineAfter, TimeUnit deadlineAfterUnits, Runnable onReadyHandler) throws IOException {
            try {
                while (writeState.get() != 1) {
                    MICROSECONDS.sleep(1);
                }
            } catch (InterruptedException e) {
                throw new IOException(e);
            }
            // move into output stream state
            writeState.getAndIncrement();
            return super.getOutput(deadlineAfter, deadlineAfterUnits, onReadyHandler);
        }
    });
    Thread expiringThread = new Thread(() -> {
        try {
            fileCache.put(new Blob(ByteString.copyFromUtf8("Hello, World"), DIGEST_UTIL));
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        fail("should not get here");
    });
    expiringThread.setUncaughtExceptionHandler((t, e) -> exRef.set(e));
    // wait for blocking state
    while (writeState.get() != 1) {
        MICROSECONDS.sleep(1);
    }
    expiringThread.start();
    while (writeState.get() != 2) {
        MICROSECONDS.sleep(1);
    }
    // expiry has been initiated, thread should be waiting
    // just trying to ensure that we've reached the future wait point
    MICROSECONDS.sleep(10);
    // hopefully this will be scheduled *after* the discharge task
    Future<Void> completedExpiration = expireService.submit(() -> null);
    // interrupt it
    expiringThread.interrupt();
    assertThat(expiringThread.isAlive()).isTrue();
    assertThat(completedExpiration.isDone()).isFalse();
    writeState.set(0);
    while (!blockingExpiration.isDone()) {
        MICROSECONDS.sleep(1);
    }
    expiringThread.join();
    // CAS should now be empty due to expiration and failed put
    while (!completedExpiration.isDone()) {
        MICROSECONDS.sleep(1);
    }
    assertThat(fileCache.size()).isEqualTo(0);
    Throwable t = exRef.get();
    assertThat(t).isNotNull();
    t = t.getCause();
    assertThat(t).isNotNull();
    assertThat(t).isInstanceOf(InterruptedException.class);
}

Example 2 with FeedbackOutputStream

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

the class CASFileCacheTest method readThroughSwitchedToLocalContinues.

@Test
public void readThroughSwitchedToLocalContinues() throws Exception {
    ByteString content = ByteString.copyFromUtf8("Hello, World");
    Blob blob = new Blob(content, DIGEST_UTIL);
    ExecutorService service = newSingleThreadExecutor();
    SettableFuture<Void> writeComplete = SettableFuture.create();
    // we need to register callbacks on the shared write future
    Write write = new NullWrite() {

        @Override
        public ListenableFuture<Long> getFuture() {
            return Futures.transform(writeComplete, result -> blob.getDigest().getSizeBytes(), directExecutor());
        }

        @Override
        public FeedbackOutputStream getOutput(long deadlineAfter, TimeUnit deadlineAfterUnits, Runnable onReadyHandler) {
            return new FeedbackOutputStream() {

                int offset = 0;

                @Override
                public void write(int b) {
                    throw new UnsupportedOperationException();
                }

                @Override
                public void write(byte[] buf, int ofs, int len) throws IOException {
                    // hangs on second read
                    if (offset == 6) {
                        service.submit(() -> writeComplete.set(null));
                        throw new ClosedChannelException();
                    }
                    offset += len;
                }

                @Override
                public boolean isReady() {
                    return true;
                }
            };
        }
    };
    when(delegate.getWrite(eq(blob.getDigest()), any(UUID.class), any(RequestMetadata.class))).thenReturn(write);
    when(delegate.newInput(eq(blob.getDigest()), eq(0L))).thenReturn(content.newInput());
    // the switch will reset to this point
    InputStream switchedIn = content.newInput();
    switchedIn.skip(6);
    when(delegate.newInput(eq(blob.getDigest()), eq(6L))).thenReturn(switchedIn);
    InputStream in = fileCache.newReadThroughInput(blob.getDigest(), 0, write);
    byte[] buf = new byte[content.size()];
    // advance to the middle of the content
    assertThat(in.read(buf, 0, 6)).isEqualTo(6);
    assertThat(ByteString.copyFrom(buf, 0, 6)).isEqualTo(content.substring(0, 6));
    verify(delegate, times(1)).newInput(blob.getDigest(), 0L);
    // read the remaining content
    int remaining = content.size() - 6;
    assertThat(in.read(buf, 6, remaining)).isEqualTo(remaining);
    assertThat(ByteString.copyFrom(buf)).isEqualTo(content);
    if (!shutdownAndAwaitTermination(service, 1, SECONDS)) {
        throw new RuntimeException("could not shut down service");
    }
}

Example 3 with FeedbackOutputStream

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

the class Worker method streamIntoWriteFuture.

private ListenableFuture<Long> streamIntoWriteFuture(InputStream in, Write write, Digest digest) throws IOException {
    SettableFuture<Long> writtenFuture = SettableFuture.create();
    int chunkSizeBytes = (int) Size.kbToBytes(128);
    // The following callback is performed each time the write stream is ready.
    // For each callback we only transfer a small part of the input stream in order to avoid
    // accumulating a large buffer.  When the file is done being transfered,
    // the callback closes the stream and prepares the future.
    FeedbackOutputStream out = write.getOutput(/* deadlineAfter=*/
    1, /* deadlineAfterUnits=*/
    DAYS, () -> {
        try {
            FeedbackOutputStream outStream = (FeedbackOutputStream) write;
            while (outStream.isReady()) {
                if (!CopyBytes(in, outStream, chunkSizeBytes)) {
                    return;
                }
            }
        } catch (IOException e) {
            if (!write.isComplete()) {
                write.reset();
                logger.log(Level.SEVERE, "unexpected error transferring file for " + digest, e);
            }
        }
    });
    write.getFuture().addListener(() -> {
        try {
            try {
                out.close();
            } catch (IOException e) {
            // ignore
            }
            long committedSize = write.getCommittedSize();
            if (committedSize != digest.getSizeBytes()) {
                logger.log(Level.WARNING, format("committed size %d did not match expectation for digestUtil", committedSize));
            }
            writtenFuture.set(digest.getSizeBytes());
        } catch (RuntimeException e) {
            writtenFuture.setException(e);
        }
    }, directExecutor());
    return writtenFuture;
}

Example 4 with FeedbackOutputStream

use of build.buildfarm.common.io.FeedbackOutputStream 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();
}

Example 5 with FeedbackOutputStream

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

the class CASFileCache method newWrite.

Write newWrite(BlobWriteKey key, ListenableFuture<Long> future) {
    Write write = new Write() {

        CancellableOutputStream out = null;

        boolean isReset = false;

        SettableFuture<Void> closedFuture = null;

        long fileCommittedSize = -1;

        @Override
        public synchronized void reset() {
            try {
                if (out != null) {
                    out.cancel();
                }
            } catch (IOException e) {
                logger.log(Level.SEVERE, "could not reset write " + DigestUtil.toString(key.getDigest()) + ":" + key.getIdentifier(), e);
            } finally {
                isReset = true;
            }
        }

        @Override
        public synchronized long getCommittedSize() {
            long committedSize = getCommittedSizeFromOutOrDisk();
            if (committedSize == 0 && out == null) {
                isReset = true;
            }
            return committedSize;
        }

        long getCommittedSizeFromOutOrDisk() {
            if (isComplete()) {
                return key.getDigest().getSizeBytes();
            }
            return getCommittedSizeFromOut();
        }

        synchronized long getCommittedSizeFromOut() {
            if (out == null) {
                if (fileCommittedSize < 0) {
                    // we need to cache this from disk until an out stream is acquired
                    String blobKey = getKey(key.getDigest(), false);
                    Path blobKeyPath = getPath(blobKey);
                    try {
                        fileCommittedSize = Files.size(blobKeyPath.resolveSibling(blobKey + "." + key.getIdentifier()));
                    } catch (IOException e) {
                        fileCommittedSize = 0;
                    }
                }
                return fileCommittedSize;
            }
            return out.getWritten();
        }

        @Override
        public synchronized boolean isComplete() {
            return getFuture().isDone() || ((closedFuture == null || closedFuture.isDone()) && containsLocal(key.getDigest(), /* result=*/
            null, (key) -> {
            }));
        }

        @Override
        public synchronized ListenableFuture<FeedbackOutputStream> getOutputFuture(long deadlineAfter, TimeUnit deadlineAfterUnits, Runnable onReadyHandler) {
            if (closedFuture == null || closedFuture.isDone()) {
                try {
                    // this isn't great, and will block when there are multiple requesters
                    return immediateFuture(getOutput(deadlineAfter, deadlineAfterUnits, onReadyHandler));
                } catch (IOException e) {
                    return immediateFailedFuture(e);
                }
            }
            return transformAsync(closedFuture, result -> getOutputFuture(deadlineAfter, deadlineAfterUnits, onReadyHandler), directExecutor());
        }

        @Override
        public synchronized FeedbackOutputStream getOutput(long deadlineAfter, TimeUnit deadlineAfterUnits, Runnable onReadyHandler) throws IOException {
            // will block until it is returned via a close.
            if (closedFuture != null) {
                try {
                    closedFuture.get();
                } catch (ExecutionException e) {
                    throw new IOException(e.getCause());
                } catch (InterruptedException e) {
                    throw new IOException(e);
                }
            }
            SettableFuture<Void> outClosedFuture = SettableFuture.create();
            UniqueWriteOutputStream uniqueOut = createUniqueWriteOutput(out, key.getDigest(), UUID.fromString(key.getIdentifier()), () -> outClosedFuture.set(null), this::isComplete, isReset);
            commitOpenState(uniqueOut.delegate(), outClosedFuture);
            return uniqueOut;
        }

        private void commitOpenState(CancellableOutputStream out, SettableFuture<Void> closedFuture) {
            // transition the Write to an open state, and modify all internal state required
            // atomically
            // this function must. not. throw.
            this.out = out;
            this.closedFuture = closedFuture;
            // they will likely write to this, so we can no longer assume isReset.
            // might want to subscribe to a write event on the stream
            isReset = false;
            // our cached file committed size is now invalid
            fileCommittedSize = -1;
        }

        @Override
        public ListenableFuture<Long> getFuture() {
            return future;
        }
    };
    write.getFuture().addListener(write::reset, directExecutor());
    return write;
}