fix: update grpc finalize on close resumable uploads to validate ack'd object size (#2572)

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Follow up to #2527

Add ITGapicUnbufferedFinalizeOnCloseResumableWritableByteChannelTest with failure scenarios

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Refactoring to clean up some request building lifecycle and variable scope

* fix: update grpc finalize on close resumable uploads to validate ack'd object size

Tests are passing

Author: BenWhitehead

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

@@ -20,11 +20,13 @@
 
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.api.gax.rpc.ClientStreamingCallable;
 import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
 import com.google.cloud.storage.UnbufferedWritableByteChannelSession.UnbufferedWritableByteChannel;
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import com.google.storage.v2.ChecksummedData;
 import com.google.storage.v2.ObjectChecksums;
@@ -33,11 +35,7 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.ClosedChannelException;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.concurrent.ExecutionException;
-import java.util.function.Consumer;
-import java.util.function.LongConsumer;
 import org.checkerframework.checker.nullness.qual.NonNull;
 
 final class GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel
@@ -55,22 +53,23 @@ final class GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel
   private boolean open = true;
   private boolean first = true;
   private boolean finished = false;
+  private volatile WriteObjectRequest lastWrittenRequest;
 
   GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
       SettableApiFuture<WriteObjectResponse> resultFuture,
       ChunkSegmenter chunkSegmenter,
       ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> write,
-      ResumableWrite requestFactory) {
-    String bucketName = requestFactory.bucketName();
+      WriteCtx<ResumableWrite> writeCtx) {
+    String bucketName = writeCtx.getRequestFactory().bucketName();
     this.resultFuture = resultFuture;
     this.chunkSegmenter = chunkSegmenter;
 
     GrpcCallContext internalContext =
         contextWithBucketName(bucketName, GrpcCallContext.createDefault());
     this.write = write.withDefaultCallContext(internalContext);
 
-    this.writeCtx = new WriteCtx<>(requestFactory);
-    this.responseObserver = new Observer(writeCtx.getConfirmedBytes()::set, resultFuture::set);
+    this.writeCtx = writeCtx;
+    this.responseObserver = new Observer(internalContext);
   }
 
   @Override
@@ -92,27 +91,24 @@ public boolean isOpen() {
 
   @Override
   public void close() throws IOException {
+    if (!open) {
+      return;
+    }
+    open = false;
     ApiStreamObserver<WriteObjectRequest> openedStream = openedStream();
-    if (!finished) {
-      WriteObjectRequest message = finishMessage();
-      try {
+    try {
+      if (!finished) {
+        WriteObjectRequest message = finishMessage();
+        lastWrittenRequest = message;
         openedStream.onNext(message);
-        openedStream.onCompleted();
         finished = true;
-      } catch (RuntimeException e) {
-        resultFuture.setException(e);
-        throw e;
-      }
-    } else {
-      try {
-        openedStream.onCompleted();
-      } catch (RuntimeException e) {
-        resultFuture.setException(e);
-        throw e;
       }
+      openedStream.onCompleted();
+      responseObserver.await();
+    } catch (RuntimeException e) {
+      resultFuture.setException(e);
+      throw e;
     }
-    open = false;
-    responseObserver.await();
   }
 
   private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, boolean finalize)
@@ -122,51 +118,54 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
     }
 
     ChunkSegment[] data = chunkSegmenter.segmentBuffers(srcs, srcsOffset, srcsLength);
-
-    List<WriteObjectRequest> messages = new ArrayList<>();
+    if (data.length == 0) {
+      return 0;
+    }
 
     ApiStreamObserver<WriteObjectRequest> openedStream = openedStream();
     int bytesConsumed = 0;
-    for (ChunkSegment datum : data) {
-      Crc32cLengthKnown crc32c = datum.getCrc32c();
-      ByteString b = datum.getB();
-      int contentSize = b.size();
-      long offset = writeCtx.getTotalSentBytes().getAndAdd(contentSize);
-      Crc32cLengthKnown cumulative =
-          writeCtx
-              .getCumulativeCrc32c()
-              .accumulateAndGet(crc32c, chunkSegmenter.getHasher()::nullSafeConcat);
-      ChecksummedData.Builder checksummedData = ChecksummedData.newBuilder().setContent(b);
-      if (crc32c != null) {
-        checksummedData.setCrc32C(crc32c.getValue());
-      }
-      WriteObjectRequest.Builder builder =
-          writeCtx
-              .newRequestBuilder()
-              .setWriteOffset(offset)
-              .setChecksummedData(checksummedData.build());
-      if (!datum.isOnlyFullBlocks()) {
-        builder.setFinishWrite(true);
-        if (cumulative != null) {
-          builder.setObjectChecksums(
-              ObjectChecksums.newBuilder().setCrc32C(cumulative.getValue()).build());
+    try {
+      for (int i = 0; i < data.length; i++) {
+        ChunkSegment datum = data[i];
+        Crc32cLengthKnown crc32c = datum.getCrc32c();
+        ByteString b = datum.getB();
+        int contentSize = b.size();
+        long offset = writeCtx.getTotalSentBytes().getAndAdd(contentSize);
+        Crc32cLengthKnown cumulative =
+            writeCtx
+                .getCumulativeCrc32c()
+                .accumulateAndGet(crc32c, chunkSegmenter.getHasher()::nullSafeConcat);
+        ChecksummedData.Builder checksummedData = ChecksummedData.newBuilder().setContent(b);
+        if (crc32c != null) {
+          checksummedData.setCrc32C(crc32c.getValue());
+        }
+        WriteObjectRequest.Builder builder = writeCtx.newRequestBuilder();
+        if (!first) {
+          builder.clearUploadId();
+          builder.clearWriteObjectSpec();
+          builder.clearObjectChecksums();
+        }
+        builder.setWriteOffset(offset).setChecksummedData(checksummedData.build());
+        if (!datum.isOnlyFullBlocks() || (finalize && i + 1 == data.length)) {
+          builder.setFinishWrite(true);
+          if (cumulative != null) {
+            builder.setObjectChecksums(
+                ObjectChecksums.newBuilder().setCrc32C(cumulative.getValue()).build());
+          }
+          finished = true;
         }
-        finished = true;
-      }
-
-      WriteObjectRequest build = possiblyPairDownRequest(builder, first).build();
-      first = false;
-      messages.add(build);
-      bytesConsumed += contentSize;
-    }
-    if (finalize && !finished) {
-      messages.add(finishMessage());
-      finished = true;
-    }
 
-    try {
-      for (WriteObjectRequest message : messages) {
-        openedStream.onNext(message);
+        WriteObjectRequest build = builder.build();
+        first = false;
+        lastWrittenRequest = build;
+        openedStream.onNext(build);
+        bytesConsumed += contentSize;
+      }
+      if (finalize && !finished) {
+        WriteObjectRequest finishMessage = finishMessage();
+        lastWrittenRequest = finishMessage;
+        openedStream.onNext(finishMessage);
+        finished = true;
       }
     } catch (RuntimeException e) {
       resultFuture.setException(e);
@@ -201,73 +200,104 @@ private ApiStreamObserver<WriteObjectRequest> openedStream() {
     return stream;
   }
 
-  /**
-   * Several fields of a WriteObjectRequest are only allowed on the "first" message sent to gcs,
-   * this utility method centralizes the logic necessary to clear those fields for use by subsequent
-   * messages.
-   */
-  private static WriteObjectRequest.Builder possiblyPairDownRequest(
-      WriteObjectRequest.Builder b, boolean firstMessageOfStream) {
-    if (firstMessageOfStream && b.getWriteOffset() == 0) {
-      return b;
-    }
-    if (b.getWriteOffset() > 0) {
-      b.clearWriteObjectSpec();
-    }
-
-    if (b.getWriteOffset() > 0 && !b.getFinishWrite()) {
-      b.clearObjectChecksums();
-    }
-    return b;
-  }
-
-  static class Observer implements ApiStreamObserver<WriteObjectResponse> {
+  class Observer implements ApiStreamObserver<WriteObjectResponse> {
 
-    private final LongConsumer sizeCallback;
-    private final Consumer<WriteObjectResponse> completeCallback;
+    private final GrpcCallContext context;
 
     private final SettableApiFuture<Void> invocationHandle;
     private volatile WriteObjectResponse last;
 
-    Observer(LongConsumer sizeCallback, Consumer<WriteObjectResponse> completeCallback) {
-      this.sizeCallback = sizeCallback;
-      this.completeCallback = completeCallback;
+    Observer(GrpcCallContext context) {
+      this.context = context;
       this.invocationHandle = SettableApiFuture.create();
     }
 
     @Override
     public void onNext(WriteObjectResponse value) {
-      // incremental update
-      if (value.hasPersistedSize()) {
-        sizeCallback.accept(value.getPersistedSize());
-      } else if (value.hasResource()) {
-        sizeCallback.accept(value.getResource().getSize());
-      }
       last = value;
     }
 
-    /**
-     * observed exceptions so far
-     *
-     * <ol>
-     *   <li>{@link com.google.api.gax.rpc.OutOfRangeException}
-     *   <li>{@link com.google.api.gax.rpc.AlreadyExistsException}
-     *   <li>{@link io.grpc.StatusRuntimeException}
-     * </ol>
-     */
     @Override
     public void onError(Throwable t) {
-      invocationHandle.setException(t);
+      if (t instanceof ApiException) {
+        // use StorageExceptions logic to translate from ApiException to our status codes ensuring
+        // things fall in line with our retry handlers.
+        // This is suboptimal, as it will initialize a second exception, however this is the
+        // unusual case, and it should not cause a significant overhead given its rarity.
+        StorageException tmp = StorageException.asStorageException((ApiException) t);
+        StorageException storageException =
+            ResumableSessionFailureScenario.toStorageException(
+                tmp.getCode(),
+                tmp.getMessage(),
+                tmp.getReason(),
+                ImmutableList.of(lastWrittenRequest),
+                null,
+                context,
+                t);
+        resultFuture.setException(storageException);
+        invocationHandle.setException(storageException);
+      } else {
+        resultFuture.setException(t);
+        invocationHandle.setException(t);
+      }
     }
 
     @Override
     public void onCompleted() {
-      if (last != null && last.hasResource()) {
-        completeCallback.accept(last);
+      boolean finalizing = lastWrittenRequest.getFinishWrite();
+      if (last == null) {
+        clientDetectedError(
+            ResumableSessionFailureScenario.toStorageException(
+                0,
+                "onComplete without preceding onNext, unable to determine success.",
+                "invalid",
+                ImmutableList.of(lastWrittenRequest),
+                null,
+                context,
+                null));
+      } else if (last.hasResource() /* && finalizing*/) {
+        long totalSentBytes = writeCtx.getTotalSentBytes().get();
+        long finalSize = last.getResource().getSize();
+        if (totalSentBytes == finalSize) {
+          ok(finalSize);
+        } else if (finalSize < totalSentBytes) {
+          clientDetectedError(
+              ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException(
+                  ImmutableList.of(lastWrittenRequest), last, context, null));
+        } else {
+          clientDetectedError(
+              ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException(
+                  ImmutableList.of(lastWrittenRequest), last, context, null));
+        }
+      } else if (!finalizing || last.hasPersistedSize()) { // unexpected incremental response
+        clientDetectedError(
+            ResumableSessionFailureScenario.toStorageException(
+                0,
+                "Unexpected incremental response for finalizing request.",
+                "invalid",
+                ImmutableList.of(lastWrittenRequest),
+                last,
+                context,
+                null));
+      } else {
+        clientDetectedError(
+            ResumableSessionFailureScenario.SCENARIO_0.toStorageException(
+                ImmutableList.of(lastWrittenRequest), last, context, null));
       }
+    }
+
+    private void ok(long persistedSize) {
+      writeCtx.getConfirmedBytes().set(persistedSize);
+      resultFuture.set(last);
       invocationHandle.set(null);
     }
 
+    private void clientDetectedError(StorageException storageException) {
+      open = false;
+      resultFuture.setException(storageException);
+      invocationHandle.setException(storageException);
+    }
+
     void await() {
       try {
         invocationHandle.get();

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

@@ -305,7 +305,7 @@ UnbufferedWritableByteChannelSession<WriteObjectResponse> build() {
                         Retrying::newCallContext);
                   } else {
                     return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
-                        result, getChunkSegmenter(), write, start);
+                        result, getChunkSegmenter(), write, new WriteCtx<>(start));
                   }
                 })
                 .andThen(StorageByteChannels.writable()::createSynchronized));
@@ -346,7 +346,7 @@ BufferedWritableByteChannelSession<WriteObjectResponse> build() {
                         Retrying::newCallContext);
                   } else {
                     return new GapicUnbufferedFinalizeOnCloseResumableWritableByteChannel(
-                        result, getChunkSegmenter(), write, start);
+                        result, getChunkSegmenter(), write, new WriteCtx<>(start));
                   }
                 })
                 .andThen(c -> new DefaultBufferedWritableByteChannel(bufferHandle, c))