fix(spanner): cancel the in-flight Commit RPC when the commit is abandoned

fornwall · fornwall · commit f30fbf53b60b · 2026-07-19T16:18:06.000+02:00
Signed-off-by: Fredrik Fornwall &lt;fredrik@fornwall.net&gt;
diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
@@ -187,6 +187,14 @@ public void removeListener(Runnable listener) {
     @GuardedBy("precommitTokenLock")
     private MultiplexedSessionPrecommitToken latestPrecommitToken;
 
+    private final Object commitCancellationLock = new Object();
+
+    @GuardedBy("commitCancellationLock")
+    private boolean commitCancelled;
+
+    @GuardedBy("commitCancellationLock")
+    private ApiFuture<com.google.spanner.v1.CommitResponse> inFlightCommitFuture;
+
     @GuardedBy("lock")
     private volatile SettableApiFuture<Void> finishedAsyncOperations = SettableApiFuture.create();
 
@@ -365,9 +373,7 @@ void commit() {
                     rpc.getCommitRetrySettings().getTotalTimeout().getSeconds() + 5,
                     TimeUnit.SECONDS);
       } catch (InterruptedException | TimeoutException e) {
-        if (commitFuture != null) {
-          commitFuture.cancel(true);
-        }
+        cancelInFlightCommit();
         if (e instanceof InterruptedException) {
           throw SpannerExceptionFactory.propagateInterrupt((InterruptedException) e);
         } else {
@@ -378,7 +384,35 @@ void commit() {
       }
     }
 
-    volatile ApiFuture<CommitResponse> commitFuture;
+    private void cancelInFlightCommit() {
+      ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
+      synchronized (commitCancellationLock) {
+        // Record the cancellation even if the RPC has not been sent yet, so that CommitRunnable
+        // cancels it as soon as it is sent instead of leaving it running server-side.
+        commitCancelled = true;
+        commitFuture = inFlightCommitFuture;
+      }
+      // Cancel outside the lock, as this runs the listeners of the future on the current thread.
+      if (commitFuture != null) {
+        commitFuture.cancel(true);
+      }
+    }
+
+    private void publishOrCancelInFlightCommit(
+        ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture) {
+      boolean cancelled;
+      synchronized (commitCancellationLock) {
+        // The commit can already have been abandoned while this RPC was being sent, in which case
+        // cancelInFlightCommit() did not see this future and it is cancelled here instead.
+        cancelled = commitCancelled;
+        if (!cancelled) {
+          inFlightCommitFuture = commitFuture;
+        }
+      }
+      if (cancelled) {
+        commitFuture.cancel(true);
+      }
+    }
 
     ApiFuture<CommitResponse> commitAsync() {
       close();
@@ -430,6 +464,16 @@ ApiFuture<CommitResponse> commitAsync() {
           new CommitRunnable(
               res, finishOps, builder, /* retryAttemptDueToCommitProtocolExtension= */ false),
           MoreExecutors.directExecutor());
+      // Cancelling the future that is returned to the caller must also cancel the Commit RPC, in
+      // the same way as the timeout/interrupt path in commit() does. This listener also runs when
+      // the commit finishes normally, hence the isCancelled() check.
+      res.addListener(
+          () -> {
+            if (res.isCancelled()) {
+              cancelInFlightCommit();
+            }
+          },
+          MoreExecutors.directExecutor());
       return res;
     }
 
@@ -491,6 +535,7 @@ public void run() {
           try (IScope ignore = tracer.withSpan(opSpan)) {
             commitFuture = rpc.commitAsync(commitRequest, getTransactionChannelHint());
           }
+          publishOrCancelInFlightCommit(commitFuture);
           session.markUsed(clock.instant());
           commitFuture.addListener(
               () -> {
diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
@@ -29,7 +30,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
 import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler;
@@ -195,6 +198,141 @@ public void testCommitWithClientContext() {
     assertEquals(clientContext, capturedOptions.clientContext());
   }
 
+  @Test
+  public void commitCancelsInFlightRpcWhenCallingThreadInterrupted() {
+    when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
+    TransactionContextImpl transaction =
+        TransactionContextImpl.newBuilder()
+            .setSession(session)
+            .setTransactionId(ByteString.copyFromUtf8("test-txn"))
+            .setOptions(Options.fromTransactionOptions())
+            .setRpc(rpc)
+            .setTracer(tracer)
+            .setSpan(span)
+            .build();
+    SettableApiFuture<CommitResponse> inFlightCommit = SettableApiFuture.create();
+    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
+        .thenAnswer(
+            invocation -> {
+              Thread.currentThread().interrupt();
+              return inFlightCommit;
+            });
+
+    try {
+      SpannerException e = assertThrows(SpannerException.class, transaction::commit);
+      assertEquals(ErrorCode.CANCELLED, e.getErrorCode());
+      assertTrue("in-flight Commit RPC was not cancelled", inFlightCommit.isCancelled());
+    } finally {
+      // Clear the interrupt flag so it cannot leak into other tests.
+      Thread.interrupted();
+    }
+  }
+
+  @Test
+  public void commitAsyncCancelsInFlightRpcWhenReturnedFutureIsCancelled() {
+    when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
+    TransactionContextImpl transaction =
+        TransactionContextImpl.newBuilder()
+            .setSession(session)
+            .setTransactionId(ByteString.copyFromUtf8("test-txn"))
+            .setOptions(Options.fromTransactionOptions())
+            .setRpc(rpc)
+            .setTracer(tracer)
+            .setSpan(span)
+            .build();
+    SettableApiFuture<CommitResponse> inFlightCommit = SettableApiFuture.create();
+    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
+        .thenReturn(inFlightCommit);
+
+    ApiFuture<com.google.cloud.spanner.CommitResponse> commitFuture = transaction.commitAsync();
+    assertTrue(commitFuture.cancel(true));
+
+    assertTrue("in-flight Commit RPC was not cancelled", inFlightCommit.isCancelled());
+  }
+
+  @Test
+  public void commitAsyncCancelsCommitRpcThatIsSentAfterCancellation() {
+    SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
+    SettableApiFuture<CommitResponse> inFlightCommit =
+        setUpTransactionThatCommitsAfterBeginTransaction(beginTransaction);
+    TransactionContextImpl transaction = newTransactionWithoutTransactionId();
+
+    ApiFuture<com.google.cloud.spanner.CommitResponse> commitFuture = transaction.commitAsync();
+    // The Commit RPC is only sent once BeginTransaction has finished, so cancelling here means
+    // that the commit is abandoned before there is any RPC to cancel.
+    assertTrue(commitFuture.cancel(true));
+    verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
+
+    beginTransaction.set(
+        Transaction.newBuilder().setId(ByteString.copyFromUtf8("test-txn")).build());
+
+    assertTrue(
+        "Commit RPC that was sent after the cancellation was not cancelled",
+        inFlightCommit.isCancelled());
+  }
+
+  @Test
+  public void commitCancelsCommitRpcThatIsSentAfterCallingThreadInterrupted() {
+    SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
+    SettableApiFuture<CommitResponse> inFlightCommit =
+        setUpTransactionThatCommitsAfterBeginTransaction(beginTransaction);
+    TransactionContextImpl transaction = newTransactionWithoutTransactionId();
+
+    try {
+      // Interrupting before commit() waits for the result means that it gives up while
+      // BeginTransaction is still pending, and therefore before the Commit RPC has been sent.
+      Thread.currentThread().interrupt();
+      SpannerException e = assertThrows(SpannerException.class, transaction::commit);
+      assertEquals(ErrorCode.CANCELLED, e.getErrorCode());
+      verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
+
+      // commit() restores the interrupt flag, and BeginTransaction finishes on a gax thread that
+      // does not have that flag. Clear it so the listeners below are not interrupted either.
+      Thread.interrupted();
+      beginTransaction.set(
+          Transaction.newBuilder().setId(ByteString.copyFromUtf8("test-txn")).build());
+
+      verify(rpc, times(1)).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
+      assertTrue(
+          "Commit RPC that was sent after the interrupt was not cancelled",
+          inFlightCommit.isCancelled());
+    } finally {
+      // Clear the interrupt flag so it cannot leak into other tests.
+      Thread.interrupted();
+    }
+  }
+
+  /**
+   * Sets up a session that only finishes BeginTransaction once the given future is set, and returns
+   * the future of the Commit RPC that is sent after that.
+   */
+  private SettableApiFuture<CommitResponse> setUpTransactionThatCommitsAfterBeginTransaction(
+      SettableApiFuture<Transaction> beginTransaction) {
+    when(session.getName()).thenReturn("projects/p/instances/i/databases/d/sessions/s");
+    when(session.beginTransactionAsync(
+            Mockito.any(Options.class),
+            Mockito.anyBoolean(),
+            Mockito.anyMap(),
+            Mockito.any(),
+            Mockito.any()))
+        .thenReturn(beginTransaction);
+    SettableApiFuture<CommitResponse> inFlightCommit = SettableApiFuture.create();
+    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
+        .thenReturn(inFlightCommit);
+    return inFlightCommit;
+  }
+
+  /** Returns a transaction that has to call BeginTransaction before it can commit. */
+  private TransactionContextImpl newTransactionWithoutTransactionId() {
+    return TransactionContextImpl.newBuilder()
+        .setSession(session)
+        .setOptions(Options.fromTransactionOptions())
+        .setRpc(rpc)
+        .setTracer(tracer)
+        .setSpan(span)
+        .build();
+  }
+
   @SuppressWarnings("unchecked")
   @Test
   public void usesPreparedTransaction() {
