Threading refactor (6/7): Tracer correctness & chain.isDone()

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

@@ -36,6 +36,11 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
     delegate.cancel(message, cause);
   }
 
+  @Override
+  public boolean isDone() {
+    return delegate.isDone();
+  }
+
   @Override
   public void requestNext() {
     delegate.requestNext();

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

@@ -69,9 +69,10 @@ public RetryingVRpc(Supplier<VRpc<ReqT, RespT>> supplier, BigtableTimer timer) {
   @Override
   public void start(ReqT req, VRpcCallContext ctx, VRpcListener<RespT> listener) {
     if (started) {
-      listener.onClose(
+      VRpcResult alreadyStarted =
           VRpcResult.createRejectedError(
-              Status.FAILED_PRECONDITION.withDescription("operation is already started")));
+              Status.FAILED_PRECONDITION.withDescription("operation is already started"));
+      ctx.getExecutor().execute(() -> listener.onClose(alreadyStarted));
       return;
     }
 
@@ -117,6 +118,11 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
                 Status.CANCELLED.withDescription(message).withCause(finalCause))));
   }
 
+  @Override
+  public boolean isDone() {
+    return currentState.isDone();
+  }
+
   @Override
   public void requestNext() {
     // Assert the op-executor affinity even though the body is dead today — when streaming lands
@@ -134,13 +140,18 @@ void onStateChange(State state) {
     if (currentState.isDone()) {
       return;
     }
+    // Give the outgoing state a chance to release per-state resources (timers, registrations,
+    // tracer pairings). Default is a no-op; states that hold cleanup-worthy resources override.
+    currentState.onExit();
     this.currentState = state;
     currentState.onStart();
   }
 
   abstract static class State {
     public abstract void onStart();
 
+    public void onExit() {}
+
     public void onCancel(String reason, Throwable throwable) {}
 
     public boolean isDone() {
@@ -164,6 +175,16 @@ public void onStart() {
   class Active extends State {
 
     private VRpc<ReqT, RespT> attempt;
+    // Tracer pairing flag scoped to this attempt. finishAttempt is idempotent via this flag so
+    // listener path, cancel path, and onExit safety net never double-fire.
+    private boolean attemptFinished = false;
+
+    private void finishAttempt(VRpcResult result) {
+      if (!attemptFinished) {
+        attemptFinished = true;
+        tracer.onAttemptFinish(result);
+      }
+    }
 
     @Override
     public void onStart() {
@@ -195,12 +216,12 @@ public void onMessage(RespT msg) {
               }
               if (userThrow != null) {
                 // Classify as USER_FAILURE (not CANCELLED, which is what the OpExecutor uncaught
-                // handler would produce via chain.cancel). Finish tracing for the in-flight
-                // attempt, cancel the underlying gRPC call so no further events arrive (its later
-                // onClose is dropped by the currentState != Active.this guard), and transition
-                // directly to Done with the user-error result.
+                // handler would produce via chain.cancel). Finish the attempt's tracer span
+                // with the user-error result, cancel the underlying gRPC call so no further
+                // events arrive (its later onClose is dropped by the currentState !=
+                // Active.this guard), and transition directly to Done.
                 VRpcResult userResult = VRpcResult.createUserError(userThrow);
-                tracer.onAttemptFinish(userResult);
+                finishAttempt(userResult);
                 attempt.cancel("User callback threw", userThrow);
                 onStateChange(new Done(userResult));
               }
@@ -209,7 +230,6 @@ public void onMessage(RespT msg) {
             @Override
             public void onClose(VRpcResult result) {
               context.getExecutor().throwIfNotInThisExecutor();
-              tracer.onAttemptFinish(result);
               if (currentState != Active.this) {
                 LOG.log(
                     Level.FINE,
@@ -218,6 +238,7 @@ public void onClose(VRpcResult result) {
                     result);
                 return;
               }
+              finishAttempt(result);
               if (shouldRetry(result)) {
                 context = context.createForNextAttempt();
                 Duration retryDelay =
@@ -240,14 +261,29 @@ public void onClose(VRpcResult result) {
 
     @Override
     public void onCancel(String reason, Throwable throwable) {
-      // attempt could be null if attemptFactory.get() throws an exception. In which case sync
-      // context uncaught exception handler will be called, which calls cancel on the current
-      // state before transition into done state.
+      // Pair the onAttemptStart fired in onStart with an onAttemptFinish at the moment we
+      // abandon the attempt — the later server onClose for the cancelled attempt is dropped by
+      // the stale-state guard, so this is the only chance to balance the tracer.
+      finishAttempt(
+          VRpcResult.createRejectedError(
+              Status.CANCELLED.withDescription(reason).withCause(throwable)));
+      // attempt could be null if attemptFactory.get() threw before assignment.
       if (attempt != null) {
         attempt.cancel(reason, throwable);
       }
     }
 
+    @Override
+    public void onExit() {
+      // Defense-in-depth: every existing exit path (listener.onClose, onMessage user-throw,
+      // onCancel) calls finishAttempt with a meaningful result before transitioning. This catches
+      // any new exit path that forgets, recording a generic 'abandoned' instead of leaking the
+      // tracer span.
+      finishAttempt(
+          VRpcResult.createRejectedError(
+              Status.CANCELLED.withDescription("attempt abandoned during transition")));
+    }
+
     boolean shouldRetry(VRpcResult result) {
       // If the error has RetryInfo, it means it comes from the server and should
       // be retried.
@@ -281,9 +317,7 @@ 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.
+    // Scheduled to a CANCELLED Done instead of silently discarding the pending timeout.
     private BigtableTimer.Registration stopHook;
 
     Scheduled(Duration retryDelay) {
@@ -304,15 +338,16 @@ public void onStart() {
                         .execute(
                             () ->
                                 grpcContext
-                                    .wrap(() -> otelContext.wrap(this::onTimerFired).run())
+                                    .wrap(
+                                        () ->
+                                            otelContext.wrap(() -> onStateChange(new Idle())).run())
                                     .run()),
                 Durations.toMillis(retryDelay),
                 TimeUnit.MILLISECONDS);
       } 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();
+        // op-executor uncaught handler does not have to. onExit will release the stopHook.
         onStateChange(
             new Done(
                 VRpcResult.createRejectedError(
@@ -323,11 +358,6 @@ 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() {
@@ -346,18 +376,14 @@ private void onTimerStopping() {
               });
     }
 
-    private void unregisterStopHook() {
+    @Override
+    public void onExit() {
+      // Consolidated cleanup: runs on every exit path (normal fire → Idle, cancel → Done,
+      // shutdown hook → Done). Both fields may be null if schedule threw an ISE before assignment.
       if (stopHook != null) {
         stopHook.unregister();
         stopHook = null;
       }
-    }
-
-    @Override
-    public void onCancel(String reason, Throwable throwable) {
-      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();
       }
@@ -374,10 +400,8 @@ class Done extends State {
 
     @Override
     public void onStart() {
-      if (!started) {
-        LOG.fine("operation is not started yet.");
-        return;
-      }
+      // Per-attempt tracer pairing is owned by Active.onExit; Done just runs the user listener
+      // and the per-operation tracer finish.
       Stopwatch appTimer = Stopwatch.createStarted();
       try {
         listener.onClose(result);

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

@@ -95,21 +95,21 @@ public void start(ReqT req, VRpcListener<RespT> listener) {
     // Register AFTER chain.start so a context-cancel that fires immediately is sequenced behind
     // start. Matches ClientCallImpl's ordering (grpc-java issue #1343).
     //
-    // Queueing the registration onto the op executor is what makes the closed-check sound: any
+    // Queueing the registration onto the op executor is what makes the isDone-check sound: any
     // onClose that chain.start enqueued during runInline drains FIRST (FIFO), so by the time this
-    // task runs wrapped.closed reflects whether onClose has already fired. If it has, we skip
-    // addListener — otherwise the listener would pin the chain on grpcContext for its lifetime
-    // (CleanupListener.onClose already called removeListener as a no-op pre-registration). If
-    // grpcContext gets cancelled between start() returning and this task running, the
-    // directExecutor below fires the listener immediately on addListener, so cancel still
-    // propagates correctly.
+    // task runs chain.isDone() reflects whether the chain has already reached terminal. If it
+    // has, we skip addListener — otherwise the listener would pin the chain on grpcContext for
+    // its lifetime (CleanupListener.onClose already called removeListener as a no-op
+    // pre-registration). If grpcContext gets cancelled between start() returning and this task
+    // running, the directExecutor below fires the listener immediately on addListener, so cancel
+    // still propagates correctly.
     //
     // runInline is the right verb here: when chain.start enqueued nothing (common path), the
     // executor is idle and the body runs inline on this thread — no extra context switch. When
     // chain.start did enqueue an onClose, runInline takes the queue branch and FIFO drains both.
     exec.runInline(
         () -> {
-          if (!wrapped.closed) {
+          if (!chain.isDone()) {
             grpcContext.addListener(cancellationListener, MoreExecutors.directExecutor());
           }
         });
@@ -123,10 +123,6 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
   private static class CleanupListener<RespT> extends ForwardListener<RespT> {
     private final Context grpcContext;
     private final Context.CancellationListener cancellationListener;
-    // Read by VOperationImpl.start on the caller thread after runInline returns. runInline runs
-    // chain.start synchronously, so any sync onClose has completed (and this flag been set) by
-    // the time start() reads it on the same thread — no synchronization needed.
-    volatile boolean closed = false;
 
     CleanupListener(
         VRpcListener<RespT> delegate,
@@ -139,7 +135,6 @@ private static class CleanupListener<RespT> extends ForwardListener<RespT> {
 
     @Override
     public void onClose(VRpcResult result) {
-      closed = true;
       grpcContext.removeListener(cancellationListener);
       super.onClose(result);
     }

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

@@ -59,6 +59,14 @@ public interface VRpc<ReqT, RespT> {
   /** Cancel a started RPC. This will be done by best effort. */
   void cancel(@Nullable String message, @Nullable Throwable cause);
 
+  /**
+   * True once a terminal result has been (or is about to be) delivered to the listener; future
+   * events on this VRpc are no-ops. Callers use this to detect a synchronous terminal during {@link
+   * #start} — e.g. VOperationImpl checks this to avoid registering a gRPC cancellation listener
+   * that would never be removed because the chain already finished.
+   */
+  boolean isDone();
+
   /**
    * TBD - server streaming rpcs. This will be used to request more data. Unlike gRPC's request(n),
    * starting a call will implicitly request the first message.

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

@@ -777,6 +777,13 @@ void cancelWithResult(VRpcResult result) {
               });
     }
 
+    @Override
+    public boolean isDone() {
+      // realCall set in drainTo's lambda; once we hand off, it's the source of truth.
+      // Pre-handoff, isCancelled tracks our own terminal state.
+      return realCall != null ? realCall.isDone() : isCancelled;
+    }
+
     @Override
     public void requestNext() {
       if (realCall != null) {

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

@@ -196,6 +196,11 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
     session.cancelRpc(rpcId, message, cause);
   }
 
+  @Override
+  public boolean isDone() {
+    return state.get() == State.CLOSED;
+  }
+
   @Override
   public void requestNext() {
     throw new UnsupportedOperationException("streamed RPCs are not supported yet");

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java

@@ -219,6 +219,11 @@ public void start(Object req, VRpcCallContext ctx, VRpcListener ignored) {
     @Override
     public void cancel(@Nullable String message, @Nullable Throwable cause) {}
 
+    @Override
+    public boolean isDone() {
+      return false;
+    }
+
     @Override
     public void requestNext() {}
   }

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java

@@ -510,6 +510,11 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) {
       throw new UnsupportedOperationException();
     }
 
+    @Override
+    public boolean isDone() {
+      return false;
+    }
+
     @Override
     public void requestNext() {
       throw new UnsupportedOperationException();

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java

@@ -57,6 +57,11 @@ public void cancel(@Nullable String msg, @Nullable Throwable cause) {
             cancelLatch.countDown();
           }
 
+          @Override
+          public boolean isDone() {
+            return false;
+          }
+
           @Override
           public void requestNext() {}
         };
@@ -88,9 +93,9 @@ public void onClose(VRpcResult result) {}
 
   @Test
   void asyncOnCloseFromChainDoesNotPropagateLaterContextCancel() throws InterruptedException {
-    // Regression for the wrapped.closed TOCTOU. When chain.start asynchronously queues an
+    // Regression for the chain-already-done TOCTOU. When chain.start asynchronously queues an
     // onClose via the op executor, the addListener task (also queued through the op executor
-    // by VOperationImpl.start) drains FIFO after the onClose and observes wrapped.closed=true,
+    // by VOperationImpl.start) drains FIFO after the onClose and observes chain.isDone()=true,
     // so the cancellationListener is NOT registered. A later grpcContext.cancel therefore has
     // no path to reach chain.cancel — which is correct because the chain has already terminated.
     Context.CancellableContext grpcContext = Context.current().withCancellation();
@@ -100,22 +105,31 @@ void asyncOnCloseFromChainDoesNotPropagateLaterContextCancel() throws Interrupte
 
     VRpc<String, String> chain =
         new VRpc<String, String>() {
+          private volatile boolean done = false;
+
           @Override
           public void start(String req, VRpcCallContext ctx, VRpcListener<String> listener) {
             // Simulate PendingVRpc pool-closed branch / VRpcImpl deadline short-circuit.
             ctx.getExecutor()
                 .execute(
-                    () ->
-                        listener.onClose(
-                            VRpcResult.createUncommitedError(
-                                Status.UNAVAILABLE.withDescription("fast-fail"))));
+                    () -> {
+                      done = true;
+                      listener.onClose(
+                          VRpcResult.createUncommitedError(
+                              Status.UNAVAILABLE.withDescription("fast-fail")));
+                    });
           }
 
           @Override
           public void cancel(@Nullable String msg, @Nullable Throwable cause) {
             chainCancelCount.incrementAndGet();
           }
 
+          @Override
+          public boolean isDone() {
+            return done;
+          }
+
           @Override
           public void requestNext() {}
         };