Threading refactor (5/7): Close ordering & shutdown hardening by igorbernstein2 · Pull Request #5 · igorbernstein2/google-cloud-java

igorbernstein2 · 2026-06-16T02:21:42Z

Summary

Phase 5 of 7. The largest phase — six fixes flushed out by the new threading model, all around Client.close / shutdown ordering and race conditions.

The big one is the Client.close restructure (3 explicit phases: initiate close → await termination → tear down infrastructure). The rest are targeted hardening of specific races.

What's in here

Commit	What it does
`fix: drain SessionPools before tearing down userCallbackExecutor on Client.close` (+ addListener serialize, Duration cleanup)	`SessionPool` gains `awaitTerminated(Duration)`. `Client.close` becomes three explicit
phases: initiate graceful close, `awaitTerminated` per pool (6-min budget), tear down. Plus: `VOperationImpl.start` queues `grpcContext.addListener` through the op executor via `runInline` — without this, an
async-queued `onClose` could drain between the closed-check and `addListener`, leaking the listener per-RPC for the lifetime of `grpcContext`. Plus `closed` AtomicBoolean + idempotent `close()`.
`fix: deliver terminal onClose to Scheduled retries pending at Client.close`	`BigtableTimer` gains an `onStop` hook. `Scheduled.onStart` registers a hook on entry; `NettyWheelTimer.stop()` runs every hook synchronously
before discarding pending wheel timeouts. Phase 3 of `Client.close` reordered so `sessionTimer.stop()` runs before `userCallbackExecutor.close()`, giving the hook-fired `onClose` tasks a live op-executor to land on.
Without this, long-delay `Scheduled` retries (server-driven `RetryInfo.retryDelay`) outlived Phase 2 and were silently discarded.
`fix: serialize open*/close so racing opens cannot create orphan pools` (+ supersedes earlier sessionPools sync)	All four `open*` / `close` sites on `Client` now hold `sessionPools`' monitor across the `closed` check,
pool construction, and insertion. `closed` downgrades to a plain `boolean` (CAS provided no value under the monitor). Prevents an orphan pool — created after `close()` snapshotted `sessionPools` — from outliving the close
with its callbacks landing on shut-down executors.
`fix: ShimImpl uses shutdownAndAwait for userCallbackExecutor`	`ShimImpl` registered its `userCallbackExecutor` with a bare `shutdown()` closer, while `Client.create`'s path uses `shutdownAndAwait` (5s drain for
in-flight callbacks). Promote `Client.shutdownAndAwait` from private-static to public-static so `ShimImpl` can reuse the same semantics.
`fix: stop heartbeat-driven force-close once session enters CLOSING`	A session in `CLOSING` is on a terminal trajectory; a missed heartbeat there reclassifies `closeReason` from `USER` to `MISSED_HEARTBEAT` and pollutes
shutdown metrics. The stuck-vRPC backstop during shutdown is now `SessionPool.awaitTerminated` / `Client.POOL_DRAIN_TIMEOUT`. Also cancel `heartbeatTimeout` on entry to `CLOSING` (was: `WAIT_SERVER_CLOSE`).
`fix: lift shim loader throws to failed futures via SessionPoolMap`	After the `Client.close` hardening, the shim loader (`ReadRowShim` / `MutateRowShim`) throws `IllegalStateException` post-close; `getUnchecked` wraps
it in `UncheckedExecutionException`, surfacing an unchecked throw out of `readRow`/`mutateRow` instead of the failed `CompletableFuture` the async API contract promises. New `SessionPoolMap` wrapper unwraps the wrapper
exception for `get()` and converts loader throws into failed futures for `apply()`.

Stack position

Base: threading-refactor-phase-4
Next: threading-refactor-phase-6 — tracer correctness & chain.isDone()

nimf

It's funny, reading the tests was more helpful for understanding this than reading comments.

nimf · 2026-06-22T07:14:28Z

+  @Test
+  void awaitTerminatedReturnsTrueAfterSessionsDrain()
+      throws InterruptedException, ExecutionException, TimeoutException {
+    // Start a real session, issue + complete a vRPC so the session is fully open, then close


maybe also worth adding a test case when pool close is issued while vrpc is in-flight.

mutianf · 2026-06-18T20:51:00Z

+  private static final Logger logger = Logger.getLogger(Client.class.getName());
+
+  // Per-pool drain budget during close. One full watchdog tick (5 min) plus 1 min buffer; if a
+  // pool can't drain in that window, something is genuinely wrong on the server side and we give


Hmmm.. I don't understand this comment

The primary focus of this PR is to ensure that all of the async work is done before close the user executor (to make sure that all user notifications have a delivery path). The watchdog is the thing that prevent sessions from hanging indefinitely. Not sure how to re-word this

mutianf · 2026-06-22T13:16:36Z

                                    () ->
                                        otelContext
-                                            .wrap(() -> onStateChange(new Idle()))
+                                            .wrap(this::onTimerFired)


I forgot.. what is this timeout timer for?

its the expoential backoff timer. ie the delay between attempts

mutianf · 2026-06-22T13:19:54Z

+
+    // Invoked from BigtableTimer.stop on the close thread. Trampoline back to the op executor so
+    // currentState reads and onStateChange are still single-threaded with the rest of the chain.
+    private void onTimerStopping() {


Hmm I don't fully understand the stop hook. Is the timer stooped when it fired or cancelled? So this hook is called on timer fired and on cancel? So we unregister it in both places to clean up?

stop hook is when the timer is being closed so that listeners that have scheduled work can be notified that their task will never run and clean up

mutianf · 2026-06-22T13:22:35Z

+
+    private void unregisterStopHook() {
+      if (stopHook != null) {
+        stopHook.unregister();


what does unregister do? NettyWheelTimer already clears the hook from the map in onStop?

Timer has 2 callbacks: to run the task when its scheduled and separate callback for being notified that the timer is being closed. stophook.unregister removes the 2nd callback

we have 2 listeners here:

for being notified when we are scheduled to run

for being notified that the timer is being closed and 1 will never run.
They are registered separately so need to be removed separately

@nested

…lient.close Client.close shut down userCallbackExecutor before draining the SessionPools that depend on it, so late listener.onClose tasks from in-flight RPCs arrived after backing was dead and got RejectedExecutionException — silently stranding the user's terminal callbacks. The earlier fix sprinkled inline-drain fallbacks inside OpExecutor; restructure shutdown instead so the race can't happen. SessionPool gains awaitTerminated(Duration), backed by a CompletableFuture SessionPoolImpl completes from onSessionClose once the pool is CLOSED and the last session has drained. close() no longer kills the watchdog — awaitTerminated takes ownership of that, so the watchdog stays alive during shutdown and can escalate any session stuck in WAIT_SERVER_CLOSE longer than its tick interval (5 min) via forceClose. Client.close becomes three explicit phases: (1) initiate graceful close on each pool, (2) awaitTerminated on each with a 6-minute per-pool budget (one full watchdog tick plus buffer), (3) tear down userCallbackExecutor / channelPool / timers in the existing order, now safely because all listener.onClose tasks are queued or drained before backing dies. VOperationImpl.start queues grpcContext.addListener through the op executor via runInline. Without this, an async-queued onClose from chain.start (PendingVRpc pool-closed fast-fail, VRpcImpl deadline-exceeded short-circuit) could drain between the CleanupListener.closed read and addListener: CleanupListener.onClose would call removeListener as a no-op pre-registration, and the caller would then register a listener with nothing to remove it. The leak is per-RPC and permanent until grpcContext cancels — for long-lived application contexts it accumulates indefinitely. FIFO ordering through the op executor makes the closed-check sound: any onClose chain.start enqueued drains first, so the check is accurate by the time we evaluate it. Add a `closed` AtomicBoolean + checkNotClosed() guard on the three openers so concurrent opens during shutdown can't create pools the close path won't see. close() is now idempotent via CAS on that flag. Tests: - ClientTest#openAfterCloseThrows / closeIsIdempotent - SessionPoolImplTest awaitTerminated* coverage - VOperationImplTest covering async onClose / context cancel ordering - SessionPoolImplTest tearDown now calls awaitTerminated so the watchdog is closed before testTimer.stop races its self-reschedule - FakeSessionPool in TableBaseTest gains a no-op awaitTerminated stub Drive-by: remove the spurious @nested annotation from SessionPoolImplTest's top-level class. @nested is meaningful only on non-static inner classes; on the outer class it caused Surefire to mis-attribute test counts.

…close Scheduled RetryingVRpcs hold no session reference, so a long-delay retry (server-driven RetryInfo.retryDelay) outlives Phase 2 drain and is silently discarded when sessionTimer.stop() runs in Phase 3. The user's listener never fires. Add an onStop hook primitive to BigtableTimer. Scheduled.onStart registers a hook on entry and unregisters on every exit path (normal fire, cancel, hook fire). NettyWheelTimer.stop() runs every hook synchronously before discarding pending wheel timeouts; hooks trampoline back through the op executor to drive Scheduled to a CANCELLED Done. Reorder Client.close Phase 3 so sessionTimer.stop() runs before userCallbackExecutor.close(), giving the hook-fired onClose tasks a live op-executor backing to land on. Also replace Scheduled.onStart's dead RejectedExecutionException catch with IllegalStateException, matching BigtableTimer.stop()'s documented post-condition.

openTableAsync / openAuthorizedViewAsync / openMaterializedViewAsync read 'closed' lock-free, then constructed the pool, then inserted it into sessionPools — close() could CAS closed=true and snapshot sessionPools in between, leaving the new pool orphaned: never closed, its callbacks landing on shut-down executors. Hold sessionPools' monitor across the closed check, the construction, and the insert. Opens are infrequent (typically once per table at app startup) so the monitor cost is negligible. Move close()'s closed flip inside the same monitor too. With every access now under the lock, downgrade 'closed' from AtomicBoolean to a plain boolean — the CAS provided no value over a plain read+write under the lock.

ShimImpl registered its userCallbackExecutor with a bare shutdown() closer, while Client.create's path uses shutdownAndAwait (which gives in-flight listener.onClose tasks a 5-second drain window before shutdownNow). On ShimImpl close, queued callbacks were abandoned mid-flight — fine for quiescent shutdowns but a regression for fast-close patterns (test boundaries, dynamic config reloads) where in-flight callbacks have not yet drained. Promote Client.shutdownAndAwait from private-static to public-static so ShimImpl (different package) can reuse the same shutdown semantics, and update ShimImpl to call it.

ReadRowShim and MutateRowShim cache per-target handles in a Guava LoadingCache whose loader calls Client.openTableAsync (and friends). After the Client-close hardening, that loader throws IllegalStateException post-close — and getUnchecked wraps it in UncheckedExecutionException, so callers see an unchecked exception thrown out of readRow / mutateRow instead of a failed CompletableFuture as the surface contract promises. Introduce SessionPoolMap, a small wrapper over the existing Util.createSessionMap cache that owns the conversion: get() unwraps UncheckedExecutionException to surface the original cause, and apply() converts a loader throw into a failed CompletableFuture for the async call paths. Replace the inline LoadingCache fields in both shim ops files. Covered by a new focused unit test.

The existing awaitTerminated tests only close an empty pool, close after the vRPC has completed, or issue the vRPC against an already-closed pool — none of them exercise SessionPoolImpl.close()'s pendingRpcs cancellation loop with a non-empty toCancel list. Add closeCancelsInFlightVRpcAndDrains: park the session-open response in a DelayedClientInterceptor so the vRPC sits in pendingRpcs, then close the pool and verify the listener gets a single CANCELLED onClose with the "SessionPool closed:" description, awaitTerminated returns true via the OPENING session draining through sessions.close(req), and no phantom second onClose arrives once the delayed responses flush.

…ailure The catch branch in awaitTerminated is reachable only if drainedFuture is ever completed exceptionally or cancelled — neither happens today, but if it ever does we should leave a telemetry breadcrumb instead of just an uncaught IllegalStateException for the caller to puzzle over. Mirror the precondition-failure pattern used elsewhere in this file (record at ERROR, then throw).

`apply` already converted lookup throws to failed futures, but a sync throw from `op.apply(v)` itself (NPE on malformed input, REE during shutdown, any RuntimeException from the wrapped call chain) propagated to the caller — contradicting the docstring's contract that "callers consistently observe failures through the future surface." Wrap the op invocation in try/catch with the same pattern as the lookup arm.

igorbernstein2 force-pushed the threading-refactor-phase-5 branch from 6a94038 to 49dd066 Compare June 17, 2026 20:02

igorbernstein2 force-pushed the threading-refactor-phase-4 branch from 64d0864 to bb3fd91 Compare June 17, 2026 20:02

nimf approved these changes Jun 22, 2026

View reviewed changes

mutianf reviewed Jun 22, 2026

View reviewed changes

igorbernstein2 force-pushed the threading-refactor-phase-5 branch from 49dd066 to de62fa1 Compare June 22, 2026 21:41

igorbernstein2 force-pushed the threading-refactor-phase-4 branch from bb3fd91 to 1d18231 Compare June 22, 2026 21:41

mutianf approved these changes Jun 23, 2026

View reviewed changes

igorbernstein2 added 7 commits June 29, 2026 16:31

igorbernstein2 force-pushed the threading-refactor-phase-4 branch from 1d18231 to d824fd4 Compare June 29, 2026 17:20

igorbernstein2 force-pushed the threading-refactor-phase-5 branch from b82a480 to 8f643ad Compare June 29, 2026 17:20

igorbernstein2 added 2 commits June 30, 2026 16:03

chore: apply fmt-maven-plugin

3508eb3

Conversation

igorbernstein2 commented Jun 16, 2026

Summary

What's in here

Stack position

Uh oh!

nimf left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants