Repository: googleapis/java-storage
Pull request #1645 (2022-09-22 16:08:44)

chore: wire automatic retries into usages of StartResumableUpload


Patch:
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
index b1e9e625..a513baeb 100644
--- 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
@@ -63,6 +63,7 @@ import com.google.cloud.storage.UnifiedOpts.ObjectSourceOpt;
 import com.google.cloud.storage.UnifiedOpts.ObjectTargetOpt;
 import com.google.cloud.storage.UnifiedOpts.Opts;
 import com.google.cloud.storage.UnifiedOpts.ProjectId;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.BaseEncoding;
@@ -310,14 +311,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions> implements Stora
           },
           this::getBlob);
     } else {
-      ApiFuture<ResumableWrite> start =
-          ResumableMedia.gapic()
-              .write()
-              .resumableWrite(
-                  storageClient
-                      .startResumableWriteCallable()
-                      .withDefaultCallContext(grpcCallContext),
-                  req);
+      ApiFuture<ResumableWrite> start = startResumableWrite(grpcCallContext, req);
       BufferedWritableByteChannelSession<WriteObjectResponse> session =
           channelSessionBuilder
               .resumable()
@@ -352,12 +346,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions> implements Stora
         opts.grpcMetadataMapper().apply(GrpcCallContext.createDefault());
     WriteObjectRequest req = getWriteObjectRequest(blobInfo, opts);
 
-    ApiFuture<ResumableWrite> start =
-        ResumableMedia.gapic()
-            .write()
-            .resumableWrite(
-                storageClient.startResumableWriteCallable().withDefaultCallContext(grpcCallContext),
-                req);
+    ApiFuture<ResumableWrite> start = startResumableWrite(grpcCallContext, req);
 
     BufferedWritableByteChannelSession<WriteObjectResponse> session =
         ResumableMedia.gapic()
@@ -744,14 +733,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions> implements Stora
         storageClient.writeObjectCallable(),
         getOptions(),
         retryAlgorithmManager.idempotent(),
-        () ->
-            ResumableMedia.gapic()
-                .write()
-                .resumableWrite(
-                    storageClient
-                        .startResumableWriteCallable()
-                        .withDefaultCallContext(grpcCallContext),
-                    req));
+        () -> startResumableWrite(grpcCallContext, req));
   }
 
   @Override
@@ -1385,6 +1367,20 @@ final class GrpcStorageImpl extends BaseService<StorageOptions> implements Stora
         .build();
   }
 
+  @VisibleForTesting
+  ApiFuture<ResumableWrite> startResumableWrite(
+      GrpcCallContext grpcCallContext, WriteObjectRequest req) {
+    Set<StatusCode.Code> codes =
+        GrpcStorageImpl.resultRetryAlgorithmToCodes(retryAlgorithmManager.getFor(req));
+    return ResumableMedia.gapic()
+        .write()
+        .resumableWrite(
+            storageClient
+                .startResumableWriteCallable()
+                .withDefaultCallContext(grpcCallContext.withRetryableCodes(codes)),
+            req);
+  }
+
   /**
    * When using the retry features of the Gapic client, we are only allowed to provide a {@link
    * Set}{@code <}{@link StatusCode.Code}{@code >}. Given {@link StatusCode.Code} is an enum, we can
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
index 426f5756..02e502dc 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java
@@ -28,6 +28,7 @@ import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.retrying.StreamResumptionStrategy;
 import com.google.api.gax.rpc.HeaderProvider;
+import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.auth.Credentials;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.ServiceFactory;
@@ -118,8 +119,9 @@ public final class GrpcStorageOptions extends StorageOptions
               .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
               .build());
     }
+    RetrySettings baseRetrySettings = getRetrySettings();
     RetrySettings readRetrySettings =
-        getRetrySettings()
+        baseRetrySettings
             .toBuilder()
             // when performing a read via ReadObject, the ServerStream will have a default relative
             // deadline set of `requestStartTime() + totalTimeout`, meaning if the specified
@@ -130,13 +132,23 @@ public final class GrpcStorageOptions extends StorageOptions
             // on idleTimeout below.
             .setLogicalTimeout(Duration.ofDays(28))
             .build();
-    Duration totalTimeout = getRetrySettings().getTotalTimeout();
-    // all retries for unary methods are handled at a different level
+    Duration totalTimeout = baseRetrySettings.getTotalTimeout();
+    Set<Code> startResumableWriteRetryableCodes =
+        builder.startResumableWriteSettings().getRetryableCodes();
+
+    // retries for unary methods are generally handled at a different level, except
+    // StartResumableWrite
     builder.applyToAllUnaryMethods(
         input -> {
           input.setSimpleTimeoutNoRetries(totalTimeout);
           return null;
         });
+
+    // configure the settings for StartResumableWrite
+    builder
+        .startResumableWriteSettings()
+        .setRetrySettings(baseRetrySettings)
+        .setRetryableCodes(startResumableWriteRetryableCodes);
     // for ReadObject we are configuring the server stream handling to do its own retries, so wire
     // things through. Retryable codes will be controlled closer to the use site as idempotency
     // considerations need to be made.


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/GrpcStorageImplUploadRetryTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/GrpcStorageImplUploadRetryTest.java
index 94dcca80..6d0ffc6f 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/GrpcStorageImplUploadRetryTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/GrpcStorageImplUploadRetryTest.java
@@ -19,6 +19,8 @@ package com.google.cloud.storage;
 import static com.google.cloud.storage.ByteSizeConstants._2MiB;
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.cloud.storage.GapicUnbufferedWritableByteChannelTest.DirectWriteService;
 import com.google.cloud.storage.Storage.BlobTargetOption;
 import com.google.cloud.storage.Storage.BlobWriteOption;
@@ -29,6 +31,7 @@ import com.google.storage.v2.Object;
 import com.google.storage.v2.ObjectChecksums;
 import com.google.storage.v2.StartResumableWriteRequest;
 import com.google.storage.v2.StartResumableWriteResponse;
+import com.google.storage.v2.StorageGrpc.StorageImplBase;
 import com.google.storage.v2.WriteObjectRequest;
 import com.google.storage.v2.WriteObjectResponse;
 import com.google.storage.v2.WriteObjectSpec;
@@ -129,6 +132,44 @@ public final class GrpcStorageImplUploadRetryTest {
     assertThat(service.returnError.get()).isFalse();
   }
 
+  @Test
+  public void startResumableWrite() throws Exception {
+
+    AtomicBoolean returnError = new AtomicBoolean(true);
+    StorageImplBase service =
+        new StorageImplBase() {
+          @Override
+          public void startResumableWrite(
+              StartResumableWriteRequest request, StreamObserver<StartResumableWriteResponse> obs) {
+            if (request.equals(Resumable.startReq)) {
+              if (returnError.get()) {
+                // clear the need to error. We only error on the first request.
+                returnError.compareAndSet(true, false);
+                obs.onError(TestUtils.apiException(Code.INTERNAL, "should retry"));
+              } else {
+                obs.onNext(Resumable.startResp);
+                obs.onCompleted();
+              }
+            } else {
+              obs.onError(
+                  TestUtils.apiException(Code.PERMISSION_DENIED, "Unexpected request chain."));
+            }
+          }
+        };
+
+    try (FakeServer server = FakeServer.of(service);
+        GrpcStorageImpl s = (GrpcStorageImpl) server.getGrpcStorageOptions().getService()) {
+      ApiFuture<ResumableWrite> f =
+          s.startResumableWrite(GrpcCallContext.createDefault(), Resumable.baseReq);
+      ResumableWrite resumableWrite = f.get();
+      StartResumableWriteResponse resp = resumableWrite.getRes();
+      assertThat(resp).isNotNull();
+      assertThat(resp.getUploadId()).isEqualTo(Resumable.uploadId);
+    }
+
+    assertThat(returnError.get()).isFalse();
+  }
+
   private static final class Direct {
     private static final Object obj =
         Object.newBuilder().setBucket(FORMATTED_BUCKET_NAME).setName("obj").build();
@@ -157,6 +198,7 @@ public final class GrpcStorageImplUploadRetryTest {
             (obs, reqs) -> {
               if (reqs.equals(ImmutableList.of(req1))) {
                 if (returnError.get()) {
+                  // clear the need to error. We only error on the first request.
                   returnError.compareAndSet(true, false);
                   obs.onError(TestUtils.apiException(Code.INTERNAL, "should retry"));
                 } else {
@@ -191,6 +233,8 @@ public final class GrpcStorageImplUploadRetryTest {
     private static final WriteObjectSpec spec =
         WriteObjectSpec.newBuilder().setResource(obj).setIfGenerationMatch(0).build();
 
+    private static final WriteObjectRequest baseReq =
+        WriteObjectRequest.newBuilder().setWriteObjectSpec(spec).build();
     private static final StartResumableWriteRequest startReq =
         StartResumableWriteRequest.newBuilder().setWriteObjectSpec(spec).build();
     private static final StartResumableWriteResponse startResp =
@@ -218,6 +262,7 @@ public final class GrpcStorageImplUploadRetryTest {
             (obs, reqs) -> {
               if (reqs.equals(ImmutableList.of(req1))) {
                 if (returnError.get()) {
+                  // clear the need to error. We only error on the first request.
                   returnError.compareAndSet(true, false);
                   obs.onError(TestUtils.apiException(Code.INTERNAL, "should retry"));
                 } else {