Skip to content

Commit 0a1f445

Browse files
whoweswhowes
committed
refactor(gax): invert resumable upload future and coordinator ownership
1 parent 34d17dc commit 0a1f445

3 files changed

Lines changed: 115 additions & 53 deletions

File tree

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

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import com.google.api.core.ApiFutureCallback;
3636
import com.google.api.core.ApiFutures;
3737
import com.google.api.core.InternalApi;
38+
import com.google.api.core.SettableApiFuture;
3839
import com.google.api.gax.resumable.ChunkUploadRequest;
3940
import com.google.api.gax.resumable.ChunkUploadResponse;
4041
import com.google.api.gax.resumable.ResumableUploadStatus;
@@ -45,6 +46,7 @@
4546
import java.util.Arrays;
4647
import java.util.concurrent.CancellationException;
4748
import org.jspecify.annotations.NullMarked;
49+
import org.jspecify.annotations.Nullable;
4850

4951
/**
5052
* Coordinates chunk transmission steps of a resumable upload session.
@@ -64,32 +66,43 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
6466
private final byte[] buffer;
6567
private final int chunkSize;
6668
private final ApiCallContext callContext;
67-
private final ResumableUploadFutureImpl<ResponseT> sessionFuture;
69+
private final SettableApiFuture<ResponseT> uploadResultFuture = SettableApiFuture.create();
70+
private volatile @Nullable ApiFuture<?> inFlightFuture;
6871

6972
ResumableUploadChunkCoordinator(
7073
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
7174
String uploadUrl,
7275
InputStream payload,
7376
int chunkSize,
74-
ApiCallContext callContext,
75-
ResumableUploadFutureImpl<ResponseT> sessionFuture) {
77+
ApiCallContext callContext) {
7678
this.uploadChunkCallable =
7779
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
7880
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
7981
this.payload = checkNotNull(payload, "payload must not be null");
8082
this.chunkSize = chunkSize;
8183
this.callContext = checkNotNull(callContext, "callContext must not be null");
82-
this.sessionFuture = checkNotNull(sessionFuture, "sessionFuture must not be null");
8384
this.buffer = new byte[chunkSize];
8485
}
8586

87+
ApiFuture<ResponseT> getFuture() {
88+
return uploadResultFuture;
89+
}
90+
8691
void start() {
92+
uploadResultFuture.addListener(
93+
() -> {
94+
ApiFuture<?> inFlight = inFlightFuture;
95+
if (uploadResultFuture.isCancelled() && inFlight != null) {
96+
inFlight.cancel(true);
97+
}
98+
},
99+
MoreExecutors.directExecutor());
87100
transmitChunk(0L);
88101
}
89102

90103
private void transmitChunk(long currentOffset) {
91104
// Abort if the session was already completed or canceled.
92-
if (sessionFuture.isDone()) {
105+
if (uploadResultFuture.isDone()) {
93106
return;
94107
}
95108

@@ -98,7 +111,7 @@ private void transmitChunk(long currentOffset) {
98111
try {
99112
bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
100113
} catch (IOException e) {
101-
sessionFuture.fail(e);
114+
uploadResultFuture.setException(e);
102115
return;
103116
}
104117

@@ -126,22 +139,23 @@ private void transmitChunk(long currentOffset) {
126139
try {
127140
ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
128141
uploadChunkCallable.futureCall(chunkRequest, callContext);
129-
sessionFuture.setInFlightFuture(chunkFuture);
142+
if (!tryRegisterInFlightFuture(chunkFuture)) {
143+
return;
144+
}
130145

131-
// Asynchronously handle the response: complete, fail, or chain the next chunk.
132146
ApiFutures.addCallback(
133147
chunkFuture,
134148
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
135149
@Override
136150
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
137-
if (sessionFuture.isDone()) {
151+
if (uploadResultFuture.isDone()) {
138152
return;
139153
}
140154
long nextOffset = currentOffset + chunkLength;
141155
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
142-
sessionFuture.succeed(response.getResponse());
156+
uploadResultFuture.set(response.getResponse());
143157
} else if (isFinal) {
144-
sessionFuture.fail(
158+
uploadResultFuture.setException(
145159
new IllegalStateException(
146160
"Upload stream ended and final chunk was transmitted, but server returned"
147161
+ " incomplete status"));
@@ -152,15 +166,28 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
152166

153167
@Override
154168
public void onFailure(Throwable t) {
155-
if (t instanceof CancellationException || sessionFuture.isDone()) {
169+
if (t instanceof CancellationException || uploadResultFuture.isDone()) {
156170
return;
157171
}
158-
sessionFuture.fail(t);
172+
uploadResultFuture.setException(t);
159173
}
160174
},
161175
MoreExecutors.directExecutor());
162176
} catch (Throwable t) {
163-
sessionFuture.fail(t);
177+
uploadResultFuture.setException(t);
178+
}
179+
}
180+
181+
/**
182+
* Registers the in-flight future for possible cancellation, returning false if the upload was
183+
* already cancelled.
184+
*/
185+
private boolean tryRegisterInFlightFuture(ApiFuture<?> future) {
186+
this.inFlightFuture = future;
187+
if (uploadResultFuture.isCancelled()) {
188+
future.cancel(true);
189+
return false;
164190
}
191+
return true;
165192
}
166193
}

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

Lines changed: 30 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
7272

7373
private volatile @Nullable String uploadSessionUrl;
7474

75+
// Tracks the current operation's Future (start, chunk upload) to propagate cancellation.
7576
@GuardedBy("lock")
7677
private @Nullable ApiFuture<?> inFlightFuture;
7778

@@ -121,23 +122,36 @@ private void start() {
121122
new ApiFutureCallback<ResumableUploadSession>() {
122123
@Override
123124
public void onSuccess(ResumableUploadSession session) {
124-
if (resultFuture.isDone()) {
125-
return;
126-
}
127-
uploadSessionUrl = session.getUploadUrl();
125+
String sessionUrl = session.getUploadUrl();
128126
ResumableUploadChunkCoordinator<ResponseT> coordinator =
129127
new ResumableUploadChunkCoordinator<>(
130-
uploadChunkCallable,
131-
session.getUploadUrl(),
132-
payload,
133-
settings.getChunkSize(),
134-
callContext,
135-
ResumableUploadFutureImpl.this);
136-
try {
137-
coordinator.start();
138-
} catch (Throwable t) {
139-
fail(t);
128+
uploadChunkCallable, sessionUrl, payload, settings.getChunkSize(), callContext);
129+
ApiFuture<ResponseT> uploadFuture = coordinator.getFuture();
130+
synchronized (lock) {
131+
if (resultFuture.isDone()) {
132+
return;
133+
}
134+
uploadSessionUrl = sessionUrl;
135+
inFlightFuture = uploadFuture;
140136
}
137+
ApiFutures.addCallback(
138+
uploadFuture,
139+
new ApiFutureCallback<ResponseT>() {
140+
@Override
141+
public void onSuccess(ResponseT response) {
142+
succeed(response);
143+
}
144+
145+
@Override
146+
public void onFailure(Throwable t) {
147+
if (t instanceof CancellationException) {
148+
return;
149+
}
150+
fail(t);
151+
}
152+
},
153+
MoreExecutors.directExecutor());
154+
coordinator.start();
141155
}
142156

143157
@Override
@@ -151,33 +165,15 @@ public void onFailure(Throwable t) {
151165
MoreExecutors.directExecutor());
152166
}
153167

154-
/**
155-
* Registers the active in-flight future for cancellation. If this session future has already been
156-
* canceled, the supplied future is canceled immediately.
157-
*/
158-
void setInFlightFuture(ApiFuture<?> inFlightFuture) {
159-
boolean shouldCancel = false;
160-
synchronized (lock) {
161-
if (resultFuture.isDone()) {
162-
shouldCancel = resultFuture.isCancelled();
163-
} else {
164-
this.inFlightFuture = inFlightFuture;
165-
}
166-
}
167-
if (shouldCancel) {
168-
inFlightFuture.cancel(true);
169-
}
170-
}
171-
172-
void succeed(@Nullable ResponseT result) {
168+
private void succeed(@Nullable ResponseT result) {
173169
synchronized (lock) {
174170
inFlightFuture = null;
175171
}
176172
closePayload();
177173
resultFuture.set(result);
178174
}
179175

180-
void fail(Throwable t) {
176+
private void fail(Throwable t) {
181177
synchronized (lock) {
182178
inFlightFuture = null;
183179
}

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

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import static org.mockito.Mockito.withSettings;
4242

4343
import com.google.api.core.ApiFutures;
44+
import com.google.api.core.ForwardingApiFuture;
4445
import com.google.api.core.SettableApiFuture;
4546
import com.google.api.gax.resumable.ChunkUploadRequest;
4647
import com.google.api.gax.resumable.ChunkUploadResponse;
@@ -57,6 +58,7 @@
5758
import java.util.concurrent.CountDownLatch;
5859
import java.util.concurrent.ExecutionException;
5960
import java.util.concurrent.TimeUnit;
61+
import java.util.concurrent.atomic.AtomicReference;
6062
import org.junit.jupiter.api.BeforeEach;
6163
import org.junit.jupiter.api.Test;
6264
import org.junit.jupiter.api.extension.ExtendWith;
@@ -230,17 +232,54 @@ void testUploadCallable_cancelInFlight_haltsUpload() throws Exception {
230232
}
231233

232234
@Test
233-
void testUploadCallable_setInFlightFutureAfterCancel_immediatelyCancelsFuture() {
235+
void testUploadCallable_cancelBeforeStartCompletes_abortsChunkUpload() {
234236
SettableApiFuture<ResumableUploadSession> startFuture = SettableApiFuture.create();
235-
when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture);
237+
// Ignore cancellation on startFuture to simulate start completing concurrently with cancel()
238+
when(mockStartCallable.futureCall(any(), any()))
239+
.thenReturn(
240+
new ForwardingApiFuture<ResumableUploadSession>(startFuture) {
241+
@Override
242+
public boolean cancel(boolean mayInterruptIfRunning) {
243+
return false;
244+
}
245+
});
246+
236247
ResumableUploadFuture<String> future =
237248
callable.futureCall("resource-path", streamOf("data"), null);
238249
assertThat(future.cancel(true)).isTrue();
239250
assertThat(future.isCancelled()).isTrue();
240251

241-
SettableApiFuture<String> lateFuture = SettableApiFuture.create();
242-
((ResumableUploadFutureImpl<String>) future).setInFlightFuture(lateFuture);
243-
assertThat(lateFuture.isCancelled()).isTrue();
252+
startFuture.set(
253+
ResumableUploadSession.newBuilder().setUploadUrl("https://upload.url/late").build());
254+
verifyNoInteractions(mockChunkCallable);
255+
}
256+
257+
@Test
258+
void testUploadCallable_cancelDuringChunkDispatch_immediatelyCancelsChunkFuture() {
259+
stubStartSession("https://upload.url/cancel-dispatch");
260+
SettableApiFuture<ChunkUploadResponse<String>> lateChunkFuture = SettableApiFuture.create();
261+
AtomicReference<ResumableUploadFuture<String>> futureRef = new AtomicReference<>();
262+
when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
263+
.thenAnswer(
264+
inv -> {
265+
futureRef.get().cancel(true);
266+
return lateChunkFuture;
267+
});
268+
269+
// Defer startFuture completion until futureRef is populated
270+
SettableApiFuture<ResumableUploadSession> startFuture = SettableApiFuture.create();
271+
when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture);
272+
273+
ResumableUploadFuture<String> future =
274+
callable.futureCall("resource-path", streamOf("data"), null);
275+
futureRef.set(future);
276+
startFuture.set(
277+
ResumableUploadSession.newBuilder()
278+
.setUploadUrl("https://upload.url/cancel-dispatch")
279+
.build());
280+
281+
assertThat(future.isCancelled()).isTrue();
282+
assertThat(lateChunkFuture.isCancelled()).isTrue();
244283
}
245284

246285
@Test

0 commit comments

Comments
 (0)