Examples with Write build.buildfarm.common.Write used on opensource projects

Example 11 with Write

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

the class StubInstanceTest method outputStreamWrites.

@Test
public void outputStreamWrites() throws IOException, InterruptedException {
    AtomicReference<ByteString> writtenContent = new AtomicReference<>();
    serviceRegistry.addService(new ByteStreamImplBase() {

        ByteString content = ByteString.EMPTY;

        boolean finished = false;

        public void queryWriteStatus(QueryWriteStatusRequest request, StreamObserver<QueryWriteStatusResponse> responseObserver) {
            responseObserver.onNext(QueryWriteStatusResponse.newBuilder().setCommittedSize(content.size()).setComplete(finished).build());
            responseObserver.onCompleted();
        }

        @Override
        public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> responseObserver) {
            return new StreamObserver<WriteRequest>() {

                @Override
                public void onNext(WriteRequest request) {
                    checkState(!finished);
                    if (request.getData().size() != 0) {
                        checkState(request.getWriteOffset() == content.size());
                        content = content.concat(request.getData());
                    }
                    finished = request.getFinishWrite();
                    if (finished) {
                        writtenContent.set(content);
                        responseObserver.onNext(WriteResponse.newBuilder().setCommittedSize(content.size()).build());
                    }
                }

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

                @Override
                public void onCompleted() {
                    responseObserver.onCompleted();
                }
            };
        }
    });
    Instance instance = newStubInstance("outputStream-test");
    String resourceName = "output-stream-test";
    ByteString content = ByteString.copyFromUtf8("test-content");
    Write operationStreamWrite = instance.getOperationStreamWrite(resourceName);
    try (OutputStream out = operationStreamWrite.getOutput(1, SECONDS, () -> {
    })) {
        content.writeTo(out);
    }
    assertThat(writtenContent.get()).isEqualTo(content);
    instance.stop();
}

Example 12 with Write

use of build.buildfarm.common.Write 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;
}

Example 13 with Write

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

the class CASFileCache method newInputFallback.

// CAS fallback methods
private InputStream newInputFallback(Digest digest, long offset) throws IOException {
    checkNotNull(delegate);
    if (digest.getSizeBytes() > maxEntrySizeInBytes) {
        return delegate.newInput(digest, offset);
    }
    Write write = getWrite(digest, UUID.randomUUID(), RequestMetadata.getDefaultInstance());
    return newReadThroughInput(digest, offset, write);
}

Example 14 with Write

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

the class ByteStreamServiceTest method uploadsCanResetInLine.

@Test
public void uploadsCanResetInLine() 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 flush() throws IOException {
            output.flush();
        }

        @Override
        public void write(byte[] b) throws IOException {
            output.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            output.write(b, off, len);
        }

        @Override
        public void write(int b) throws IOException {
            output.write(b);
        }

        @Override
        public boolean isReady() {
            return true;
        }
    };
    Write write = mock(Write.class);
    doAnswer((Answer<Void>) invocation -> {
        output.reset();
        return null;
    }).when(write).reset();
    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.onNext(WriteRequest.newBuilder().setWriteOffset(0).setData(content).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(1)).reset();
    verify(write, times(1)).getFuture();
}

Example 15 with Write

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

the class Worker method createFuseExecFileSystem.

private ExecFileSystem createFuseExecFileSystem(InputStreamFactory remoteInputStreamFactory, ContentAddressableStorage storage) {
    InputStreamFactory storageInputStreamFactory = (digest, offset) -> storage.get(digest).getData().substring((int) offset).newInput();
    InputStreamFactory localPopulatingInputStreamFactory = (blobDigest, offset) -> {
        // FIXME use write
        ByteString content = ByteString.readFrom(remoteInputStreamFactory.newInput(blobDigest, offset));
        if (offset == 0) {
            // extra computations
            Blob blob = new Blob(content, digestUtil);
            // here's hoping that our digest matches...
            storage.put(blob);
        }
        return content.newInput();
    };
    return new FuseExecFileSystem(root, new FuseCAS(root, new EmptyInputStreamFactory(new FailoverInputStreamFactory(storageInputStreamFactory, localPopulatingInputStreamFactory))), storage);
}