Skip to content

refactor(gax): remove circular ref between resumable upload future and chunk coordinator #14421

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
whowes wants to merge 1 commit into from
+64 −53
Draft

refactor(gax): remove circular ref between resumable upload future and chunk coordinator #14421

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 28 additions & 15 deletions ...va/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import com.google.api.core.ApiFutureCallback;
import com.google.api.core.ApiFutures;
import com.google.api.core.InternalApi;
import com.google.api.core.SettableApiFuture;
import com.google.api.gax.resumable.ChunkUploadRequest;
import com.google.api.gax.resumable.ChunkUploadResponse;
import com.google.api.gax.resumable.ResumableUploadStatus;
Expand All @@ -45,6 +46,7 @@
import java.util.Arrays;
import java.util.concurrent.CancellationException;
import org.jspecify.annotations.NullMarked;
import org.jspecify.annotations.Nullable;

/**
* Coordinates chunk transmission steps of a resumable upload session.
Expand All @@ -64,32 +66,40 @@
private final byte[] buffer;
private final int chunkSize;
private final ApiCallContext callContext;
private final ResumableUploadFutureImpl<ResponseT> sessionFuture;
private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
private volatile @Nullable ApiFuture<?> currentChunkFuture;

Check warning on line 70 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java

View check run for this annotation

SonarQubeCloud / [gapic-generator-java-root] SonarCloud Code Analysis

Use a thread-safe type; adding "volatile" is not enough to make this field thread-safe. 


See more on https://sonarcloud.io/project/issues?id=googleapis_google-cloud-java_showcase&issues=AaC7_0ryzZY4JqY3nQ_v&open=AaC7_0ryzZY4JqY3nQ_v&pullRequest=14421

ResumableUploadChunkCoordinator(
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
String uploadUrl,
InputStream payload,
int chunkSize,
ApiCallContext callContext,
ResumableUploadFutureImpl<ResponseT> sessionFuture) {
ApiCallContext callContext) {
this.uploadChunkCallable =
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
this.payload = checkNotNull(payload, "payload must not be null");
this.chunkSize = chunkSize;
this.callContext = checkNotNull(callContext, "callContext must not be null");
this.sessionFuture = checkNotNull(sessionFuture, "sessionFuture must not be null");
this.buffer = new byte[chunkSize];
}

void start() {
ApiFuture<ResponseT> start() {
result.addListener(

@blakeli0 blakeli0 Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: rename result to something more readable, otherwise it's not easy to comprehend. There might be a more concise name for it, but IIUC, it's basically a uploadAllChunksResultFuture?

() -> {
ApiFuture<?> chunk = currentChunkFuture;
if (result.isCancelled() && chunk != null) {
chunk.cancel(true);
}
},
MoreExecutors.directExecutor());
transmitChunk(0L);
return result;
}

private void transmitChunk(long currentOffset) {

Check failure on line 100 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java

View check run for this annotation

SonarQubeCloud / [gapic-generator-java-root] SonarCloud Code Analysis

Refactor this method to reduce its Cognitive Complexity from 16 to the 15 allowed. 


See more on https://sonarcloud.io/project/issues?id=googleapis_google-cloud-java_showcase&issues=AaDFuw5TfAcdo_VXvqSO&open=AaDFuw5TfAcdo_VXvqSO&pullRequest=14421
// Abort if the session was already completed or canceled.
if (sessionFuture.isDone()) {
if (result.isDone()) {
return;
}

Expand All @@ -98,7 +108,7 @@
try {
bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
} catch (IOException e) {
sessionFuture.fail(e);
result.setException(e);
return;
}

Expand Down Expand Up @@ -126,22 +136,25 @@
try {
ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
uploadChunkCallable.futureCall(chunkRequest, callContext);
sessionFuture.setInFlightFuture(chunkFuture);
this.currentChunkFuture = chunkFuture;

@gemini-code-assist gemini-code-assist Bot Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There is a race condition where result can be cancelled after result.isDone() is checked at the beginning of transmitChunk, but before currentChunkFuture is assigned. In this scenario, the cancellation listener registered in start() will have already executed (finding currentChunkFuture to be null or a previous chunk), and the newly created chunkFuture will never be cancelled, leading to a leaked background upload task.

To prevent this, check if result has been cancelled immediately after assigning currentChunkFuture and cancel the chunk future if so.

      this.currentChunkFuture = chunkFuture;
      if (result.isCancelled()) {
        chunkFuture.cancel(true);
      }

@whowes whowes Sep 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed.

if (result.isCancelled()) {
chunkFuture.cancel(true);
return;
}

// Asynchronously handle the response: complete, fail, or chain the next chunk.
ApiFutures.addCallback(
chunkFuture,
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
@Override
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
if (sessionFuture.isDone()) {
if (result.isDone()) {
return;
}
long nextOffset = currentOffset + chunkLength;
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
sessionFuture.succeed(response.getResponse());
result.set(response.getResponse());
} else if (isFinal) {
sessionFuture.fail(
result.setException(
new IllegalStateException(
"Upload stream ended and final chunk was transmitted, but server returned"
+ " incomplete status"));
Expand All @@ -152,15 +165,15 @@

@Override
public void onFailure(Throwable t) {
if (t instanceof CancellationException || sessionFuture.isDone()) {
if (t instanceof CancellationException || result.isDone()) {
return;
}
sessionFuture.fail(t);
result.setException(t);
}
},
MoreExecutors.directExecutor());
} catch (Throwable t) {
sessionFuture.fail(t);
result.setException(t);
}
}
}
60 changes: 36 additions & 24 deletions ...orm-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,16 +128,46 @@ public void onSuccess(ResumableUploadSession session) {
ResumableUploadChunkCoordinator<ResponseT> coordinator =
new ResumableUploadChunkCoordinator<>(
uploadChunkCallable,
session.getUploadUrl(),
uploadSessionUrl,
payload,
settings.getChunkSize(),
callContext,
ResumableUploadFutureImpl.this);
callContext);
ApiFuture<ResponseT> uploadFuture;
try {
coordinator.start();
uploadFuture = coordinator.start();
} catch (Throwable t) {
fail(t);
return;
}
boolean alreadyDone = false;
synchronized (lock) {
if (resultFuture.isDone()) {

@blakeli0 blakeli0 Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess the only possibility that resultFuture is already done at this moment is that it is cancelled?

alreadyDone = true;
} else {
inFlightFuture = uploadFuture;

@blakeli0 blakeli0 Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the purpose of inFlightFuture to keep track of the in flight future so we know which future to cancel when cancel() is called?

In general, I feel we have a lot of complex logics just to support correct cancellation of this future. Maybe we can restructure the code to make it simpler.

}
}
if (alreadyDone) {
uploadFuture.cancel(true);
return;
}
Comment on lines +143 to 153

@gemini-code-assist gemini-code-assist Bot Sep 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

With the removal of setInFlightFuture, there is a race condition where resultFuture can be cancelled while coordinator.start() is executing. If this happens, resultFuture.isDone() will be true when entering the synchronized block, and the method will return early without cancelling the newly started uploadFuture. This can leak the upload process in the background.

To fix this, check if resultFuture was cancelled when it is done, and propagate the cancellation to uploadFuture accordingly.

            boolean shouldCancel = false;
            synchronized (lock) {
              if (resultFuture.isDone()) {
                shouldCancel = resultFuture.isCancelled();
              } else {
                inFlightFuture = uploadFuture;
              }
            }
            if (shouldCancel) {
              uploadFuture.cancel(true);
              return;
            }
References
  1. When concurrent operations (such as cancellation and lazy initialization) are protected by a common lock (e.g., synchronized (this)), atomic state transitions (like compareAndSet) are not strictly necessary as the synchronization already prevents concurrent interleaving.

@whowes whowes Sep 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed.

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());
}

@Override
Expand All @@ -151,33 +181,15 @@ public void onFailure(Throwable t) {
MoreExecutors.directExecutor());
}

/**
* Registers the active in-flight future for cancellation. If this session future has already been
* canceled, the supplied future is canceled immediately.
*/
void setInFlightFuture(ApiFuture<?> inFlightFuture) {
boolean shouldCancel = false;
synchronized (lock) {
if (resultFuture.isDone()) {
shouldCancel = resultFuture.isCancelled();
} else {
this.inFlightFuture = inFlightFuture;
}
}
if (shouldCancel) {
inFlightFuture.cancel(true);
}
}

void succeed(@Nullable ResponseT result) {
private void succeed(@Nullable ResponseT result) {
synchronized (lock) {
inFlightFuture = null;
}
closePayload();
resultFuture.set(result);
}

void fail(Throwable t) {
private void fail(Throwable t) {
synchronized (lock) {
inFlightFuture = null;
}
Expand Down
14 changes: 0 additions & 14 deletions ...va/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,20 +229,6 @@ void testUploadCallable_cancelInFlight_haltsUpload() throws Exception {
assertThrows(CancellationException.class, future::get);
}

@Test
void testUploadCallable_setInFlightFutureAfterCancel_immediatelyCancelsFuture() {
SettableApiFuture<ResumableUploadSession> startFuture = SettableApiFuture.create();
when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture);
ResumableUploadFuture<String> future =
callable.futureCall("resource-path", streamOf("data"), null);
assertThat(future.cancel(true)).isTrue();
assertThat(future.isCancelled()).isTrue();

SettableApiFuture<String> lateFuture = SettableApiFuture.create();
((ResumableUploadFutureImpl<String>) future).setInFlightFuture(lateFuture);
assertThat(lateFuture.isCancelled()).isTrue();
}

@Test
void testUploadCallable_startFailure_failsFuture() {
when(mockStartCallable.futureCall(any(), any()))
Expand Down
Loading