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

fornwall · fornwall · commit a151ebf7c2b8 · 2026-07-19T16:25:00.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,41 @@ 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
+        // skips sending it, 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);
+      }
+    }
+
+    private boolean isCommitCancelled() {
+      synchronized (commitCancellationLock) {
+        return commitCancelled;
+      }
+    }
 
     ApiFuture<CommitResponse> commitAsync() {
       close();
@@ -430,6 +470,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;
     }
 
@@ -484,13 +534,21 @@ public void run() {
                 "Retrying commit operation with a new precommit token obtained from the previous"
                     + " CommitResponse");
           }
+          if (isCommitCancelled()) {
+            res.setException(
+                newSpannerException(
+                    ErrorCode.CANCELLED,
+                    "The commit was cancelled before the Commit RPC was sent"));
+            return;
+          }
           final CommitRequest commitRequest = requestBuilder.build();
           span.addAnnotation("Starting Commit");
           final ApiFuture<com.google.spanner.v1.CommitResponse> commitFuture;
           final ISpan opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span);
           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,135 @@ 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 commitAsyncSkipsCommitRpcWhenCancelledBeforeItIsSent() {
+    SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
+    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());
+
+    // The commit was already abandoned, so the Commit RPC should not be sent at all.
+    verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
+  }
+
+  @Test
+  public void commitSkipsCommitRpcWhenCallingThreadInterruptedBeforeItIsSent() {
+    SettableApiFuture<Transaction> beginTransaction = SettableApiFuture.create();
+    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());
+
+      // The commit was already abandoned, so the Commit RPC should not be sent at all.
+      verify(rpc, never()).commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap());
+    } 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. The Commit
+   * RPC is also stubbed, so that a commit that is unexpectedly sent gets a well-defined future
+   * instead of a {@code null} from Mockito, and fails the {@code verify(rpc, never())} assertions.
+   */
+  private void 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);
+    when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap()))
+        .thenReturn(SettableApiFuture.create());
+  }
+
+  /** 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() {
