diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/ForwardingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/ForwardingVRpc.java index 05173e6a9ef9..47058986a5b5 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/ForwardingVRpc.java +++ b/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(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java index 614f45eb2eb4..d11c509798d4 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/RetryingVRpc.java +++ b/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> supplier, BigtableTimer timer) { @Override public void start(ReqT req, VRpcCallContext ctx, VRpcListener 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,6 +140,9 @@ 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(); } @@ -141,6 +150,8 @@ void onStateChange(State state) { 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 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); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImpl.java index a11a1e5cbbd1..062ed34f9a36 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImpl.java +++ b/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 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 extends ForwardListener { 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 delegate, @@ -139,7 +135,6 @@ private static class CleanupListener extends ForwardListener { @Override public void onClose(VRpcResult result) { - closed = true; grpcContext.removeListener(cancellationListener); super.onClose(result); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java index f45d523b5efd..8b700066a064 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/VRpc.java +++ b/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 { /** 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. diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java index 04f31e33c335..d3e65975a394 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImpl.java +++ b/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) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java index 83228fd43456..b6ac222336da 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java +++ b/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"); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java index 60c706a7af9d..b5a136f9a193 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/TableBaseTest.java +++ b/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() {} } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java index b9765d4e777b..4602c0941822 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/VRpcTracerTest.java +++ b/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(); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java index a69bed4c0462..337115e7a6d9 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java +++ b/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,15 +105,19 @@ void asyncOnCloseFromChainDoesNotPropagateLaterContextCancel() throws Interrupte VRpc chain = new VRpc() { + private volatile boolean done = false; + @Override public void start(String req, VRpcCallContext ctx, VRpcListener 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 @@ -116,6 +125,11 @@ public void cancel(@Nullable String msg, @Nullable Throwable cause) { chainCancelCount.incrementAndGet(); } + @Override + public boolean isDone() { + return done; + } + @Override public void requestNext() {} };