diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java index 7b14bdd04741..c84a42d81a32 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java @@ -110,6 +110,16 @@ ApiFuture 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() { diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java index 85019e5f90d4..b434d73e5a43 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java @@ -80,8 +80,7 @@ final class ResumableUploadFutureImpl 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 resultFuture = SettableApiFuture.create(); private volatile @Nullable InputStream payload; @@ -120,7 +119,8 @@ static ResumableUploadFutureImpl create( queryStatusCallable, options, clientContext, - recoveryAlgorithm); + recoveryAlgorithm, + ResumableUploadProgressTracker.create()); try { future.start(payloadSupplier); } catch (Throwable t) { @@ -135,7 +135,8 @@ private ResumableUploadFutureImpl( UnaryCallable> 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 coordinator = - new ResumableUploadChunkCoordinator<>( - uploadChunkCallable, - queryStatusCallable, - sessionUrl, - payload, - chunkSize, - callContext, - recoveryAlgorithm, - executor, - progressTracker); - ApiFuture 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() { - @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 createAndRegisterCoordinator( + String sessionUrl, int chunkSize) { + ResumableUploadChunkCoordinator coordinator = + new ResumableUploadChunkCoordinator<>( + uploadChunkCallable, + queryStatusCallable, + sessionUrl, + checkNotNull(payload), + chunkSize, + callContext, + recoveryAlgorithm, + executor, + progressTracker); + ApiFuture uploadFuture = coordinator.getFuture(); + synchronized (lock) { + if (inFlightFuture == null) { + return null; + } + effectiveChunkSize = chunkSize; + uploadSessionUrl = sessionUrl; + inFlightFuture = uploadFuture; + } + ApiFutures.addCallback( + uploadFuture, + new ApiFutureCallback() { + @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) { diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java index 98354f972089..df2887effe20 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java +++ b/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) { diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadProgressTrackerTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadProgressTrackerTest.java index 1de6650e0902..3d238f4481dc 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadProgressTrackerTest.java +++ b/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 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 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 received = new ArrayList<>(); @@ -135,7 +135,7 @@ void testTerminalState_noNotificationsAfterFinalized() { @Test void testListenerReentrancy_canQueryStatusInsideCallback() { - ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker(); + ResumableUploadProgressTracker tracker = ResumableUploadProgressTracker.create(); AtomicReference 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);