diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java index dd3581e00a87..55244cbb516c 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java @@ -71,6 +71,7 @@ class ChunkAttemptCallable implements Callable implements Callable implements Callable implements Callable implements Callable> attemptFuture, ApiCallContext attemptContext, RetryingFuture> currentRetryingFuture) { + progressTracker.onRecovering(lastFailure); + // Per GAX-R7: query uses sensible unary defaults trimmed to the remaining global deadline. long remainingNanos = deadlineNanos == Long.MAX_VALUE @@ -254,6 +261,7 @@ private void handleQuerySuccess( // Normal path: realign buffer to committedOffset, compact and top up. try { buffer.realignTo(committedOffset); + progressTracker.onOffsetReceived(committedOffset); } catch (Throwable e) { failAttempt(attemptFuture, e); return; 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 882c3f8e81d9..9a77ee962856 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 @@ -52,6 +52,7 @@ import java.io.InputStream; import java.time.Duration; import java.util.concurrent.CancellationException; +import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; @@ -94,6 +95,7 @@ final class ResumableUploadChunkCoordinator { private final int chunkSize; private final ApiCallContext callContext; private final ClientContext clientContext; + private final UploadProgressTracker progressTracker = new UploadProgressTracker(); private volatile @Nullable String uploadSessionUrl; private volatile @Nullable RewindableStreamBuffer buffer; @@ -166,6 +168,7 @@ public void onSuccess(ResumableUploadSession session) { } } uploadSessionUrl = session.getUploadUrl(); + progressTracker.onStarted(uploadSessionUrl); buffer = new RewindableStreamBuffer(payload, chunkSize, uploadSessionUrl); scheduleNextChunk(0L); } @@ -181,6 +184,14 @@ public void onFailure(Throwable t) { MoreExecutors.directExecutor()); } + void addProgressListener(ResumableUploadProgressListener listener, Executor executor) { + progressTracker.addListener(listener, executor); + } + + ResumableUploadStatus getStatus() { + return progressTracker.getStatus(); + } + private void onTimeout() { synchronized (lock) { if (done) { @@ -233,6 +244,7 @@ void cancel(boolean mayInterruptIfRunning) { if (inFlight != null) { inFlight.cancel(mayInterruptIfRunning); } + progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl); closePayload(); } @@ -257,11 +269,15 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) { } IOException closeError = closePayload(); if (error == null) { + long totalBytes = + buffer != null ? buffer.getBufferBaseOffset() + buffer.getPayloadLength() : 0L; + progressTracker.onFinalized(totalBytes); result.set(response); } else { if (closeError != null) { error.addSuppressed(closeError); } + progressTracker.onFailed(error, uploadSessionUrl); result.setException(error); } } @@ -367,6 +383,7 @@ private void transmitSingleChunk(long currentOffset) { chunkRequest, chunkCallContext, command, + progressTracker, deadlineNanos, clientContext.getClock()); @@ -386,6 +403,7 @@ public void onSuccess(ChunkUploadResponse response) { } } long nextOffset = streamBuffer.getBufferBaseOffset() + streamBuffer.getPayloadLength(); + progressTracker.onChunkUploaded(nextOffset); if (response.isComplete()) { finish(response.getResponse(), null); } else if (streamBuffer.isFinal()) { diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java index 4282ef89fa33..a09a542486c1 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java @@ -31,6 +31,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.BetaApi; +import java.util.concurrent.Executor; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -48,4 +49,18 @@ public interface ResumableUploadFuture extends ApiFuture { /** Returns the upload session URL, or {@code null} if session initiation is in progress. */ @Nullable String getUploadSessionUrl(); + + /** + * Registers a listener to receive progress and state transition notifications for this upload. + * + *

A snapshot of the current upload status is dispatched to the listener immediately upon + * subscription on the provided executor. Subsequent status updates are delivered in order. + * + * @param listener callback listener to receive progress notifications + * @param executor executor on which the listener callbacks are dispatched + */ + void addProgressListener(ResumableUploadProgressListener listener, Executor executor); + + /** Returns the current status snapshot of the upload session. */ + ResumableUploadStatus getStatus(); } 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 62454177e566..b7e8c912f9d3 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 @@ -93,6 +93,16 @@ static ResumableUploadFutureImpl create( return coordinator.getUploadSessionUrl(); } + @Override + public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) { + coordinator.addProgressListener(listener, executor); + } + + @Override + public ResumableUploadStatus getStatus() { + return coordinator.getStatus(); + } + @Override public void addListener(Runnable listener, Executor executor) { result.addListener(listener, executor); @@ -115,12 +125,12 @@ public boolean isDone() { } @Override - public ResponseT get() throws InterruptedException, ExecutionException { + public @Nullable ResponseT get() throws InterruptedException, ExecutionException { return result.get(); } @Override - public ResponseT get(long timeout, TimeUnit unit) + public @Nullable ResponseT get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { return result.get(timeout, unit); } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java index 8ebbd20c172b..6ed00729373c 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java @@ -130,7 +130,8 @@ void call_successfulChunk_setsAttemptFuture() throws Exception { "https://upload.url/test", request, callContext, - ResumableUploadCommand.UPLOAD); + ResumableUploadCommand.UPLOAD, + new UploadProgressTracker()); callable.setRetryingFuture(mockExternalFuture); ChunkUploadResponse callResult = callable.call(); @@ -178,7 +179,8 @@ void call_returnsWithoutBlocking_andPropagatesCancellation() throws Exception { "https://upload.url/test", request, callContext, - ResumableUploadCommand.UPLOAD); + ResumableUploadCommand.UPLOAD, + new UploadProgressTracker()); List listeners = new ArrayList<>(); doAnswer( @@ -251,7 +253,8 @@ void call_perAttemptDeadline_appliesRpcTimeoutToCallContext() throws Exception { "https://upload.url/test", request, callContext, - ResumableUploadCommand.UPLOAD); + ResumableUploadCommand.UPLOAD, + new UploadProgressTracker()); callable.setRetryingFuture(mockExternalFuture); callable.call(); @@ -292,7 +295,8 @@ void call_nonBlockingExecution_callingThreadMakesImmediateProgress() throws Exce "https://upload.url/test", request, callContext, - ResumableUploadCommand.UPLOAD); + ResumableUploadCommand.UPLOAD, + new UploadProgressTracker()); callable.setRetryingFuture(mockExternalFuture); diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java index 7de11039fcc7..8d637ae81c61 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java @@ -57,15 +57,20 @@ import java.io.InputStream; import java.nio.charset.StandardCharsets; import java.time.Duration; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import org.jspecify.annotations.Nullable; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -1067,6 +1072,281 @@ void testGlobalTimeout_coversStartSessionTimeout() throws Exception { assertThat(hungStartFuture.isCancelled()).isTrue(); } + @Test + void testProgressListener_snapshotOnSubscribe_postsExactlyOneImmediateUpdate() throws Exception { + SettableApiFuture hungStartFuture = SettableApiFuture.create(); + when(mockStartCallable.futureCall(any(), any())).thenReturn(hungStartFuture); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + List statuses = new CopyOnWriteArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + future.addProgressListener( + status -> { + statuses.add(status); + latch.countDown(); + }, + executor); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(statuses).hasSize(1); + ResumableUploadStatus snapshot = statuses.get(0); + assertThat(snapshot.getState()).isEqualTo(ResumableUploadStatus.State.STARTING); + assertThat(snapshot.getBytesUploaded()).isEqualTo(0L); + assertThat(snapshot.getUploadUrl()).isNull(); + } + + @Test + void testProgressListener_prescribedStateTransitions() throws Exception { + SettableApiFuture startFuture = SettableApiFuture.create(); + when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture); + + // 20 bytes with chunkSize = 8 -> 3 chunks: [0..8), [8..16), [16..20) + // Chunk 1 succeeds -> [0..8) + // Chunk 2 fails with Cat-2 400 + // Query succeeds -> committed offset = 8 + // Chunk 2 resend succeeds -> [8..16) + // Chunk 3 succeeds and finalizes -> [16..20) + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "done"))); + + when(mockQueryCallable.futureCall(any(QueryStatusRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + QueryStatusResponse.newBuilder() + .setComplete(false) + .setCommittedOffset(8L) + .setUploadStatus("active") + .build())); + + List receivedStatuses = new CopyOnWriteArrayList<>(); + CountDownLatch finalizedLatch = new CountDownLatch(1); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("01234567890123456789"), null); + future.addProgressListener( + status -> { + receivedStatuses.add(status); + if (status.getState() == ResumableUploadStatus.State.FINALIZED) { + finalizedLatch.countDown(); + } + }, + executor); + + startFuture.set( + ResumableUploadSession.newBuilder() + .setUploadUrl("https://upload.url/progress-transitions") + .setUploadStatus("active") + .build()); + + assertThat(future.get()).isEqualTo("done"); + assertThat(finalizedLatch.await(5, TimeUnit.SECONDS)).isTrue(); + + List states = new ArrayList<>(); + for (ResumableUploadStatus s : receivedStatuses) { + states.add(s.getState()); + } + + assertThat(states) + .containsAtLeast( + ResumableUploadStatus.State.STARTED, + ResumableUploadStatus.State.UPLOADING, + ResumableUploadStatus.State.RECOVERING, + ResumableUploadStatus.State.OFFSET_RECEIVED, + ResumableUploadStatus.State.FINALIZED) + .inOrder(); + + long lastBytes = 0; + for (ResumableUploadStatus s : receivedStatuses) { + assertThat(s.getBytesUploaded()).isAtLeast(lastBytes); + lastBytes = s.getBytesUploaded(); + if (s.getState() != ResumableUploadStatus.State.STARTING) { + assertThat(s.getUploadUrl()).isEqualTo("https://upload.url/progress-transitions"); + } + } + assertThat(lastBytes).isEqualTo(20L); + } + + @Test + void testProgressListener_throwingListener_doesNotBreakUpload() throws Exception { + stubStartSession("https://upload.url/throwing-listener"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "ok"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + future.addProgressListener( + status -> { + throw new RuntimeException("boom from listener"); + }, + executor); + + assertThat(future.get()).isEqualTo("ok"); + assertThat(future.isDone()).isTrue(); + } + + @Test + void testProgressListener_subscribingAfterCompletion_yieldsOneTerminalSnapshot() + throws Exception { + stubStartSession("https://upload.url/post-completion"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "completed-ok"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + assertThat(future.get()).isEqualTo("completed-ok"); + + List postStatuses = new CopyOnWriteArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + future.addProgressListener( + status -> { + postStatuses.add(status); + latch.countDown(); + }, + executor); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + executor.submit(() -> {}).get(5, TimeUnit.SECONDS); + assertThat(postStatuses).hasSize(1); + ResumableUploadStatus snapshot = postStatuses.get(0); + assertThat(snapshot.getState()).isEqualTo(ResumableUploadStatus.State.FINALIZED); + assertThat(snapshot.getUploadUrl()).isEqualTo("https://upload.url/post-completion"); + assertThat(snapshot.getBytesUploaded()).isEqualTo(5L); + } + + @Test + void testProgressListener_subscribingAfterFailure_yieldsOneTerminalFailedSnapshot() + throws Exception { + when(mockStartCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(401, StatusCode.Code.UNAUTHENTICATED))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + assertThrows(ExecutionException.class, future::get); + + List postStatuses = new CopyOnWriteArrayList<>(); + CountDownLatch latch = new CountDownLatch(1); + future.addProgressListener( + status -> { + postStatuses.add(status); + latch.countDown(); + }, + executor); + + assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue(); + executor.submit(() -> {}).get(5, TimeUnit.SECONDS); + assertThat(postStatuses).hasSize(1); + ResumableUploadStatus snapshot = postStatuses.get(0); + assertThat(snapshot.getState()).isEqualTo(ResumableUploadStatus.State.FAILED); + assertThat(snapshot.getException()).isInstanceOf(ApiException.class); + } + + @Test + void testProgressListener_futureCancelFromInsideListenerBody_worksWithoutDeadlock() + throws Exception { + stubStartSession("https://upload.url/cancel-inside-listener"); + SettableApiFuture> hungChunk = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())).thenReturn(hungChunk); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + CountDownLatch cancelAttemptedLatch = new CountDownLatch(1); + future.addProgressListener( + status -> { + if (status.getState() == ResumableUploadStatus.State.STARTED) { + future.cancel(true); + cancelAttemptedLatch.countDown(); + } + }, + executor); + + assertThat(cancelAttemptedLatch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(future.isCancelled()).isTrue(); + assertThat(hungChunk.isCancelled()).isTrue(); + } + + @Test + void testProgressListener_getStatus_reflectsCurrentState() throws Exception { + stubStartSession("https://upload.url/get-status"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "ok"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("ok"); + ResumableUploadStatus status = future.getStatus(); + assertThat(status.getState()).isEqualTo(ResumableUploadStatus.State.FINALIZED); + assertThat(status.getUploadUrl()).isEqualTo("https://upload.url/get-status"); + assertThat(status.getBytesUploaded()).isEqualTo(5L); + } + + @Test + void testProgressListener_orderingUnderConcurrency_pinsSequentialExecutor() throws Exception { + ExecutorService multiThreadedExecutor = Executors.newFixedThreadPool(8); + try { + stubStartSession("https://upload.url/concurrency-order"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "finished"))); + + List events = new CopyOnWriteArrayList<>(); + AtomicInteger concurrentExecutions = new AtomicInteger(0); + AtomicBoolean concurrencyDetected = new AtomicBoolean(false); + CountDownLatch finalizedLatch = new CountDownLatch(1); + + ResumableUploadFuture future = + callable.futureCall( + "resource-path", streamOf("0123456789012345678901234567890123456789"), null); + + future.addProgressListener( + status -> { + int inProgress = concurrentExecutions.incrementAndGet(); + if (inProgress > 1) { + concurrencyDetected.set(true); + } + try { + Thread.sleep(10); + events.add(status); + if (status.getState() == ResumableUploadStatus.State.FINALIZED) { + finalizedLatch.countDown(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + concurrentExecutions.decrementAndGet(); + } + }, + multiThreadedExecutor); + + assertThat(future.get(10, TimeUnit.SECONDS)).isEqualTo("finished"); + assertThat(finalizedLatch.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(concurrencyDetected.get()).isFalse(); + + long lastBytes = 0; + for (ResumableUploadStatus s : events) { + assertThat(s.getBytesUploaded()).isAtLeast(lastBytes); + lastBytes = s.getBytesUploaded(); + } + assertThat(lastBytes).isEqualTo(40L); + } finally { + multiThreadedExecutor.shutdownNow(); + } + } + private static class HttpStatusStatusCode implements StatusCode { private final int httpStatus; private final StatusCode.Code code; @@ -1138,7 +1418,7 @@ private static void assertChunk( } private static class TrackableStream extends ByteArrayInputStream { - int closeCount = 0; + volatile int closeCount = 0; TrackableStream(String content) { super(content.getBytes(StandardCharsets.UTF_8)); diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java index 0ca00f757eb6..9cf7f5607200 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java @@ -45,6 +45,9 @@ import com.google.common.util.concurrent.MoreExecutors; import java.io.ByteArrayInputStream; import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -114,6 +117,16 @@ public void close() throws IOException { callContext, clientContext); + List terminalStatuses = Collections.synchronizedList(new ArrayList<>()); + coordinator.addProgressListener( + status -> { + if (status.getState() == ResumableUploadStatus.State.FINALIZED + || status.getState() == ResumableUploadStatus.State.FAILED) { + terminalStatuses.add(status); + } + }, + MoreExecutors.directExecutor()); + AtomicInteger completionListenerCount = new AtomicInteger(0); result.addListener(completionListenerCount::incrementAndGet, MoreExecutors.directExecutor()); @@ -153,5 +166,6 @@ public void close() throws IOException { assertThat(result.isDone()).isTrue(); assertThat(closeCount.get()).isEqualTo(1); assertThat(completionListenerCount.get()).isEqualTo(1); + assertThat(terminalStatuses).hasSize(1); } }