fix: deliver terminal onClose to Scheduled retries pending at Client.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.

Author: igorbernstein2

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java

@@ -273,17 +273,22 @@ public void close() {
       }
     }
 
-    // Phase 3: tear down infrastructure. By this point all listener.onClose tasks for in-flight
-    // RPCs are queued on their op executors (which run on userCallbackExecutor), and no new
-    // session responses are coming since every session is CLOSED. The 5s await inside
-    // userCallbackExecutor.close() is therefore just a guard for tasks in flight — it should
-    // return immediately in the typical case.
+    // Phase 3: tear down infrastructure.
+    //
+    // sessionTimer.stop() runs FIRST so its onStop hooks can drive any pending Scheduled retries
+    // to a terminal Done — that delivery hops through op executor → userCallbackExecutor, both
+    // of which must still be alive at this moment.
+    //
+    // userCallbackExecutor.close() next, with a 5s drain to catch the listener.onClose tasks
+    // queued by both the session drain (Phase 2) and the just-fired retry shutdowns.
+    //
+    // backgroundExecutor must close last because it's the timer's dispatcher and the op
+    // executor's chain ultimately runs ScheduledExecutorService tasks here.
+    sessionTimer.stop();
     userCallbackExecutor.close();
     metrics.close();
     channelPool.close();
     configManager.close();
-    // Stop the timer before tearing down backgroundExecutor (the timer's dispatcher).
-    sessionTimer.stop();
     backgroundExecutor.close();
   }
 

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java

@@ -25,7 +25,6 @@
 import io.grpc.Context;
 import io.grpc.Status;
 import java.util.Optional;
-import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import java.util.logging.Level;
@@ -281,6 +280,11 @@ boolean shouldRetry(VRpcResult result) {
   class Scheduled extends State {
     private final Duration retryDelay;
     private BigtableTimer.Timeout future;
+    // Registered with the timer on entry so a Client.close that stops the timer drives this
+    // Scheduled to a CANCELLED Done instead of silently discarding the pending timeout. Cleared
+    // on every exit path (normal fire, cancel, hook fire) to avoid accumulating dead entries on
+    // a long-lived Client.
+    private BigtableTimer.Registration stopHook;
 
     Scheduled(Duration retryDelay) {
       this.retryDelay = retryDelay;
@@ -289,6 +293,7 @@ class Scheduled extends State {
     @Override
     public void onStart() {
       try {
+        stopHook = timer.onStop(this::onTimerStopping);
         // Wraps go innermost so the captured gRPC + OpenTelemetry contexts are re-established at
         // the moment the body runs, not just while the dispatcher is invoking the outer task.
         future =
@@ -299,13 +304,15 @@ public void onStart() {
                         .execute(
                             () ->
                                 grpcContext
-                                    .wrap(
-                                        () ->
-                                            otelContext.wrap(() -> onStateChange(new Idle())).run())
+                                    .wrap(() -> otelContext.wrap(this::onTimerFired).run())
                                     .run()),
                 Durations.toMillis(retryDelay),
                 TimeUnit.MILLISECONDS);
-      } catch (RejectedExecutionException e) {
+      } catch (IllegalStateException e) {
+        // Timer was stopped between Active.onClose deciding to retry and this task running on the
+        // op executor. Race window is narrow (post-drain shutdown), but cover it cleanly so the
+        // op-executor uncaught handler does not have to.
+        unregisterStopHook();
         onStateChange(
             new Done(
                 VRpcResult.createRejectedError(
@@ -316,11 +323,38 @@ public void onStart() {
       }
     }
 
+    private void onTimerFired() {
+      unregisterStopHook();
+      onStateChange(new Idle());
+    }
+
+    // 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() {
+      context.getExecutor().execute(() -> {
+        if (currentState != Scheduled.this) {
+          return; // already transitioned out via normal fire or cancel
+        }
+        onStateChange(
+            new Done(
+                VRpcResult.createRejectedError(
+                    Status.CANCELLED.withDescription(
+                        "Client closing while retry pending"))));
+      });
+    }
+
+    private void unregisterStopHook() {
+      if (stopHook != null) {
+        stopHook.unregister();
+        stopHook = null;
+      }
+    }
+
     @Override
     public void onCancel(String reason, Throwable throwable) {
-      // future can be null if schedule throws an exception that's not RejectedExecutionException.
-      // In which case sync context uncaught exception handler will be called, which calls cancel on
-      // the current state before transition into done state.
+      unregisterStopHook();
+      // future can be null if schedule throws and we end up here via the op-executor uncaught
+      // path.
       if (future != null && !future.isCancelled()) {
         future.cancel();
       }

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java

@@ -47,14 +47,34 @@ public interface BigtableTimer {
 
   /**
    * Releases the tick thread and discards any pending timeouts. Idempotent. After {@code stop()},
-   * subsequent calls to {@link #newTimeout} throw {@link IllegalStateException}.
+   * subsequent calls to {@link #newTimeout} or {@link #onStop} throw {@link
+   * IllegalStateException}.
+   *
+   * <p>Before releasing the tick thread, invokes every hook registered via {@link #onStop} on the
+   * caller thread. Hooks fire in unspecified order; a hook that throws is logged and other hooks
+   * still fire.
    */
   void stop();
 
+  /**
+   * Registers a hook to run during {@link #stop()}. Use this to drive caller-owned state (e.g. a
+   * scheduled retry waiting on the timer) to a terminal state before the timer is torn down,
+   * instead of letting a pending timeout silently disappear.
+   *
+   * <p>The returned {@link Registration} unregisters the hook; call it when the hook is no longer
+   * needed (e.g. the scheduled work fired normally or was cancelled) so the hook set does not
+   * accumulate stale entries.
+   */
+  Registration onStop(Runnable hook);
+
   interface Timeout {
     /** Cancels the scheduled task. Returns true if the task had not yet fired. */
     boolean cancel();
 
     boolean isCancelled();
   }
+
+  interface Registration {
+    void unregister();
+  }
 }

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java

@@ -17,8 +17,13 @@
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import io.grpc.netty.shaded.io.netty.util.HashedWheelTimer;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 /**
  * {@link BigtableTimer} backed by Netty's {@code HashedWheelTimer}, accessed via the shaded copy
@@ -29,6 +34,8 @@
  * with an in-tree implementation that does not reach into gRPC's shaded internals.
  */
 public final class NettyWheelTimer implements BigtableTimer {
+  private static final Logger LOG = Logger.getLogger(NettyWheelTimer.class.getName());
+
   // 10 ms tick × 512 buckets ≈ 5 s per rotation. Heartbeat (100 ms), deadlines (sub-second to
   // seconds), and watchdog (5 min) all sit comfortably inside this resolution.
   private static final long TICK_DURATION_MS = 10;
@@ -37,6 +44,11 @@ public final class NettyWheelTimer implements BigtableTimer {
   private final HashedWheelTimer delegate;
   private final Executor dispatcher;
 
+  // ConcurrentHashMap-backed Set so onStop/Registration.unregister can run from any thread without
+  // blocking newTimeout. Stop drains it once, then refuses further registrations.
+  private final Set<Runnable> stopHooks = ConcurrentHashMap.newKeySet();
+  private volatile boolean stopped = false;
+
   public NettyWheelTimer(String name, Executor dispatcher) {
     this.dispatcher = dispatcher;
     this.delegate =
@@ -49,11 +61,39 @@ public NettyWheelTimer(String name, Executor dispatcher) {
 
   @Override
   public Timeout newTimeout(Runnable task, long delay, TimeUnit unit) {
-    return new TimeoutHandle(delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit));
+    if (stopped) {
+      throw new IllegalStateException("timer stopped");
+    }
+    return new TimeoutHandle(
+        delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit));
+  }
+
+  @Override
+  public Registration onStop(Runnable hook) {
+    if (stopped) {
+      throw new IllegalStateException("timer stopped");
+    }
+    stopHooks.add(hook);
+    return () -> stopHooks.remove(hook);
   }
 
   @Override
   public void stop() {
+    if (stopped) {
+      return;
+    }
+    stopped = true;
+    // Snapshot then clear so hooks can no longer be unregistered while we iterate (and so a hook
+    // that re-enters onStop sees stopped=true and fails fast).
+    Set<Runnable> hooks = new HashSet<>(stopHooks);
+    stopHooks.clear();
+    for (Runnable hook : hooks) {
+      try {
+        hook.run();
+      } catch (Throwable t) {
+        LOG.log(Level.WARNING, "stop hook threw; continuing", t);
+      }
+    }
     delegate.stop();
   }
 

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java

@@ -843,6 +843,11 @@ public boolean isCancelled() {
       };
     }
 
+    @Override
+    public Registration onStop(Runnable hook) {
+      return delegate.onStop(hook);
+    }
+
     @Override
     public void stop() {
       delegate.stop();

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java

@@ -84,6 +84,11 @@ public boolean isCancelled() {
       };
     }
 
+    @Override
+    public Registration onStop(Runnable hook) {
+      return () -> {};
+    }
+
     @Override
     public void stop() {}
   }