chore: do a metadata lookup if we try to do a reconnect and don't have a generation

JesseLovelace · BenWhitehead · commit bee43087ec13 · 2025-03-14T17:38:20.000-04:00
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicBidiUnbufferedAppendableWriteableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicBidiUnbufferedAppendableWriteableByteChannel.java
@@ -22,7 +22,9 @@
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.api.gax.rpc.BidiStreamingCallable;
 import com.google.api.gax.rpc.ErrorDetails;
+import com.google.api.gax.rpc.NotFoundException;
 import com.google.api.gax.rpc.OutOfRangeException;
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.storage.ChunkSegmenter.ChunkSegment;
 import com.google.cloud.storage.Conversions.Decoder;
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
@@ -33,12 +35,15 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.FieldMask;
 import com.google.storage.v2.AppendObjectSpec;
 import com.google.storage.v2.BidiWriteHandle;
 import com.google.storage.v2.BidiWriteObjectRedirectedError;
 import com.google.storage.v2.BidiWriteObjectRequest;
 import com.google.storage.v2.BidiWriteObjectResponse;
 import com.google.storage.v2.ChecksummedData;
+import com.google.storage.v2.GetObjectRequest;
+import com.google.storage.v2.Object;
 import com.google.storage.v2.ObjectChecksums;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -61,6 +66,7 @@
 final class GapicBidiUnbufferedAppendableWritableByteChannel
     implements UnbufferedWritableByteChannel {
   private final BidiStreamingCallable<BidiWriteObjectRequest, BidiWriteObjectResponse> write;
+  private final UnaryCallable<GetObjectRequest, Object> get;
   private final RetrierWithAlg retrier;
   private final SettableApiFuture<BidiWriteObjectResponse> resultFuture;
   private final ChunkSegmenter chunkSegmenter;
@@ -87,16 +93,17 @@ final class GapicBidiUnbufferedAppendableWritableByteChannel
 
   GapicBidiUnbufferedAppendableWritableByteChannel(
       BidiStreamingCallable<BidiWriteObjectRequest, BidiWriteObjectResponse> write,
+      UnaryCallable<GetObjectRequest, Object> get,
       RetrierWithAlg retrier,
       SettableApiFuture<BidiWriteObjectResponse> resultFuture,
       ChunkSegmenter chunkSegmenter,
       BidiWriteCtx<BidiAppendableWrite> writeCtx,
       Supplier<GrpcCallContext> baseContextSupplier) {
     this.write = write;
+    this.get = get;
     this.retrier = retrier;
     this.resultFuture = resultFuture;
     this.chunkSegmenter = chunkSegmenter;
-
     this.writeCtx = writeCtx;
     this.responseObserver = new RedirectHandlingResponseObserver(new BidiObserver());
     this.baseContextSupplier = baseContextSupplier;
@@ -163,16 +170,24 @@ void restart() {
     if (!resultFuture.isDone()) {
       ApiStreamObserver<BidiWriteObjectRequest> requestStream1 =
           openedStream(reconnectArguments.getCtx());
-      requestStream1.onNext(req);
-      lastWrittenRequest = req;
-      responseObserver.await();
-      first = false;
+      if (req != null) {
+        requestStream1.onNext(req);
+        lastWrittenRequest = req;
+        responseObserver.await();
+        first = false;
+      } else {
+        // This means we did a metadata lookup and determined that GCS never received the initial
+        // WriteObjectSpec,
+        // So we can just start over and send it again
+        first = true;
+      }
     }
   }
 
   public void startAppendableTakeoverStream() {
     BidiWriteObjectRequest req =
         writeCtx.newRequestBuilder().setFlush(true).setStateLookup(true).build();
+    generation.set(req.getAppendObjectSpec().getGeneration());
     this.messages = Collections.singletonList(req);
     flush();
     first = false;
@@ -259,9 +274,9 @@ private BidiWriteObjectRequest finishMessage() {
     Crc32cLengthKnown crc32cValue = writeCtx.getCumulativeCrc32c().get();
 
     BidiWriteObjectRequest.Builder b = writeCtx.newRequestBuilder();
-    if (!first) {
-      b.clearUploadId().clearObjectChecksums().clearWriteObjectSpec().clearAppendObjectSpec();
-    }
+
+    b.clearUploadId().clearObjectChecksums().clearWriteObjectSpec().clearAppendObjectSpec();
+
     b.setFinishWrite(true).setWriteOffset(offset);
     if (crc32cValue != null) {
       b.setObjectChecksums(ObjectChecksums.newBuilder().setCrc32C(crc32cValue.getValue()).build());
@@ -299,6 +314,7 @@ private void flush() {
           try {
             ApiStreamObserver<BidiWriteObjectRequest> opened = openedStream(context);
             for (BidiWriteObjectRequest message : this.messages) {
+
               opened.onNext(message);
               lastWrittenRequest = message;
             }
@@ -509,6 +525,7 @@ public void onCompleted() {
 
     private void ok(BidiWriteObjectResponse value) {
       last = value;
+      first = false;
       sem.release();
     }
 
@@ -685,6 +702,42 @@ ReconnectArguments getReconnectArguments() {
       long generation = this.generation.get();
       if (generation > 0) {
         spec.setGeneration(generation);
+      } else {
+        GetObjectRequest req =
+            GetObjectRequest.newBuilder()
+                .setBucket(spec.getBucket())
+                .setObject(spec.getObject())
+                .setReadMask(
+                    FieldMask.newBuilder()
+                        .addPaths(Storage.BlobField.GENERATION.getGrpcName())
+                        .build())
+                .build();
+        boolean objectNotFound = false;
+        try {
+          retrier.run(
+              () -> {
+                this.generation.set(get.call(req).getGeneration());
+                return null;
+              },
+              Decoder.identity());
+        } catch (Throwable t) {
+          if (t.getCause() instanceof NotFoundException) {
+            objectNotFound = true;
+          } else {
+            t.addSuppressed(new AsyncStorageTaskException());
+            throw t;
+          }
+        }
+        generation = this.generation.get();
+        if (generation > 0) {
+          spec.setGeneration(generation);
+        } else if (objectNotFound) {
+          // If the object wasn't found, that means GCS never saw the initial WriteObjectSpec, which
+          // means we'll need
+          // to send it again. We can process this retry by just starting over again
+          return ReconnectArguments.of(
+              baseContextSupplier.get().withExtraHeaders(getHeaders()), null);
+        }
       }
 
       BidiWriteHandle bidiWriteHandle = this.bidiWriteHandle.get();
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicBidiWritableByteChannelSessionBuilder.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicBidiWritableByteChannelSessionBuilder.java
@@ -21,11 +21,14 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.BidiStreamingCallable;
+import com.google.api.gax.rpc.UnaryCallable;
 import com.google.cloud.storage.ChannelSession.BufferedWriteSession;
 import com.google.cloud.storage.Retrying.RetrierWithAlg;
 import com.google.cloud.storage.UnbufferedWritableByteChannelSession.UnbufferedWritableByteChannel;
 import com.google.storage.v2.BidiWriteObjectRequest;
 import com.google.storage.v2.BidiWriteObjectResponse;
+import com.google.storage.v2.GetObjectRequest;
+import com.google.storage.v2.Object;
 import com.google.storage.v2.ServiceConstants.Values;
 import java.nio.ByteBuffer;
 import java.util.function.BiFunction;
@@ -192,10 +195,9 @@ BufferedAppendableUploadBuilder buffered(BufferHandle bufferHandle) {
     }
 
     final class BufferedAppendableUploadBuilder {
-
       private final BufferHandle bufferHandle;
-
       private ApiFuture<BidiAppendableWrite> start;
+      private UnaryCallable<GetObjectRequest, Object> get;
 
       BufferedAppendableUploadBuilder(BufferHandle bufferHandle) {
         this.bufferHandle = bufferHandle;
@@ -210,6 +212,12 @@ BufferedAppendableUploadBuilder setStartAsync(ApiFuture<BidiAppendableWrite> sta
         return this;
       }
 
+      public BufferedAppendableUploadBuilder setGetCallable(
+          UnaryCallable<GetObjectRequest, Object> get) {
+        this.get = get;
+        return this;
+      }
+
       BufferedWritableByteChannelSession<BidiWriteObjectResponse> build() {
         // it is theoretically possible that the setter methods for the following variables could
         // be called again between when this method is invoked and the resulting function is
@@ -220,6 +228,8 @@ BufferedWritableByteChannelSession<BidiWriteObjectResponse> build() {
         ByteStringStrategy boundStrategy = byteStringStrategy;
         Hasher boundHasher = hasher;
         RetrierWithAlg boundRetrier = retrier;
+        UnaryCallable<GetObjectRequest, Object> boundGet =
+            requireNonNull(get, "get must be non null");
         return new BufferedWriteSession<>(
             requireNonNull(start, "start must be non null"),
             ((BiFunction<
@@ -229,6 +239,7 @@ BufferedWritableByteChannelSession<BidiWriteObjectResponse> build() {
                     (start, resultFuture) ->
                         new GapicBidiUnbufferedAppendableWritableByteChannel(
                             write,
+                            boundGet,
                             boundRetrier,
                             resultFuture,
                             new ChunkSegmenter(
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageImpl.java
@@ -1472,6 +1472,7 @@ public boolean shouldRetry(
                     }))
             .buffered(BufferHandle.allocate(bufferSize))
             .setStartAsync(startAppendableWrite)
+            .setGetCallable(storageClient.getObjectCallable())
             .build();
     DefaultBlobWriteSessionConfig.DecoratedWritableByteChannelSession<
             BufferedWritableByteChannel, BidiWriteObjectResponse>
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/ITAppendableUploadFakeTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/ITAppendableUploadFakeTest.java
