fix: update BlobAppendableUpload implementation to periodically flush for large writes (#3278)

This main idea here is to allow async incremental clearing of the outbound queue even when large writes are performed.

Previously, when using the MinFlushStrategy, if a large write was performed (larger than maxPendingBytes) a single `flush: true state_lookup: true` would be sent to GCS, thereby making it so that no new writes could be accepted until the full `maxPendingBytes` where ack'd. This change updates so that if a write is larger than `minFlushSize` a message will be annotated `flush: true state_lookup: true`. This doesn't necessarily mean that a flush will be done every `minFlushSize` as the message packed can be up to 2MiB, this will simply annotate a message as `flush: true state_lookup: true` if it has been at least `minFlushSize` bytes since we sent a flush.

Author: BenWhitehead

google-cloud-storage/src/main/java/com/google/cloud/storage/BidiAppendableUnbufferedWritableByteChannel.java

@@ -31,23 +31,30 @@ final class BidiAppendableUnbufferedWritableByteChannel implements UnbufferedWri
 
   private final BidiUploadStreamingStream stream;
   private final ChunkSegmenter chunkSegmenter;
+  private final long flushInterval;
 
   private boolean open;
   private long writeOffset;
   private volatile boolean nextWriteShouldFinalize;
   private boolean writeCalledAtLeastOnce;
+  private long lastFlushOffset;
 
   /** 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) {
+      BidiUploadStreamingStream stream,
+      ChunkSegmenter chunkSegmenter,
+      long flushInterval,
+      long writeOffset) {
     this.stream = stream;
     this.chunkSegmenter = chunkSegmenter;
+    this.flushInterval = flushInterval;
     this.open = true;
     this.writeOffset = writeOffset;
     this.nextWriteShouldFinalize = false;
     this.writeThrewError = false;
+    this.lastFlushOffset = writeOffset;
   }
 
   @Override
@@ -141,8 +148,9 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength) th
     for (int i = 0, len = data.length, lastIdx = len - 1; i < len; i++) {
       ChunkSegment datum = data[i];
       int size = datum.getB().size();
+      boolean shouldFlush = writeOffset + size >= lastFlushOffset + flushInterval;
       boolean appended;
-      if (i < lastIdx) {
+      if (i < lastIdx && !shouldFlush) {
         appended = stream.append(datum);
       } else if (i == lastIdx && nextWriteShouldFinalize) {
         appended = stream.appendAndFinalize(datum);
@@ -152,6 +160,9 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength) th
       if (appended) {
         bytesConsumed += size;
         writeOffset += size;
+        if (shouldFlush) {
+          lastFlushOffset = writeOffset;
+        }
       } 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.
@@ -171,6 +182,7 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength) th
 
   private void awaitResultFuture() throws IOException {
     try {
+      stream.awaitAckOf(writeOffset);
       stream.getResultFuture().get(10_717, TimeUnit.MILLISECONDS);
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();

google-cloud-storage/src/main/java/com/google/cloud/storage/BidiUploadState.java

@@ -327,9 +327,9 @@ abstract static class BaseUploadState extends BidiUploadState {
     protected @NonNull State state;
     protected @MonotonicNonNull BidiWriteObjectResponse lastResponseWithResource;
     protected @Nullable State stateToReturnToAfterRetry;
-    protected boolean finalFlushSignaled;
+    protected long finalFlushOffset;
     protected boolean finalFlushSent;
-    protected boolean finishWriteSignaled;
+    protected long finishWriteOffset;
     protected boolean finishWriteSent;
     protected @MonotonicNonNull OpenArguments lastOpenArguments;
     protected @Nullable SettableApiFuture<Void> pendingReconciliation;
@@ -356,6 +356,8 @@ private BaseUploadState(
       this.totalSentBytes = 0;
       this.confirmedBytes = -1;
       this.state = startingState;
+      this.finalFlushOffset = -1;
+      this.finishWriteOffset = -1;
     }
 
     @Override
@@ -463,17 +465,17 @@ public boolean finalFlush(long totalLength) {
 
         BidiWriteObjectRequest currentLast = peekLast();
         boolean equals = flush.equals(currentLast);
-        if (equals && finalFlushSignaled) {
+        if (equals && finalFlushOffset == totalLength) {
           return true;
         } else if (equals && lastSentRequestIndex == queue.size() - 1) {
-          finalFlushSignaled = true;
+          finalFlushOffset = totalLength;
           finalFlushSent = true;
           return true;
         }
 
         boolean offered = internalOffer(flush);
         if (offered) {
-          finalFlushSignaled = true;
+          finalFlushOffset = totalLength;
         }
         return offered;
       } finally {
@@ -561,7 +563,9 @@ final void updateStateFromResponse(BidiWriteObjectResponse response) {
           } else if (peek.hasOneof(FIRST_MESSAGE_DESCRIPTOR)) {
             poll();
           } else if (peek.getFlush()) {
-            if (finalFlushSent && persistedSize == totalSentBytes) {
+            if (finalFlushSent
+                && persistedSize == totalSentBytes
+                && persistedSize == finalFlushOffset) {
               setConfirmedBytes(persistedSize);
               signalTerminalSuccess = true;
               poll();
@@ -575,7 +579,9 @@ final void updateStateFromResponse(BidiWriteObjectResponse response) {
             checkState(
                 enqueuedBytes == 0,
                 "attempting to evict finish_write: true while bytes are still enqueued");
-            if (response.hasResource() && persistedSize == totalSentBytes) {
+            if (response.hasResource()
+                && persistedSize == totalSentBytes
+                && persistedSize == finishWriteOffset) {
               setConfirmedBytes(persistedSize);
               if (response.getResource().hasFinalizeTime()) {
                 signalTerminalSuccess = true;
@@ -697,7 +703,7 @@ final void retrying() {
     final boolean isFinalizing() {
       lock.lock();
       try {
-        return finishWriteSignaled && finishWriteSent;
+        return finishWriteOffset >= 0 && finishWriteSent;
       } finally {
         lock.unlock();
       }
@@ -753,7 +759,7 @@ final void sendVia(Consumer<BidiWriteObjectRequest> consumer) {
         if (prev != null) {
           if (prev.getFinishWrite()) {
             finishWriteSent = true;
-          } else if (prev.getFlush() && prev.getStateLookup() && finalFlushSignaled) {
+          } else if (prev.getFlush() && prev.getStateLookup() && finalFlushOffset > -1) {
             finalFlushSent = true;
           }
           consumer.accept(prev);
@@ -823,7 +829,7 @@ protected final void validateCurrentStateIsOneOf(State... allowed) {
 
     private void checkNotFinalizing() {
       checkState(
-          !finishWriteSignaled,
+          finishWriteOffset == -1,
           "Attempting to append bytes even though finalization has previously been signaled.");
     }
 
@@ -835,23 +841,26 @@ protected final boolean internalOffer(BidiWriteObjectRequest e) {
         }
         add = this::prepend;
       }
-      if (e.getFinishWrite()) {
-        finishWriteSignaled = true;
-      }
 
-      if (e.hasChecksummedData() && !finishWriteSignaled) {
+      boolean appended = false;
+      if (e.hasChecksummedData() && finishWriteOffset == -1) {
         ChecksummedData checksummedData = e.getChecksummedData();
         int size = checksummedData.getContent().size();
         if (size <= availableCapacity()) {
           totalSentBytes += size;
           add.accept(e);
-          return true;
+          appended = true;
         }
-        return false;
       } else {
         add.accept(e);
-        return true;
+        appended = true;
       }
+
+      if (e.getFinishWrite()) {
+        finishWriteOffset = totalSentBytes;
+      }
+
+      return appended;
     }
 
     @Nullable

google-cloud-storage/src/main/java/com/google/cloud/storage/BlobAppendableUploadConfig.java

@@ -26,6 +26,7 @@
 import com.google.cloud.storage.BidiUploadState.TakeoverAppendableUploadState;
 import com.google.cloud.storage.BlobAppendableUpload.AppendableUploadWriteableByteChannel;
 import com.google.cloud.storage.BlobAppendableUploadImpl.AppendableObjectBufferedWritableByteChannel;
+import com.google.cloud.storage.FlushPolicy.MinFlushSizeFlushPolicy;
 import com.google.cloud.storage.Storage.BlobWriteOption;
 import com.google.cloud.storage.TransportCompatibility.Transport;
 import com.google.cloud.storage.UnifiedOpts.ObjectTargetOpt;
@@ -266,9 +267,14 @@ BlobAppendableUpload create(GrpcStorageImpl storage, BlobInfo info, Opts<ObjectT
                     stream.awaitTakeoverStateReconciliation();
                     c =
                         new BidiAppendableUnbufferedWritableByteChannel(
-                            stream, chunkSegmenter, state.getConfirmedBytes());
+                            stream,
+                            chunkSegmenter,
+                            flushInterval(flushPolicy),
+                            state.getConfirmedBytes());
                   } else {
-                    c = new BidiAppendableUnbufferedWritableByteChannel(stream, chunkSegmenter, 0);
+                    c =
+                        new BidiAppendableUnbufferedWritableByteChannel(
+                            stream, chunkSegmenter, flushInterval(flushPolicy), 0);
                   }
                   return new AppendableObjectBufferedWritableByteChannel(
                       flushPolicy.createBufferedChannel(c, /* blocking= */ false),
@@ -282,6 +288,15 @@ BlobAppendableUpload create(GrpcStorageImpl storage, BlobInfo info, Opts<ObjectT
             build, BidiBlobWriteSessionConfig.Factory.WRITE_OBJECT_RESPONSE_BLOB_INFO_DECODER));
   }
 
+  private static long flushInterval(FlushPolicy fp) {
+    if (fp instanceof MinFlushSizeFlushPolicy) {
+      MinFlushSizeFlushPolicy min = (MinFlushSizeFlushPolicy) fp;
+      return min.getMinFlushSize();
+    } else {
+      return fp.getMaxPendingBytes();
+    }
+  }
+
   private static final class AppendableSession
       extends ChannelSession<
           AppendableUploadState,

google-cloud-storage/src/test/java/com/google/cloud/storage/ITAppendableUploadFakeTest.java

@@ -489,6 +489,7 @@ public void testFlushMultipleSegments_failsHalfway_partialFlush() throws Excepti
                   3,
                   storage.storageDataClient.retryContextProvider.create()),
               smallSegmenter,
+              3,
               0);
       ChecksummedTestContent content = ChecksummedTestContent.of(ALL_OBJECT_BYTES, 0, 10);
       StorageChannelUtils.blockingEmptyTo(ByteBuffer.wrap(content.getBytes()), channel);
@@ -638,6 +639,7 @@ public void testFlushMultipleSegmentsTwice_firstSucceeds_secondFailsHalfway_part
                   3,
                   storage.storageDataClient.retryContextProvider.create()),
               smallSegmenter,
+              3,
               0);
       ChecksummedTestContent content1 = ChecksummedTestContent.of(ALL_OBJECT_BYTES, 0, 10);
       ChecksummedTestContent content2 = ChecksummedTestContent.of(ALL_OBJECT_BYTES, 10, 10);
@@ -791,7 +793,7 @@ public void testFlushMultipleSegments_200ResponsePartialFlushHalfway() throws Ex
               3,
               storage.storageDataClient.retryContextProvider.create());
       BidiAppendableUnbufferedWritableByteChannel channel =
-          new BidiAppendableUnbufferedWritableByteChannel(stream, smallSegmenter, 0);
+          new BidiAppendableUnbufferedWritableByteChannel(stream, smallSegmenter, 3, 0);
       ChecksummedTestContent content = ChecksummedTestContent.of(ALL_OBJECT_BYTES, 0, 10);
       StorageChannelUtils.blockingEmptyTo(ByteBuffer.wrap(content.getBytes()), channel);
       channel.nextWriteShouldFinalize();
@@ -1049,7 +1051,7 @@ private static void runTestFlushMultipleSegments(FakeStorage fake) throws Except
               3,
               storage.storageDataClient.retryContextProvider.create());
       BidiAppendableUnbufferedWritableByteChannel channel =
-          new BidiAppendableUnbufferedWritableByteChannel(stream, smallSegmenter, 0);
+          new BidiAppendableUnbufferedWritableByteChannel(stream, smallSegmenter, 32, 0);
       StorageChannelUtils.blockingEmptyTo(ByteBuffer.wrap(content.getBytes()), channel);
       channel.nextWriteShouldFinalize();
       channel.close();