refactor(gax): extract resumable upload coordinator and future helpers

Author: whowes

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java

@@ -110,6 +110,16 @@ ApiFuture<ResponseT> getFuture() {
 
   void start() {
     progressTracker.onStarted(uploadUrl, chunkSize);
+    registerCancellationListener();
+    try {
+      buffer.fill();
+      transmitChunk();
+    } catch (Throwable t) {
+      uploadResultFuture.setException(t);
+    }
+  }
+
+  private void registerCancellationListener() {
     uploadResultFuture.addListener(
         () -> {
           Future<?> inFlight = inFlightFuture;
@@ -118,12 +128,6 @@ void start() {
           }
         },
         MoreExecutors.directExecutor());
-    try {
-      buffer.fill();
-      transmitChunk();
-    } catch (Throwable t) {
-      uploadResultFuture.setException(t);
-    }
   }
 
   private void transmitChunk() {

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java

@@ -80,8 +80,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final ApiCallContext callContext;
   private final ScheduledExecutorService executor;
   private final ExponentialRetryAlgorithm recoveryAlgorithm;
-  private final ResumableUploadProgressTracker progressTracker =
-      new ResumableUploadProgressTracker();
+  private final ResumableUploadProgressTracker progressTracker;
   private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
 
   private volatile @Nullable InputStream payload;
@@ -120,7 +119,8 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
             queryStatusCallable,
             options,
             clientContext,
-            recoveryAlgorithm);
+            recoveryAlgorithm,
+            ResumableUploadProgressTracker.create());
     try {
       future.start(payloadSupplier);
     } catch (Throwable t) {
@@ -135,7 +135,8 @@ private ResumableUploadFutureImpl(
       UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       ResumableUploadOptions options,
       ClientContext clientContext,
-      ExponentialRetryAlgorithm recoveryAlgorithm) {
+      ExponentialRetryAlgorithm recoveryAlgorithm,
+      ResumableUploadProgressTracker progressTracker) {
     this.startFuture = checkNotNull(startFuture, "startFuture must not be null");
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
@@ -147,6 +148,7 @@ private ResumableUploadFutureImpl(
     this.callContext = clientContext.getDefaultCallContext();
     this.executor = checkNotNull(clientContext.getExecutor(), "executor must not be null");
     this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null");
+    this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
     this.inFlightFuture = startFuture;
     ResumableUploadProgressListener progressListener = options.getProgressListener();
     if (progressListener != null) {
@@ -155,10 +157,7 @@ private ResumableUploadFutureImpl(
   }
 
   private void start(InputStreamSupplier payloadSupplier) throws IOException {
-    Duration timeout = firstNonNull(options.getGlobalTimeout(), DEFAULT_GLOBAL_TIMEOUT);
-    ScheduledFuture<?> timeoutFuture =
-        executor.schedule(this::onTimeout, timeout.toMillis(), TimeUnit.MILLISECONDS);
-    resultFuture.addListener(() -> timeoutFuture.cancel(false), MoreExecutors.directExecutor());
+    setUpGlobalTimeout();
     this.payload = checkNotNull(payloadSupplier.get(), "payload supplier returned null");
     ApiFutures.addCallback(
         startFuture,
@@ -174,43 +173,10 @@ public void onSuccess(ResumableUploadSession session) {
               return;
             }
             ResumableUploadChunkCoordinator<ResponseT> coordinator =
-                new ResumableUploadChunkCoordinator<>(
-                    uploadChunkCallable,
-                    queryStatusCallable,
-                    sessionUrl,
-                    payload,
-                    chunkSize,
-                    callContext,
-                    recoveryAlgorithm,
-                    executor,
-                    progressTracker);
-            ApiFuture<ResponseT> uploadFuture = coordinator.getFuture();
-            synchronized (lock) {
-              if (inFlightFuture == null) {
-                return;
-              }
-              effectiveChunkSize = chunkSize;
-              uploadSessionUrl = sessionUrl;
-              inFlightFuture = uploadFuture;
+                createAndRegisterCoordinator(sessionUrl, chunkSize);
+            if (coordinator != null) {
+              coordinator.start();
             }
-            ApiFutures.addCallback(
-                uploadFuture,
-                new ApiFutureCallback<ResponseT>() {
-                  @Override
-                  public void onSuccess(ResponseT response) {
-                    succeed(response);
-                  }
-
-                  @Override
-                  public void onFailure(Throwable t) {
-                    if (t instanceof CancellationException) {
-                      return;
-                    }
-                    fail(t);
-                  }
-                },
-                MoreExecutors.directExecutor());
-            coordinator.start();
           }
 
           @Override
@@ -224,6 +190,55 @@ public void onFailure(Throwable t) {
         MoreExecutors.directExecutor());
   }
 
+  private void setUpGlobalTimeout() {
+    Duration timeout = firstNonNull(options.getGlobalTimeout(), DEFAULT_GLOBAL_TIMEOUT);
+    ScheduledFuture<?> timeoutFuture =
+        executor.schedule(this::onTimeout, timeout.toMillis(), TimeUnit.MILLISECONDS);
+    resultFuture.addListener(() -> timeoutFuture.cancel(false), MoreExecutors.directExecutor());
+  }
+
+  private @Nullable ResumableUploadChunkCoordinator<ResponseT> createAndRegisterCoordinator(
+      String sessionUrl, int chunkSize) {
+    ResumableUploadChunkCoordinator<ResponseT> coordinator =
+        new ResumableUploadChunkCoordinator<>(
+            uploadChunkCallable,
+            queryStatusCallable,
+            sessionUrl,
+            checkNotNull(payload),
+            chunkSize,
+            callContext,
+            recoveryAlgorithm,
+            executor,
+            progressTracker);
+    ApiFuture<ResponseT> uploadFuture = coordinator.getFuture();
+    synchronized (lock) {
+      if (inFlightFuture == null) {
+        return null;
+      }
+      effectiveChunkSize = chunkSize;
+      uploadSessionUrl = sessionUrl;
+      inFlightFuture = uploadFuture;
+    }
+    ApiFutures.addCallback(
+        uploadFuture,
+        new ApiFutureCallback<ResponseT>() {
+          @Override
+          public void onSuccess(ResponseT response) {
+            succeed(response);
+          }
+
+          @Override
+          public void onFailure(Throwable t) {
+            if (t instanceof CancellationException) {
+              return;
+            }
+            fail(t);
+          }
+        },
+        MoreExecutors.directExecutor());
+    return coordinator;
+  }
+
   private void onTimeout() {
     String message;
     if (getUploadSessionUrl() != null) {

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java

@@ -93,8 +93,12 @@ private void dispatchSafely(ResumableUploadProgress status) {
   @GuardedBy("lock")
   private ResumableUploadProgress currentStatus;
 
-  ResumableUploadProgressTracker() {
-    this.currentStatus = ResumableUploadProgress.newBuilder().build();
+  private ResumableUploadProgressTracker(ResumableUploadProgress initialStatus) {
+    this.currentStatus = checkNotNull(initialStatus, "initialStatus must not be null");
+  }
+
+  static ResumableUploadProgressTracker create() {
+    return new ResumableUploadProgressTracker(ResumableUploadProgress.newBuilder().build());
   }
 
   void addListener(ResumableUploadProgressListener listener, Executor executor) {

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadProgressTrackerTest.java

@@ -46,7 +46,7 @@ class ResumableUploadProgressTrackerTest {
 
   @Test
   void testInitialSnapshotOnSubscribe() {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     List<ResumableUploadProgress> statuses = new ArrayList<>();
     tracker.addListener(statuses::add, MoreExecutors.directExecutor());
 
@@ -60,7 +60,7 @@ void testInitialSnapshotOnSubscribe() {
 
   @Test
   void testProgressUpdates() {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     tracker.onStarted("https://upload.url/test", 256);
 
     List<Long> byteUpdates = new ArrayList<>();
@@ -79,7 +79,7 @@ void testProgressUpdates() {
 
   @Test
   void testListenerExceptionSafety_doesNotDisruptSubsequentUpdates() {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     tracker.onStarted("https://upload.url/test", 256);
 
     AtomicInteger errorCount = new AtomicInteger(0);
@@ -111,7 +111,7 @@ void testListenerExceptionSafety_doesNotDisruptSubsequentUpdates() {
 
   @Test
   void testTerminalState_noNotificationsAfterFinalized() {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     tracker.onStarted("https://upload.url/test", 512);
 
     List<String> received = new ArrayList<>();
@@ -135,7 +135,7 @@ void testTerminalState_noNotificationsAfterFinalized() {
 
   @Test
   void testListenerReentrancy_canQueryStatusInsideCallback() {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     AtomicReference<ResumableUploadProgress> observed = new AtomicReference<>();
 
     tracker.addListener(
@@ -149,7 +149,7 @@ void testListenerReentrancy_canQueryStatusInsideCallback() {
 
   @Test
   void testCustomExecutorDispatch() throws Exception {
-    ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
+    ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create();
     ExecutorService executor = Executors.newSingleThreadExecutor();
     try {
       CountDownLatch latch = new CountDownLatch(2);