diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiAppendableUnbufferedWritableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiAppendableUnbufferedWritableByteChannel.java new file mode 100644 index 0000000000..1f9eb7698d --- /dev/null +++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/BidiAppendableUnbufferedWritableByteChannel.java @@ -0,0 +1,193 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.storage; + +import com.google.cloud.BaseServiceException; +import com.google.cloud.storage.ChunkSegmenter.ChunkSegment; +import com.google.cloud.storage.UnbufferedWritableByteChannelSession.UnbufferedWritableByteChannel; +import java.io.IOException; +import java.io.InterruptedIOException; +import java.nio.ByteBuffer; +import java.nio.channels.ClosedChannelException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +final class BidiAppendableUnbufferedWritableByteChannel implements UnbufferedWritableByteChannel { + private static final Logger LOGGER = + LoggerFactory.getLogger(BidiAppendableUnbufferedWritableByteChannel.class); + + private final BidiUploadStreamingStream stream; + private final ChunkSegmenter chunkSegmenter; + + private boolean open; + private long writeOffset; + private volatile boolean nextWriteShouldFinalize; + private boolean writeCalledAtLeastOnce; + + /** If write throws an error, don't attempt to finalize things when {@link #close()} is called. */ + private boolean writeThrewError; + + BidiAppendableUnbufferedWritableByteChannel( + BidiUploadStreamingStream stream, ChunkSegmenter chunkSegmenter, long writeOffset) { + this.stream = stream; + this.chunkSegmenter = chunkSegmenter; + this.open = true; + this.writeOffset = writeOffset; + this.nextWriteShouldFinalize = false; + this.writeThrewError = false; + } + + @Override + public long write(ByteBuffer[] srcs, int srcsOffset, int srcsLength) throws IOException { + return internalWrite(srcs, srcsOffset, srcsLength); + } + + @Override + public long writeAndClose(ByteBuffer[] srcs, int offset, int length) throws IOException { + long totalRemaining = Buffers.totalRemaining(srcs, offset, length); + long written = 0; + do { + written += internalWrite(srcs, offset, length); + } while (written < totalRemaining); + close(); + return written; + } + + @Override + public boolean isOpen() { + return open; + } + + @Override + public void close() throws IOException { + if (!open) { + return; + } + try { + if (writeThrewError) { + return; + } + + if (!writeCalledAtLeastOnce) { + stream.flush(); + } + if (nextWriteShouldFinalize) { + //noinspection StatementWithEmptyBody + while (!stream.finishWrite(writeOffset)) {} + } else { + //noinspection StatementWithEmptyBody + while (!stream.closeStream(writeOffset)) {} + } + + awaitResultFuture(); + } finally { + stream.sendClose(); + open = false; + } + } + + public void nextWriteShouldFinalize() { + this.nextWriteShouldFinalize = true; + } + + private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength) throws IOException { + if (!open) { + throw new ClosedChannelException(); + } + // error early. if the result future is already failed, await it to throw the error + if (stream.getResultFuture().isDone()) { + awaitResultFuture(); + return 0; + } + writeCalledAtLeastOnce = true; + + long availableCapacity = stream.availableCapacity(); + if (availableCapacity <= 0) { + return 0; + } + RewindableContent rewindableContent = RewindableContent.of(srcs, srcsOffset, srcsLength); + long totalBufferRemaining = rewindableContent.getLength(); + + ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength, true); + if (data.length == 0) { + return 0; + } + + long bytesConsumed = 0; + for (int i = 0, len = data.length, lastIdx = len - 1; i < len; i++) { + ChunkSegment datum = data[i]; + int size = datum.getB().size(); + boolean appended; + if (i < lastIdx) { + appended = stream.append(datum); + } else if (i == lastIdx && nextWriteShouldFinalize) { + appended = stream.appendAndFinalize(datum); + } else { + appended = stream.appendAndFlush(datum); + } + if (appended) { + bytesConsumed += size; + writeOffset += size; + } else { + // if we weren't able to trigger a flush by reaching the end of the array and calling + // appendAndFlush, explicitly call flush here so that some progress can be made. + // we prefer appendAndFlush so a separate message is not needed, but an extra message + // in order to make progress and free buffer space is better than ending up in a live-lock. + stream.flush(); + break; + } + } + + if (bytesConsumed != totalBufferRemaining) { + rewindableContent.rewindTo(bytesConsumed); + } + + return bytesConsumed; + } + + private void awaitResultFuture() throws IOException { + try { + stream.getResultFuture().get(10_717, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + InterruptedIOException ioe = new InterruptedIOException(); + ioe.initCause(e); + writeThrewError = true; + throw ioe; + } catch (ExecutionException e) { + BaseServiceException coalesce = StorageException.coalesce(e.getCause()); + String message = coalesce.getMessage(); + String ioExceptionMessage = message; + // if the failure is an upload scenario we detect client side, it's message will be + // verbose. To avoid duplication, select the first line only for the io exception + int firstNewLineIndex = message != null ? message.indexOf('\n') : -1; + if (firstNewLineIndex > -1) { + ioExceptionMessage = message.substring(0, firstNewLineIndex); + } + IOException ioException = new IOException(ioExceptionMessage, coalesce); + // ioException.addSuppressed(new AsyncStorageTaskException()); + writeThrewError = true; + throw ioException; + } catch (TimeoutException e) { + writeThrewError = true; + throw new IOException(e); + } + } +}