fix(spanner): prevent statement cancellation race in AbstractBaseUnitOfWork - #14283
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces an isCancelled boolean flag in AbstractBaseUnitOfWork to handle race conditions during statement cancellation, and fixes a synchronization issue in the future's listener. However, the review feedback highlights a critical bug introduced by this change: late or stray cancellation signals can set isCancelled to true when no statement is running, which will then immediately cancel the next statement executed on the same connection. The reviewer recommends assigning currentlyRunningStatementFuture before submitting the task to the executor to safely resolve the race condition without relying on a stateful flag.
| this.isCancelled = true; | ||
| if (currentlyRunningStatementFuture != null | ||
| && !currentlyRunningStatementFuture.isDone() | ||
| && !currentlyRunningStatementFuture.isCancelled()) { |
There was a problem hiding this comment.
Introducing the isCancelled boolean flag to resolve the race condition introduces a new bug: stray or late cancellations can poison and immediately cancel subsequent statements executed on the same connection/transaction.
The Bug Scenario:
- Statement A starts and completes successfully.
- The future listener runs, setting
currentlyRunningStatementFuture = nullandisCancelled = false. - A late/stray timeout task or cancellation signal for Statement A fires (which is common in concurrent environments where cancelling the timeout task itself has a slight delay).
cancel()is called. SincecurrentlyRunningStatementFutureisnull, it does not cancel any active future, but it setsthis.isCancelled = true.- Statement B is initiated on the same connection.
- In
executeStatementAsync,currentlyRunningStatementFutureis assigned the new future. Sincethis.isCancelledistrue,future.cancel(true)is immediately called, cancelling Statement B before it even runs.
Recommended Solution:
Instead of introducing a stateful isCancelled flag, a much cleaner and safer approach is to assign currentlyRunningStatementFuture before submitting the task to the executor (or before the asynchronous boundary).
By doing so:
currentlyRunningStatementFutureis guaranteed to be non-null as soon as the statement execution flow begins.- Any call to
cancel()will immediately find the active future and cancel it, eliminating the race condition entirely. - We avoid the need for the
isCancelledflag and the risk of poisoning subsequent statements.
…OfWork In `AbstractBaseUnitOfWork.executeStatementAsync`, statements were submitted asynchronously to `statementExecutor` before storing the resulting future in `currentlyRunningStatementFuture`. Under thread contention, the query RPC can reach the mock server and trigger `cancel()` before the future is assigned, causing `cancel()` to silently drop the cancellation request and resulting in statement timeouts in CI tests such as `StatementTimeoutTest#testCancelReadWriteAutocommitMultipleStatements`. Fix this by assigning `currentlyRunningStatementFuture` inside the synchronized block together with `statementExecutor.submit(...)`, ensuring any concurrent cancellation observes the active future without risk of poisoning subsequent statements.
02506d4 to
174a0dd
Compare
| MoreExecutors.directExecutor()); | ||
| final ApiFuture<T> future; | ||
| synchronized (this) { | ||
| ApiFuture<T> f = statementExecutor.submit(context.wrap(callable)); |
There was a problem hiding this comment.
Calling statementExecutor.submit(...) inside synchronized (this) introduces a lock-contention issue when StatementExecutorType.DIRECT_EXECUTOR is used (which is the default configuration for both Spanner JDBC and PGAdapter).
With DIRECT_EXECUTOR, MoreExecutors.newDirectExecutorService().submit(callable) executes callable.call() synchronously inline on the calling thread. Placing submit(...) inside synchronized (this) causes the calling thread to hold the unit-of-work lock for the entire duration of the Spanner query RPC.
We can fix the race condition without holding the lock during submit(...) by pre-registering a SettableApiFuture as the coordination future before submitting:
final SpannerAsyncExecutionException caller =
callType == CallType.ASYNC
? new SpannerAsyncExecutionException(statement.getStatement())
: null;
final SettableApiFuture<T> statementFuture = SettableApiFuture.create();
synchronized (this) {
this.currentlyRunningStatementFuture = statementFuture;
}
final ApiFuture<T> f;
try {
f = statementExecutor.submit(context.wrap(callable));
} catch (Throwable t) {
synchronized (this) {
if (this.currentlyRunningStatementFuture == statementFuture) {
this.currentlyRunningStatementFuture = null;
}
}
statementFuture.setException(t);
throw t;
}
final ApiFuture<T> future =
ApiFutures.catching(
f,
Throwable.class,
input -> {
if (caller != null) {
input.addSuppressed(caller);
}
throw SpannerExceptionFactory.asSpannerException(input);
},
MoreExecutors.directExecutor());
ApiFutures.addCallback(
future,
new ApiFutureCallback<T>() {
@Override
public void onFailure(Throwable t) {
statementFuture.setException(t);
}
@Override
public void onSuccess(T result) {
statementFuture.set(result);
}
},
MoreExecutors.directExecutor());
statementFuture.addListener(
() -> {
if (statementFuture.isCancelled()) {
future.cancel(true);
}
synchronized (AbstractBaseUnitOfWork.this) {
if (currentlyRunningStatementFuture == statementFuture) {
currentlyRunningStatementFuture = null;
}
}
if (isSingleUse()) {
endUnitOfWorkSpan();
}
},
MoreExecutors.directExecutor());
return statementFuture;There was a problem hiding this comment.
Thanks @olavloite! That makes complete sense regarding DIRECT_EXECUTOR and lock contention. Updated the implementation to pre-register SettableApiFuture under synchronization and submit the statement task outside the lock, propagating cancellation and listeners via callbacks as suggested.
…Future Pre-register a SettableApiFuture before submitting statement execution outside the monitor lock to avoid lock contention with direct executors while preventing cancellation race conditions.
🤖 I have created a release *beep* *boop* --- <details><summary>1.92.0</summary> ## [1.92.0](v1.91.0...v1.92.0) (2026-09-23) ### Features * **bigquery-jdbc:** add `EnableTimestampPicos` connection property and its plumbing ([#14284](#14284)) ([b4aa5ac](b4aa5ac)) * **bigquery-jdbc:** implement picosecond temporal math and formatting engine ([#14286](#14286)) ([2a9612a](2a9612a)) * **bigquery-jdbc:** support picosecond in REST JSON path and nested types ([#14334](#14334)) ([15ffe4a](15ffe4a)) * **bigquery-jdbc:** support picosecond in `PreparedStatement` parameters and batching ([#14373](#14373)) ([c1aac66](c1aac66)) * **bigquery-jdbc:** support picosecond timestamp in `ResultSetMetaData` and `DatabaseMetaData` ([#14358](#14358)) ([43acdd3](43acdd3)) * **bigquery-jdbc:** support picosecond timestamps in Arrow Storage Read API and nested types ([#14332](#14332)) ([b5d9aca](b5d9aca)) * **bigquery-jdbc:** support qualified project delimiter in `DefaultDataset` property ([#14240](#14240)) ([6e8d6c8](6e8d6c8)) * **bigquery:** accelerate row-based query() with Arrow wire format ([#14405](#14405)) ([8d12a8f](8d12a8f)) * **bigquery:** add ArrowDeserializer helper utility ([#13943](#13943)) ([d9a298b](d9a298b)) * **bigquery:** add ArrowQueryPageFetcher for Arrow query result pagination ([#14404](#14404)) ([615409f](615409f)) * **bigquery:** add ArrowQueryResult and ArrowQueryResultImpl for Arrow result streaming ([#13944](#13944)) ([a62fdf8](a62fdf8)) * **bigquery:** add Storage Read API slow-path fallback for row-based query() ([#14409](#14409)) ([26e568a](26e568a)) * **bigquery:** add zero-copy queryArrow API for Arrow VectorSchemaRoot streaming ([#14402](#14402)) ([b44ffe8](b44ffe8)) * **bigquery:** make BigQuery AutoCloseable with default no-op close method ([#14434](#14434)) ([00bf3de](00bf3de)) * **firestore:** add support for BSON types ([#13189](#13189)) ([8a123d9](8a123d9)) * **gax:** add ApiCallContext and request-level settings overloads to ResumableUploadCallable ([#14251](#14251)) ([e8cbd42](e8cbd42)) * **gax:** add globalTimeout settings field to ResumableUploadCallSettings ([#14253](#14253)) ([438cda6](438cda6)) * **gax:** add resumable upload error classification and retry algorithm ([#14419](#14419)) ([b70396d](b70396d)) * **gax:** add ResumableUploadCallable creation to Callables and HttpJsonCallableFactory ([#14242](#14242)) ([7de24de](7de24de)) * **gax:** implement baseline Callable and Future for resumable uploads ([#14241](#14241)) ([5a54db9](5a54db9)) * **generator:** add model flag and allowlist parser for resumable upload RPCs ([#14317](#14317)) ([acc1856](acc1856)) * **generator:** emit resumable upload client surface ([#14319](#14319)) ([a9fed00](a9fed00)) * **generator:** emit resumable upload settings and HttpJson upload stub ([#14321](#14321)) ([c122474](c122474)) * **generator:** enable resumable upload generation for showcase ([#14325](#14325)) ([f9ebd79](f9ebd79)) * **generator:** switch resumable upload specialized stubs to package private ([#14471](#14471)) ([0d4e875](0d4e875)) * **generator:** wire transport stub delegation to resumable upload stubs ([#14322](#14322)) ([cc4b980](cc4b980)) * **google/cloud/backupdr/v1beta:** add backupdr ([#14410](#14410)) ([a4a47da](a4a47da)) * **google/cloud/networkservices/v1beta1:** add networkservices ([#14407](#14407)) ([21c4955](21c4955)) * **pubsub:** add publish telemetry headers for publish attempt observability ([#14338](#14338)) ([c167ab8](c167ab8)) * **pubsub:** implement publish hedging to reduce tail latency ([#13735](#13735)) ([b302615](b302615)) * **spanner:** Support dynamic TLS certificate and key rotation for Spanner Omni ([#14456](#14456)) ([ffc745c](ffc745c)) * **storage/control:** add delete folder recursive sample ([#13642](#13642)) ([f4b1b46](f4b1b46)) * **storage/control:** add delete folder recursive sample ([#14397](#14397)) ([2c01d55](2c01d55)) ### Bug Fixes * **auth:** restore transportFactory upon deserialization in InternalAwsSecurityCredentialsSupplier ([#14340](#14340)) ([beea42f](beea42f)) * **bigquery-jdbc:** ensure row ordering in PCNT IT ([#14330](#14330)) ([a16f048](a16f048)) * **bigquery-jdbc:** fix htapi fallback due to permission logic ([#14418](#14418)) ([21e6dc8](21e6dc8)) * **bigquery-jdbc:** fix Timestamp assertions ([#14290](#14290)) ([533ba14](533ba14)) * **bigquery-jdbc:** handle null parameters in Storage Write API bulk inserts ([#14270](#14270)) ([dd2c41a](dd2c41a)), refs [#14066](#14066) * **bigquery-jdbc:** handle SQL NULLs in ResultSet primitive getters ([#14383](#14383)) ([8e464fe](8e464fe)), refs [#14371](#14371) * **bigquery:** default Arrow pagination stream location to US instead of global ([#14458](#14458)) ([2775eb1](2775eb1)) * **bigquery:** preserve page token and paginate correctly in Arrow query when maxResults is set ([#14469](#14469)) ([f5601f4](f5601f4)) * **bigquery:** use first page row count for Arrow query pagination offset ([#14466](#14466)) ([9d10dd0](9d10dd0)) * **bigtable:** don't notify config listeners while holding the manager lock ([#14294](#14294)) ([4426ccd](4426ccd)) * **bigtable:** fall back to classic path when per-RPC CallCredentials are set on session path ([#14477](#14477)) ([57bacb0](57bacb0)) * **bigtable:** fix abnormal session closures and scale-up in session pool ([#14431](#14431)) ([6361ecd](6361ecd)) * **biqguery:** fix undeclared QueryParameter wiring in QueryStatistics ([#14401](#14401)) ([64cf1d3](64cf1d3)) * **bom:** restore google-cloud-spanner-jdbc to libraries-bom ([#14362](#14362)) ([bc7be5e](bc7be5e)), refs [#14347](#14347) * **spanner:** honor maxAttempts and totalTimeout in streaming resume loop ([#14370](#14370)) ([305f47d](305f47d)) * **spanner:** only set snapshot isolation read timestamp for SI or optimistic txns in CloudClientExecutor ([#14346](#14346)) ([54c0d0f](54c0d0f)) * **spanner:** prevent statement cancellation race in AbstractBaseUnitOfWork ([#14283](#14283)) ([d9a8eef](d9a8eef)) * **spanner:** re-enable ITInstanceAdminTest on cloud-devel and cloud-staging ([#14281](#14281)) ([89a8268](89a8268)) ### Performance Improvements * **spanner:** stop re-parsing the request id on every RPC ([#14353](#14353)) ([46108f4](46108f4)) ### Documentation * Add a Http/Json Post-Quantum Cryptography Guide ([#13963](#13963)) ([fcc65b0](fcc65b0)) * **bigquery:** add QueryArrow code sample and document JDK 17+ JVM requirements ([#14437](#14437)) ([bd363f6](bd363f6)) </details> --- This PR was generated with [Release Please](https://github.com/googleapis/release-please). See [documentation](https://github.com/googleapis/release-please#release-please). --------- Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>
In
AbstractBaseUnitOfWork.executeStatementAsync, statements were submitted asynchronously tostatementExecutorbefore storing the resulting future incurrentlyRunningStatementFuture. Under thread contention, the query RPC can reach the mock server and triggercancel()before the future is assigned, causingcancel()to silently drop the cancellation request and resulting in statement timeouts in CI tests such asStatementTimeoutTest#testCancelReadWriteAutocommitMultipleStatements.Furthermore, submitting statements while holding the unit-of-work monitor lock causes lock contention when
StatementExecutorType.DIRECT_EXECUTORis used (such as in Spanner JDBC or PGAdapter), because direct executors execute synchronously on the calling thread.This change:
SettableApiFutureascurrentlyRunningStatementFutureunder synchronization before submitting to the executor, ensuring any immediatecancel()invocation is captured without racing.statementExecutoroutside the synchronized block to avoid holding the unit-of-work lock during direct synchronous execution.SettableApiFuture, and registers a listener on theSettableApiFutureto propagate cancellation to the underlying future and clearcurrentlyRunningStatementFuture.