Repository: googleapis/java-storage
Pull request #2146 (2023-07-28 03:25:07)

fix: optimize resumable uploads to allow sending bytes during finalization

Add new methods to UnbufferedWritableByteChannel to allow writing and closing in a single call
* writeAndClose(ByteBuffer)
* writeAndClose(ByteBuffer[])
* writeAndClose(ByteBuffer[], int, int)

Update grpc and json implementation to leverage new methods and to write and finalize in the same call.

DefaultBufferedWritableByteChannel will use the new methods as appropriate.




Patch:
diff --git a/google-cloud-storage/clirr-ignored-differences.xml b/google-cloud-storage/clirr-ignored-differences.xml
index 84d99070..c7189a7a 100644
--- a/google-cloud-storage/clirr-ignored-differences.xml
+++ b/google-cloud-storage/clirr-ignored-differences.xml
@@ -8,4 +8,11 @@
     <method>com.google.cloud.storage.BlobWriteSession blobWriteSession(com.google.cloud.storage.BlobInfo, com.google.cloud.storage.Storage$BlobWriteOption[])</method>
   </difference>
 
+  <!-- Not breaking, new method has a default implementation -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/storage/UnbufferedWritableByteChannelSession$UnbufferedWritableByteChannel</className>
+    <method>* writeAndClose(*)</method>
+  </difference>
+
 </differences>
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedWritableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedWritableByteChannel.java
index 8201fbc6..da138572 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedWritableByteChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedWritableByteChannel.java
@@ -57,6 +57,40 @@ final class ApiaryUnbufferedWritableByteChannel implements UnbufferedWritableByt
 
   @Override
   public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
+    return internalWrite(srcs, offset, length, false);
+  }
+
+  @Override
+  public long writeAndClose(ByteBuffer[] srcs, int offset, int length) throws IOException {
+    long write = internalWrite(srcs, offset, length, true);
+    close();
+    return write;
+  }
+
+  @Override
+  public boolean isOpen() {
+    return open;
+  }
+
+  @Override
+  public void close() throws IOException {
+    open = false;
+    if (!finished) {
+      try {
+        ResumableOperationResult<@Nullable StorageObject> operationResult =
+            session.put(RewindableContent.empty(), HttpContentRange.of(cumulativeByteCount));
+        long persistedSize = operationResult.getPersistedSize();
+        committedBytesCallback.accept(persistedSize);
+        result.set(operationResult.getObject());
+      } catch (Exception e) {
+        result.setException(e);
+        throw StorageException.coalesce(e);
+      }
+    }
+  }
+
+  private long internalWrite(ByteBuffer[] srcs, int offset, int length, boolean finalize)
+      throws ClosedChannelException {
     if (!open) {
       throw new ClosedChannelException();
     }
@@ -65,9 +99,13 @@ final class ApiaryUnbufferedWritableByteChannel implements UnbufferedWritableByt
     long newFinalByteOffset = cumulativeByteCount + available;
     final HttpContentRange header;
     ByteRangeSpec rangeSpec = ByteRangeSpec.explicit(cumulativeByteCount, newFinalByteOffset);
-    if (available % ByteSizeConstants._256KiB == 0) {
+    boolean quantumAligned = available % ByteSizeConstants._256KiB == 0;
+    if (quantumAligned && finalize) {
+      header = HttpContentRange.of(rangeSpec, newFinalByteOffset);
+      finished = true;
+    } else if (quantumAligned) {
       header = HttpContentRange.of(rangeSpec);
-    } else {
+    } else { // not quantum aligned, have to finalize
       header = HttpContentRange.of(rangeSpec, newFinalByteOffset);
       finished = true;
     }
@@ -87,26 +125,4 @@ final class ApiaryUnbufferedWritableByteChannel implements UnbufferedWritableByt
       throw StorageException.coalesce(e);
     }
   }
-
-  @Override
-  public boolean isOpen() {
-    return open;
-  }
-
-  @Override
-  public void close() throws IOException {
-    open = false;
-    if (!finished) {
-      try {
-        ResumableOperationResult<@Nullable StorageObject> operationResult =
-            session.put(RewindableContent.empty(), HttpContentRange.of(cumulativeByteCount));
-        long persistedSize = operationResult.getPersistedSize();
-        committedBytesCallback.accept(persistedSize);
-        result.set(operationResult.getObject());
-      } catch (Exception e) {
-        result.setException(e);
-        throw StorageException.coalesce(e);
-      }
-    }
-  }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
index 3514b7d6..f1b14d06 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedWritableByteChannel.java
@@ -160,8 +160,17 @@ final class DefaultBufferedWritableByteChannel implements BufferedWritableByteCh
 
   @Override
   public void close() throws IOException {
-    try (UnbufferedWritableByteChannel ignored = channel) {
-      flush();
+    if (enqueuedBytes()) {
+      ByteBuffer buffer = handle.get();
+      Buffers.flip(buffer);
+      channel.writeAndClose(buffer);
+      if (buffer.hasRemaining()) {
+        buffer.compact();
+      } else {
+        Buffers.clear(buffer);
+      }
+    } else {
+      channel.close();
     }
   }
 
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultStorageRetryStrategy.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultStorageRetryStrategy.java
index fb3ac380..cf85b740 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultStorageRetryStrategy.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultStorageRetryStrategy.java
@@ -145,7 +145,7 @@ final class DefaultStorageRetryStrategy implements StorageRetryStrategy {
     public RetryResult beforeEval(Exception exception) {
       if (exception instanceof IllegalArgumentException) {
         IllegalArgumentException illegalArgumentException = (IllegalArgumentException) exception;
-        if (illegalArgumentException.getMessage().equals("no JSON input found")) {
+        if ("no JSON input found".equals(illegalArgumentException.getMessage())) {
           return RetryResult.RETRY;
         }
       }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedWritableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedWritableByteChannel.java
index 951a52ed..ebe375b5 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedWritableByteChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedWritableByteChannel.java
@@ -63,12 +63,60 @@ final class GapicUnbufferedWritableByteChannel<
   }
 
   @Override
-  public long write(ByteBuffer[] srcs, int srcsOffset, int srcLength) throws IOException {
+  public long write(ByteBuffer[] srcs, int srcsOffset, int srcsLength) throws IOException {
+    return internalWrite(srcs, srcsOffset, srcsLength, false);
+  }
+
+  @Override
+  public long writeAndClose(ByteBuffer[] srcs, int srcsOffset, int srcsLength) throws IOException {
+    long write = internalWrite(srcs, srcsOffset, srcsLength, true);
+    close();
+    return write;
+  }
+
+  @Override
+  public boolean isOpen() {
+    return open;
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (!finished) {
+      long offset = writeCtx.getTotalSentBytes().get();
+      Crc32cLengthKnown crc32cValue = writeCtx.getCumulativeCrc32c().get();
+
+      WriteObjectRequest.Builder b =
+          writeCtx.newRequestBuilder().setFinishWrite(true).setWriteOffset(offset);
+      if (crc32cValue != null) {
+        b.setObjectChecksums(
+            ObjectChecksums.newBuilder().setCrc32C(crc32cValue.getValue()).build());
+      }
+      WriteObjectRequest message = b.build();
+      try {
+        flusher.close(message);
+        finished = true;
+      } catch (RuntimeException e) {
+        resultFuture.setException(e);
+        throw e;
+      }
+    } else {
+      flusher.close(null);
+    }
+    open = false;
+  }
+
+  @VisibleForTesting
+  WriteCtx<RequestFactoryT> getWriteCtx() {
+    return writeCtx;
+  }
+
+  private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, boolean finalize)
+      throws ClosedChannelException {
     if (!open) {
       throw new ClosedChannelException();
     }
 
-    ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcLength);
+    ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength);
 
     List<WriteObjectRequest> messages = new ArrayList<>();
 
@@ -91,7 +139,7 @@ final class GapicUnbufferedWritableByteChannel<
               .newRequestBuilder()
               .setWriteOffset(offset)
               .setChecksummedData(checksummedData.build());
-      if (!datum.isOnlyFullBlocks()) {
+      if (!datum.isOnlyFullBlocks() || finalize) {
         builder.setFinishWrite(true);
         if (cumulative != null) {
           builder.setObjectChecksums(
@@ -114,40 +162,4 @@ final class GapicUnbufferedWritableByteChannel<
 
     return bytesConsumed;
   }
-
-  @Override
-  public boolean isOpen() {
-    return open;
-  }
-
-  @Override
-  public void close() throws IOException {
-    if (!finished) {
-      long offset = writeCtx.getTotalSentBytes().get();
-      Crc32cLengthKnown crc32cValue = writeCtx.getCumulativeCrc32c().get();
-
-      WriteObjectRequest.Builder b =
-          writeCtx.newRequestBuilder().setFinishWrite(true).setWriteOffset(offset);
-      if (crc32cValue != null) {
-        b.setObjectChecksums(
-            ObjectChecksums.newBuilder().setCrc32C(crc32cValue.getValue()).build());
-      }
-      WriteObjectRequest message = b.build();
-      try {
-        flusher.close(message);
-        finished = true;
-      } catch (RuntimeException e) {
-        resultFuture.setException(e);
-        throw e;
-      }
-    } else {
-      flusher.close(null);
-    }
-    open = false;
-  }
-
-  @VisibleForTesting
-  WriteCtx<RequestFactoryT> getWriteCtx() {
-    return writeCtx;
-  }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcResumableSession.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcResumableSession.java
index 8de44fb6..2c60c9fb 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcResumableSession.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcResumableSession.java
@@ -91,6 +91,7 @@ final class GrpcResumableSession {
             if (query.getObject() != null) {
               return query;
             } else {
+              handle.get().clear();
               content.rewindTo(query.getPersistedSize());
             }
           }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageByteChannels.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageByteChannels.java
index c3ff8013..fb14d2e8 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageByteChannels.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageByteChannels.java
@@ -184,6 +184,22 @@ final class StorageByteChannels {
       return delegate.write(srcs, offset, length);
     }
 
+    @Override
+    public synchronized int writeAndClose(ByteBuffer src) throws IOException {
+      return delegate.writeAndClose(src);
+    }
+
+    @Override
+    public synchronized long writeAndClose(ByteBuffer[] srcs) throws IOException {
+      return delegate.writeAndClose(srcs);
+    }
+
+    @Override
+    public synchronized long writeAndClose(ByteBuffer[] srcs, int offset, int length)
+        throws IOException {
+      return delegate.writeAndClose(srcs, offset, length);
+    }
+
     @Override
     public boolean isOpen() {
       return delegate.isOpen();
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/UnbufferedWritableByteChannelSession.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/UnbufferedWritableByteChannelSession.java
index 8affde6b..d7a5fcef 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/UnbufferedWritableByteChannelSession.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/UnbufferedWritableByteChannelSession.java
@@ -28,12 +28,26 @@ interface UnbufferedWritableByteChannelSession<ResultT>
   interface UnbufferedWritableByteChannel extends WritableByteChannel, GatheringByteChannel {
     @Override
     default int write(ByteBuffer src) throws IOException {
-      return Math.toIntExact(write(new ByteBuffer[] {src}));
+      return Math.toIntExact(write(new ByteBuffer[] {src}, 0, 1));
     }
 
     @Override
     default long write(ByteBuffer[] srcs) throws IOException {
       return write(srcs, 0, srcs.length);
     }
+
+    default int writeAndClose(ByteBuffer src) throws IOException {
+      return Math.toIntExact(writeAndClose(new ByteBuffer[] {src}, 0, 1));
+    }
+
+    default long writeAndClose(ByteBuffer[] srcs) throws IOException {
+      return writeAndClose(srcs, 0, srcs.length);
+    }
+
+    default long writeAndClose(ByteBuffer[] srcs, int offset, int length) throws IOException {
+      long write = write(srcs, offset, length);
+      close();
+      return write;
+    }
   }
 }


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
index 23d5b984..2653d27f 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/DefaultBufferedWritableByteChannelTest.java
@@ -17,9 +17,11 @@
 package com.google.cloud.storage;
 
 import static com.google.cloud.storage.ChunkSegmenterTest.TestData.fmt;
+import static com.google.cloud.storage.TestUtils.xxd;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.fail;
 
 import com.google.cloud.storage.BufferedWritableByteChannelSession.BufferedWritableByteChannel;
 import com.google.cloud.storage.UnbufferedWritableByteChannelSession.UnbufferedWritableByteChannel;
@@ -36,6 +38,7 @@ import java.util.Arrays;
 import java.util.Deque;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import net.jqwik.api.Arbitraries;
 import net.jqwik.api.Arbitrary;
 import net.jqwik.api.Combinators;
@@ -343,6 +346,60 @@ public final class DefaultBufferedWritableByteChannelTest {
     assertThat(actual).isEqualTo(expected);
   }
 
+  @Example
+  @SuppressWarnings("JUnit5AssertionsConverter")
+  void callingCloseWithBufferedDataShouldCallWriteAndClose() throws IOException {
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+
+    AtomicBoolean closed = new AtomicBoolean(false);
+    UnbufferedWritableByteChannel delegate =
+        new UnbufferedWritableByteChannel() {
+          @Override
+          public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
+            fail("unexpected write(ByteBuffer[], int, int) call");
+            return 0;
+          }
+
+          @Override
+          public long writeAndClose(ByteBuffer[] srcs, int offset, int length) throws IOException {
+            long total = 0;
+            try (WritableByteChannel out = Channels.newChannel(baos)) {
+              for (ByteBuffer src : srcs) {
+                total += out.write(src);
+              }
+            }
+            closed.compareAndSet(false, true);
+            return total;
+          }
+
+          @Override
+          public boolean isOpen() {
+            return !closed.get();
+          }
+
+          @Override
+          public void close() throws IOException {
+            fail("unexpected close() call");
+          }
+        };
+    DefaultBufferedWritableByteChannel test =
+        new DefaultBufferedWritableByteChannel(BufferHandle.allocate(20), delegate);
+
+    byte[] bytes = DataGenerator.base64Characters().genBytes(10);
+    String expected = xxd(bytes);
+
+    int write = test.write(ByteBuffer.wrap(bytes));
+    assertThat(write).isEqualTo(10);
+
+    assertThat(closed.get()).isFalse();
+
+    test.close();
+
+    String actual = xxd(baos.toByteArray());
+    assertThat(actual).isEqualTo(expected);
+    assertThat(closed.get()).isTrue();
+  }
+
   @Property
   void bufferAllocationShouldOnlyHappenWhenNeeded(@ForAll("BufferSizes") WriteOps writeOps)
       throws IOException {
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/ITGapicUnbufferedWritableByteChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/ITGapicUnbufferedWritableByteChannelTest.java
index 08d228c2..3b2bf1bf 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/ITGapicUnbufferedWritableByteChannelTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/ITGapicUnbufferedWritableByteChannelTest.java
@@ -17,6 +17,7 @@
 package com.google.cloud.storage;
 
 import static com.google.cloud.storage.TestUtils.apiException;
+import static com.google.cloud.storage.TestUtils.getChecksummedData;
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.api.core.SettableApiFuture;
@@ -25,6 +26,7 @@ import com.google.api.gax.rpc.DataLossException;
 import com.google.api.gax.rpc.PermissionDeniedException;
 import com.google.cloud.storage.Retrying.RetryingDependencies;
 import com.google.cloud.storage.WriteCtx.WriteObjectRequestBuilderFactory;
+import com.google.cloud.storage.WriteFlushStrategy.Flusher;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -48,10 +50,13 @@ import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.logging.Logger;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.junit.Test;
 
 public final class ITGapicUnbufferedWritableByteChannelTest {
@@ -307,6 +312,52 @@ public final class ITGapicUnbufferedWritableByteChannelTest {
     assertThat(writeCtx.getConfirmedBytes().get()).isEqualTo(40);
   }
 
+  @Test
+  public void resumableUpload_finalizeWhenWriteAndCloseCalledEvenWhenQuantumAligned()
+      throws IOException, InterruptedException, ExecutionException {
+    int quantum = 10;
+    ChunkSegmenter segmenter =
+        new ChunkSegmenter(Hasher.noop(), ByteStringStrategy.copy(), 50, quantum);
+    SettableApiFuture<WriteObjectResponse> result = SettableApiFuture.create();
+
+    AtomicReference<List<WriteObjectRequest>> actualFlush = new AtomicReference<>();
+    WriteObjectRequest closeRequestSentinel =
+        WriteObjectRequest.newBuilder().setUploadId("sentinel").build();
+    AtomicReference<WriteObjectRequest> actualClose = new AtomicReference<>(closeRequestSentinel);
+    GapicUnbufferedWritableByteChannel<?> c =
+        new GapicUnbufferedWritableByteChannel<>(
+            result,
+            segmenter,
+            reqFactory,
+            (bucketName, committedTotalBytesCallback, onSuccessCallback) ->
+                new Flusher() {
+                  @Override
+                  public void flush(@NonNull List<WriteObjectRequest> segments) {
+                    actualFlush.compareAndSet(null, segments);
+                  }
+
+                  @Override
+                  public void close(@Nullable WriteObjectRequest req) {
+                    actualClose.compareAndSet(closeRequestSentinel, req);
+                  }
+                });
+
+    byte[] bytes = DataGenerator.base64Characters().genBytes(quantum);
+
+    long written = c.writeAndClose(ByteBuffer.wrap(bytes));
+    WriteObjectRequest expectedRequest =
+        WriteObjectRequest.newBuilder()
+            .setUploadId(uploadId)
+            .setChecksummedData(getChecksummedData(ByteString.copyFrom(bytes), Hasher.noop()))
+            .setFinishWrite(true)
+            .build();
+
+    assertThat(written).isEqualTo(10);
+    assertThat(actualFlush.get()).isEqualTo(ImmutableList.of(expectedRequest));
+    // calling close is okay, as long as the provided request is null
+    assertThat(actualClose.get()).isAnyOf(closeRequestSentinel, null);
+  }
+
   static class DirectWriteService extends StorageImplBase {
     private static final Logger LOGGER = Logger.getLogger(DirectWriteService.class.getName());
     private final BiConsumer<StreamObserver<WriteObjectResponse>, List<WriteObjectRequest>> c;