diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 0d63a07e866e..56d517f021a7 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java @@ -171,7 +171,7 @@ public static Client create(ClientSettings settings) throws IOException { Resource.createOwned(metrics, metrics::close), Resource.createOwned(configManager, configManager::close), Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown), - Resource.createOwned(userCallbackExecutor, userCallbackExecutor::shutdown)); + Resource.createOwned(userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor))); } public Client( @@ -225,13 +225,17 @@ public void close() { .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER) .setDescription("Client closing") .build())); + // Drain user-callback first so pool.close's cancelWithResult listener notifications complete + // before we tear down the surrounding executors and timer. Without this, the late onClose + // submissions race the shutdown and get RejectedExecutionException, silently dropping the + // user's terminal onClose. + userCallbackExecutor.close(); metrics.close(); channelPool.close(); configManager.close(); // Stop the timer before tearing down backgroundExecutor (the timer's dispatcher). sessionTimer.stop(); backgroundExecutor.close(); - userCallbackExecutor.close(); } public TableAsync openTableAsync(String tableId, Permission permission) { @@ -289,8 +293,10 @@ public MaterializedViewAsync openMaterializedViewAsync( } public static class Resource { - private T value; - private Runnable closer; + private final T value; + private final Runnable closer; + private final java.util.concurrent.atomic.AtomicBoolean closed = + new java.util.concurrent.atomic.AtomicBoolean(false); public static Resource createOwned(T value, Runnable closer) { return new Resource<>(value, closer); @@ -305,12 +311,29 @@ private Resource(T value, Runnable closer) { this.closer = closer; } + /** Idempotent. Repeat calls are no-ops. */ public void close() { - this.closer.run(); + if (closed.compareAndSet(false, true)) { + this.closer.run(); + } } public T get() { return value; } } + + // Drain in-flight listener.onClose tasks before the executor is shut down; bound the wait at 5s + // so close() doesn't hang the caller on a pathological listener. + private static void shutdownAndAwait(ExecutorService exec) { + exec.shutdown(); + try { + if (!exec.awaitTermination(5, TimeUnit.SECONDS)) { + exec.shutdownNow(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + exec.shutdownNow(); + } + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java index 9a60a271ea7f..402151676c7c 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/channels/ChannelPoolDpImpl.java @@ -573,7 +573,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return com.google.common.base.Objects.hashCode(id); + return Long.hashCode(id); } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/OpExecutor.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/OpExecutor.java index 7622647bcc38..4eb0115558dc 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/OpExecutor.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/middleware/OpExecutor.java @@ -16,17 +16,27 @@ package com.google.cloud.bigtable.data.v2.internal.middleware; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.ArrayDeque; import java.util.concurrent.Executor; +import javax.annotation.concurrent.GuardedBy; /** - * Per-op serializing executor. Wraps a delegate {@link Executor} (typically a per-call {@code - * SerializingExecutor} over the user-callback pool) and tracks which thread is currently running a - * task, so callers can assert they are on the executor (analogous to {@code + * Per-op serializing executor that tracks which thread is currently draining it, so callers can + * assert they are running inside the executor (analogous to {@code * SynchronizationContext#throwIfNotInThisSynchronizationContext}). * - *

If a task throws, the registered {@link UncaughtExceptionHandler} is invoked — this is the - * last-resort recovery point for the chain. Without it, a throw from a user-installed tracer or a - * listener callback would silently drop and the caller's future would never complete. + *

{@link #runInline} executes synchronously on the caller thread when the executor is idle, or + * falls back to a queued submission when busy. Use it to avoid the queue+drain round-trip for the + * first task of a fresh op executor (e.g. the start() dispatch). + * + *

If a task throws, the registered {@link UncaughtExceptionHandler} is invoked on the same task + * thread (still inside the drain wrapper, so {@link #throwIfNotInThisExecutor} still passes). This + * is the last-resort recovery point for the op chain — without it, an exception inside a callback + * is silently dropped and the caller's listener never sees a terminal close. The handler's own + * throws propagate (caught by the backing executor in production; surfaced to the calling thread + * when the backing is {@link MoreExecutors#directExecutor()}, which is what makes fail-fast test + * handlers like {@code t -> throw new AssertionError(t)} work). */ public final class OpExecutor implements Executor { @@ -36,6 +46,15 @@ public interface UncaughtExceptionHandler { private final Executor backing; private final UncaughtExceptionHandler handler; + + // queue is the lock. runningThread is volatile (not @GuardedBy): throwIfNotInThisExecutor() + // reads it lock-free; writes happen inside synchronized(queue) to piggy-back the memory barrier. + @GuardedBy("queue") + private final ArrayDeque queue = new ArrayDeque<>(); + + @GuardedBy("queue") + private boolean drainScheduled = false; + private volatile Thread runningThread; public OpExecutor(Executor backing, UncaughtExceptionHandler handler) { @@ -45,18 +64,84 @@ public OpExecutor(Executor backing, UncaughtExceptionHandler handler) { @Override public void execute(Runnable r) { - backing.execute( - () -> { - Thread prev = runningThread; - runningThread = Thread.currentThread(); - try { - r.run(); - } catch (Throwable t) { - handler.uncaught(t); - } finally { - runningThread = prev; - } - }); + synchronized (queue) { + queue.add(r); + if (!drainScheduled && runningThread == null) { + scheduleDrainLocked(); + } + } + } + + /** + * Runs {@code r} synchronously on the caller thread if this executor is idle, otherwise queues it + * for later drain on the backing executor. Either way, FIFO ordering with other tasks is + * preserved. + */ + public void runInline(Runnable r) { + synchronized (queue) { + if (drainScheduled || runningThread != null || !queue.isEmpty()) { + queue.add(r); + if (!drainScheduled) { + scheduleDrainLocked(); + } + return; + } + runningThread = Thread.currentThread(); + } + try { + try { + r.run(); + } catch (Throwable t) { + handler.uncaught(t); + } + } finally { + synchronized (queue) { + runningThread = null; + if (!queue.isEmpty() && !drainScheduled) { + scheduleDrainLocked(); + } + } + } + } + + // Schedule a drain on the backing executor. If the backing throws (e.g. + // RejectedExecutionException + // during shutdown), reset drainScheduled before propagating so the next execute() can retry + // instead of wedging the executor with no drainer. + @GuardedBy("queue") + private void scheduleDrainLocked() { + drainScheduled = true; + try { + backing.execute(this::drain); + } catch (Throwable t) { + drainScheduled = false; + throw t; + } + } + + private void drain() { + while (true) { + Runnable r; + synchronized (queue) { + r = queue.poll(); + if (r == null) { + drainScheduled = false; + return; + } + runningThread = Thread.currentThread(); + } + try { + try { + r.run(); + } catch (Throwable t) { + handler.uncaught(t); + } + } finally { + synchronized (queue) { + runningThread = null; + } + } + } } public void throwIfNotInThisExecutor() { 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 8efd123137da..0c9d2887d87a 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 @@ -32,6 +32,8 @@ import java.util.logging.Logger; import javax.annotation.Nullable; +// All mutable state is owned by the op executor; VOperationImpl trampolines every inbound call +// onto it, so no synchronization is needed here. public class RetryingVRpc implements VRpc { private static final Logger LOG = Logger.getLogger(RetryingVRpc.class.getName()); @@ -47,11 +49,9 @@ public class RetryingVRpc implements VRpc { private final BigtableTimer timer; - // current state and all the flags don't need to be volatile because they're only updated within - // the op executor. private State currentState; private boolean started; - // Breaks the loop if uncaught exception happens during op-executor execution. + // Breaks the loop on uncaught exception during cancel. private boolean isCancelling; public RetryingVRpc(Supplier> supplier, BigtableTimer timer) { @@ -69,64 +69,69 @@ public RetryingVRpc(Supplier> supplier, BigtableTimer timer) { @Override public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { - ctx.getExecutor() - .execute( - () -> { - if (started) { - listener.onClose( - VRpcResult.createRejectedError( - Status.FAILED_PRECONDITION.withDescription( - "operation is already started"))); - return; - } - started = true; - - this.request = req; - this.listener = listener; - this.context = ctx; - this.tracer = context.getTracer(); + if (started) { + listener.onClose( + VRpcResult.createRejectedError( + Status.FAILED_PRECONDITION.withDescription("operation is already started"))); + return; + } - tracer.onOperationStart(); - currentState.onStart(); - }); + // Publish the fields BEFORE the try block. If anything below throws and we recover via + // cancel(), cancel() reads this.context / this.listener — they must be set already, or + // we trade the original failure for an NPE inside the recovery path. + this.request = req; + this.listener = listener; + this.context = ctx; + this.tracer = context.getTracer(); + + tracer.onOperationStart(); + started = true; + + try { + currentState.onStart(); + } catch (Throwable t) { + cancel("Unexpected error in start", t); + } } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { - context - .getExecutor() - .execute( - () -> { - if (currentState.isDone() || isCancelling) { - LOG.fine("Ignoring cancel because the vRPC is already cancelled or done."); - return; - } - // Prevents infinite loop if there's any error thrown during this phase. - isCancelling = true; - Throwable finalCause = cause; - try { - currentState.onCancel(message, cause); - } catch (Throwable t) { - if (finalCause != null) { - finalCause.addSuppressed(t); - } else { - finalCause = t; - } - } - onStateChange( - new Done( - VRpcResult.createRejectedError( - Status.CANCELLED.withDescription(message).withCause(finalCause)))); - }); + if (currentState.isDone() || isCancelling) { + LOG.fine("Ignoring cancel because the vRPC is already cancelled or done."); + return; + } + // Prevents infinite loop if there's any error thrown during this phase. + isCancelling = true; + Throwable finalCause = cause; + try { + currentState.onCancel(message, cause); + } catch (Throwable t) { + if (finalCause != null) { + finalCause.addSuppressed(t); + } else { + finalCause = t; + } + } + onStateChange( + new Done( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(message).withCause(finalCause)))); } @Override public void requestNext() { + // Assert the op-executor affinity even though the body is dead today — when streaming lands + // and this becomes real, the missing assertion would silently allow off-thread access. + // Guarded on context being set so a misuse before start() still throws + // UnsupportedOperationException + // rather than NPE on the assertion. + if (context != null) { + context.getExecutor().throwIfNotInThisExecutor(); + } throw new UnsupportedOperationException("request next is not supported in unary"); } void onStateChange(State state) { - context.getExecutor().throwIfNotInThisExecutor(); if (currentState.isDone()) { return; } @@ -169,10 +174,9 @@ public void onStart() { request, context, new VRpcListener() { - // VRpcImpl dispatches its callbacks via ctx.getExecutor() already, so these methods - // run inside the op-executor task — no need to re-dispatch here. @Override public void onMessage(RespT msg) { + context.getExecutor().throwIfNotInThisExecutor(); if (currentState != Active.this) { LOG.log( Level.FINE, @@ -182,15 +186,30 @@ public void onMessage(RespT msg) { } tracer.onResponseReceived(); Stopwatch appTimer = Stopwatch.createStarted(); + Throwable userThrow = null; try { listener.onMessage(msg); + } catch (Throwable t) { + userThrow = t; } finally { tracer.recordApplicationBlockingLatencies(appTimer.elapsed()); } + 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. + VRpcResult userResult = VRpcResult.createUserError(userThrow); + tracer.onAttemptFinish(userResult); + attempt.cancel("User callback threw", userThrow); + onStateChange(new Done(userResult)); + } } @Override public void onClose(VRpcResult result) { + context.getExecutor().throwIfNotInThisExecutor(); tracer.onAttemptFinish(result); if (currentState != Active.this) { LOG.log( @@ -214,6 +233,7 @@ public void onClose(VRpcResult result) { } return; } + onStateChange(new Done(result)); } }); @@ -271,8 +291,6 @@ public void onStart() { try { // 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. - // The executor may queue the inner runnable for a later drain on a different thread; an - // outer wrap's scope would already be closed by then. future = timer.newTimeout( () -> @@ -300,9 +318,9 @@ public void onStart() { @Override public void onCancel(String reason, Throwable throwable) { - // future can be null if newTimeout 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. + // 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. if (future != null && !future.isCancelled()) { future.cancel(); } 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 8e8a1554bea8..67d7209a7cf4 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 @@ -30,8 +30,9 @@ import javax.annotation.Nullable; /** - * The single edge between the user and the VRpc middleware chain. Constructs the per-op {@link - * VRpcCallContext} and owns the gRPC {@link Context} cancellation listener. + * The single edge between the user and the VRpc middleware chain. Trampolines all inbound user + * calls onto opExecutor and owns the gRPC {@link Context} cancellation listener so that every layer + * below is single-threaded on opExecutor. * *

Precondition: {@link #cancel} must not be called before {@link #start}. */ @@ -45,6 +46,10 @@ public class VOperationImpl implements VOperation { private final boolean idempotent; private final Context.CancellationListener cancellationListener; + // Written in start() on the caller thread before the listener is registered and before cancel() + // is reachable from any external thread. Volatile for safe publication to those threads. + private volatile OpExecutor opExecutor; + public VOperationImpl( VRpc chain, Context grpcContext, @@ -63,7 +68,7 @@ public VOperationImpl( boolean deadlineExceeded = Optional.ofNullable(c.getDeadline()).map(Deadline::isExpired).orElse(false); deadlineExceeded = deadlineExceeded && c.cancellationCause() instanceof TimeoutException; - // Let VRpc machinery handle deadline exceeded. + // Let VRpc machinery handle deadline exceeded if (!deadlineExceeded) { cancel("gRPC context cancelled", c.cancellationCause()); } @@ -72,27 +77,47 @@ public VOperationImpl( @Override public void start(ReqT req, VRpcListener listener) { - // Per-call SerializingExecutor over the shared user-callback pool. The handler is the - // last-resort recovery: if any op-executor task throws (typically a user-installed tracer or - // a listener callback escape), drive the chain to a terminal state so the caller's listener - // still receives an onClose. RetryingVRpc.cancel is idempotent so cascades collapse safely. + // Last-resort recovery: if any op-executor task throws (typically a user-installed tracer, + // or a listener callback that escapes RetryingVRpc's existing per-state try/catches), drive + // the chain to a terminal state so the caller's listener still receives an onClose. The + // handler runs on the failed task's wrapper, so chain.cancel() — which calls + // OpExecutor#throwIfNotInThisExecutor — passes. RetryingVRpc.cancel is idempotent + // (isCancelling / currentState.isDone() guards), so a cascade of failures collapses to a + // single Done. OpExecutor exec = new OpExecutor( - MoreExecutors.newSequentialExecutor(userCallbackExecutor), - t -> chain.cancel("Uncaught exception in op executor task", t)); + userCallbackExecutor, t -> chain.cancel("Uncaught exception in op executor task", t)); + this.opExecutor = exec; VRpcCallContext ctx = VRpcCallContext.create(deadline, idempotent, tracer, exec); - grpcContext.addListener(cancellationListener, MoreExecutors.directExecutor()); - chain.start(req, ctx, new CleanupListener<>(listener, grpcContext, cancellationListener)); + CleanupListener wrapped = + new CleanupListener<>(listener, grpcContext, cancellationListener); + exec.runInline(() -> chain.start(req, ctx, wrapped)); + // Register AFTER chain.start so a context-cancel that fires immediately is sequenced behind + // start. runInline runs chain.start synchronously, so it has fully completed by the time the + // listener is registered. Matches ClientCallImpl's ordering (grpc-java issue #1343). + // + // If the chain reached a terminal onClose synchronously inside runInline (uncaught-handler + // recovery, immediate failure), CleanupListener already tried to remove a listener that was + // never registered (no-op). Skip the addListener in that case — otherwise we'd register a + // listener on grpcContext that nothing will ever remove, pinning the entire chain for the + // lifetime of the (potentially long-lived) gRPC Context. + if (!wrapped.closed) { + grpcContext.addListener(cancellationListener, MoreExecutors.directExecutor()); + } } @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { - chain.cancel(message, cause); + opExecutor.execute(() -> chain.cancel(message, 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, @@ -105,6 +130,7 @@ 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/session/SessionList.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java index 9beccf40fd85..6f8a3395c9f4 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java @@ -170,19 +170,35 @@ void onSessionStarted() { } /** - * The session is returned to the pool after use. This undoes what SessionList#checkoutSession + * The session is returned to the pool after a vRPC that reached the wire completes. Updates the + * picker's per-AFE latency stats on success; non-OK results skip the latency update so a + * fast-failing or cancelled-mid-flight AFE doesn't look the fastest. */ void onVRpcFinish(Duration elapsed, VRpcResult result) { - // Guaranteed to be set - vrpc can only start after the session is ready + releaseToPool(); + if (result.getStatus().isOk()) { + // Guaranteed to be set - vrpc can only start after the session is ready + this.afe.get().updateLatency(elapsed, result.getBackendLatency()); + } + } + + /** + * The session is returned to the pool after a pending vRPC was drained but cancelled before it + * could be attached to a real call (e.g. user cancelled or deadline expired between + * checkoutSession and drainTo). No latency is reported because the vRPC never reached the wire. + */ + void onPendingVRpcCancelled() { + releaseToPool(); + } + + // Shared bookkeeping for both completion paths. Undoes what SessionList#checkoutSession did. + private void releaseToPool() { + // Guaranteed to be set - checkoutSession only runs after the session is ready AfeHandle afeHandle = this.afe.get(); poolStats.inUseCount--; inUseSessions.remove(this); - if (result.getStatus().isOk()) { - afeHandle.updateLatency(elapsed, result.getBackendLatency()); - } - if (session.getState() == SessionState.READY) { poolStats.readyCount++; afeHandle.sessions.add(this); 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 8f2d6e8bd85d..2dfd73f1ba7f 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 @@ -236,6 +236,33 @@ public SessionPoolImpl( }); } + @GuardedBy("this") + private void scheduleNextAfePrune() { + if (closed) { + return; + } + afeListPruneTimeout = + timer.newTimeout( + this::runAfePruneAndReschedule, + SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(), + TimeUnit.MILLISECONDS); + } + + private void runAfePruneAndReschedule() { + synchronized (SessionPoolImpl.this) { + try { + if (closed) { + return; + } + sessions.prune(); + } catch (Throwable t) { + logger.log(Level.WARNING, "AFE prune tick threw; continuing", t); + } finally { + scheduleNextAfePrune(); + } + } + } + @Override public SessionPoolInfo getInfo() { return info; @@ -245,6 +272,7 @@ public SessionPoolInfo getInfo() { public void close(CloseSessionRequest req) { configListenerHandle.close(); + List> toCancel; synchronized (this) { if (poolState == PoolState.CLOSED) { logger.fine(String.format("Tried to close a closed SessionPool %s", info.getLogName())); @@ -255,9 +283,8 @@ public void close(CloseSessionRequest req) { poolState = PoolState.CLOSED; closed = true; - for (PendingVRpc pendingRpc : pendingRpcs) { - pendingRpc.cancel("SessionPool closed: " + req, null); - } + toCancel = new ArrayList<>(pendingRpcs); + pendingRpcs.clear(); if (afeListPruneTimeout != null) { afeListPruneTimeout.cancel(); afeListPruneTimeout = null; @@ -270,36 +297,14 @@ public void close(CloseSessionRequest req) { sessions.close(req); } - // Timer is owned by the Client and shared across pools; do not stop it here. - } - - // Self-rescheduling AFE prune. Pattern matches Watchdog: tick dispatches to the pool executor, - // executor takes the lock, prunes, schedules the next tick. Tolerates body exceptions so a - // transient fault does not permanently disable pruning. - @GuardedBy("this") - private void scheduleNextAfePrune() { - if (closed) { - return; - } - afeListPruneTimeout = - timer.newTimeout( - this::runAfePruneAndReschedule, - SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(), - TimeUnit.MILLISECONDS); - } - - private void runAfePruneAndReschedule() { - synchronized (SessionPoolImpl.this) { - try { - if (closed) { - return; - } - sessions.prune(); - } catch (Throwable t) { - logger.log(Level.WARNING, "AFE prune tick threw; continuing", t); - } finally { - scheduleNextAfePrune(); - } + // cancelWithResult trampolines through ctx.getExecutor() — required because the public + // cancel(String, Throwable) path asserts opExecutor affinity, but close() runs on the + // caller thread. + VRpcResult closeResult = + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription("SessionPool closed: " + req)); + for (PendingVRpc pendingRpc : toCancel) { + pendingRpc.cancelWithResult(closeResult); } } @@ -456,7 +461,18 @@ private synchronized void onSessionReady(SessionHandle handle, OpenSessionRespon private synchronized void onVRpcComplete( SessionHandle handle, Duration elapsed, VRpcResult result) { handle.onVRpcFinish(elapsed, result); + afterRelease(handle); + } + + // Called when a pending vRPC was drained but cancelled before it attached to a real call. + // Mirrors onVRpcComplete; the handle reports no latency because nothing ran on it. + private synchronized void onPendingVRpcCancelled(SessionHandle handle) { + handle.onPendingVRpcCancelled(); + afterRelease(handle); + } + @GuardedBy("this") + private void afterRelease(SessionHandle handle) { // pool is shutting down, dont try to drain vrpcs if (poolState != PoolState.STARTED) { return; @@ -542,16 +558,7 @@ private void onSessionClose( status, trailers))); for (PendingVRpc vrpc : toBeClosed) { try { - vrpc.ctx - .getExecutor() - .execute( - () -> { - try { - vrpc.getListener().onClose(result); - } catch (Throwable t) { - logger.log(Level.WARNING, "Exception when closing request", t); - } - }); + vrpc.cancelWithResult(result); } catch (Throwable t) { logger.log(Level.WARNING, "Exception dispatching close to op executor", t); } @@ -562,10 +569,6 @@ private void onSessionClose( @GuardedBy("this") private void tryDrainPendingRpcs() { while (!pendingRpcs.isEmpty()) { - if (pendingRpcs.peek().isCancelled) { - pendingRpcs.pop(); - continue; - } Optional handle = picker.pickSession(); if (!handle.isPresent()) { break; @@ -581,11 +584,8 @@ private void tryDrainPendingRpcs() { Iterator> iter = pendingRpcs.iterator(); while (iter.hasNext()) { PendingVRpc vrpc = iter.next(); - // vrpcs that have started on a session gets closed in SessionImpl. Do not double close. - if (!vrpc.isCancelled && vrpc.realCall == null) { - iter.remove(); - toBeClosed.add(vrpc); - } + iter.remove(); + toBeClosed.add(vrpc); } return toBeClosed; } @@ -606,7 +606,6 @@ public synchronized VRpc(desc); } - @GuardedBy("this") private VRpc newRealCall( VRpcDescriptor desc, SessionHandle handle) { @@ -667,7 +666,6 @@ public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { this.req = req; this.ctx = ctx; this.listener = listener; - this.deadlineMonitor = monitorDeadline(ctx.getOperationInfo().getDeadline()); synchronized (SessionPoolImpl.this) { if (SessionPoolImpl.this.poolState != PoolState.STARTED) { @@ -677,6 +675,10 @@ public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { ctx.getExecutor().execute(() -> listener.onClose(result)); return; } + // Only arm the deadline monitor after we've committed to queueing; otherwise the + // fast-fail early return above would leak a timer that fires later and emits a phantom + // tracer.onAttemptFinish on the Active state's stale listener. + this.deadlineMonitor = monitorDeadline(ctx.getOperationInfo().getDeadline()); pendingRpcs.add(this); if (logger.isLoggable(Level.FINE)) { @@ -696,9 +698,6 @@ public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { } } - // It's safe to call cancel on a vrpc more than once. It'll be a noop after the initial - // call. Cancelled vrpcs are removed from the pending vrpc queue the next time we - // drain the queue. @Override public void cancel(@Nullable String message, @Nullable Throwable cause) { Status status = Status.CANCELLED; @@ -711,31 +710,35 @@ public void cancel(@Nullable String message, @Nullable Throwable cause) { cancel(status, false); } - // Cancel could race with drainTo which sets the real call. Assign realCall to a NOOP_CALL - // so if drainTo gets called at the same time, it'll just get swallowed and we're only calling - // onClose once on the listener. The cancel could also come from deadline monitor when - // the deadline expires. In this case if the real call is already set, we want to real call - // to handle the deadline and return early. + // cancel() and drainTo() are sequenced via ctx.getExecutor() (a per-op SerializingExecutor), + // so isCancelled and realCall are owned exclusively by that executor — no pool lock needed. private void cancel(Status status, boolean onlyCancelPendingCall) { - boolean delegateToRealCall = true; synchronized (SessionPoolImpl.this) { - if (isCancelled) { - return; - } - isCancelled = true; - if (realCall == null) { - this.realCall = NOOP_CALL; - delegateToRealCall = false; - } else if (onlyCancelPendingCall) { - return; - } - } - if (delegateToRealCall) { - realCall.cancel(status.getDescription(), status.getCause()); - } else { - VRpcResult result = VRpcResult.createRejectedError(status); - ctx.getExecutor().execute(() -> listener.onClose(result)); + pendingRpcs.remove(this); // eager removal; no-op if already drained } + ctx.getExecutor() + .execute( + () -> { + if (isCancelled) return; + isCancelled = true; + if (realCall != null) { + if (!onlyCancelPendingCall) { + realCall.cancel(status.getDescription(), status.getCause()); + } + } else { + listener.onClose(VRpcResult.createRejectedError(status)); + } + }); + } + + void cancelWithResult(VRpcResult result) { + ctx.getExecutor() + .execute( + () -> { + if (isCancelled) return; + isCancelled = true; + listener.onClose(result); + }); } @Override @@ -748,15 +751,19 @@ public void requestNext() { } private void drainTo(SessionHandle handle) { - synchronized (SessionPoolImpl.this) { - if (realCall == null) { - this.realCall = newRealCall(desc, handle); - } - } - this.realCall.start(req, ctx, listener); if (deadlineMonitor != null) { deadlineMonitor.cancel(); } + ctx.getExecutor() + .execute( + () -> { + if (isCancelled) { + SessionPoolImpl.this.onPendingVRpcCancelled(handle); + return; + } + realCall = newRealCall(desc, handle); + realCall.start(req, ctx, listener); + }); } private VRpcListener getListener() { @@ -907,16 +914,4 @@ public void close() { } } } - - private static final VRpc NOOP_CALL = - new VRpc() { - @Override - public void start(Object req, VRpcCallContext ctx, VRpcListener listener) {} - - @Override - public void cancel(@Nullable String message, @Nullable Throwable cause) {} - - @Override - public void requestNext() {} - }; } diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java index 641a63f0a573..e27f20c809d1 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/api/ClientTest.java @@ -55,7 +55,9 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +@Timeout(30) public class ClientTest { private ClientConfiguration defaultConfig; 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 70f2b20f8922..c3e4adaa2141 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 @@ -40,10 +40,12 @@ import javax.annotation.Nullable; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; +@Timeout(30) @ExtendWith(MockitoExtension.class) public class TableBaseTest { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java index a777a366e6aa..7aef572c3e7b 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -70,19 +70,23 @@ import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Answers; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +@Timeout(30) @ExtendWith(MockitoExtension.class) public class SessionImplTest { private ScheduledExecutorService executor; @@ -678,9 +682,8 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception { void abortFiresWhenListenerOnReadyThrows() throws Exception { SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); - java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); - java.util.concurrent.atomic.AtomicReference capturedStatus = - new java.util.concurrent.atomic.AtomicReference<>(); + CountDownLatch onCloseLatch = new CountDownLatch(1); + AtomicReference capturedStatus = new AtomicReference<>(); Session.Listener throwingListener = new Session.Listener() { @@ -719,8 +722,8 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai void abortDoesNotHangWhenListenerOnCloseThrows() throws Exception { SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); - java.util.concurrent.CountDownLatch onReadyLatch = new java.util.concurrent.CountDownLatch(1); - java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); + CountDownLatch onReadyLatch = new CountDownLatch(1); + CountDownLatch onCloseLatch = new CountDownLatch(1); Session.Listener throwingListener = new Session.Listener() { @@ -772,7 +775,7 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai void abortDoesNotInfiniteLoopWhenRecoveryListenerAlsoThrows() throws Exception { SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); - java.util.concurrent.CountDownLatch onCloseInvoked = new java.util.concurrent.CountDownLatch(1); + CountDownLatch onCloseInvoked = new CountDownLatch(1); Session.Listener doublyThrowingListener = new Session.Listener() { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java index def4986a95c1..442df965e809 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionListTest.java @@ -147,6 +147,40 @@ void testReadyInUseToIdle() { assertThat(stats.getExpectedCapacity()).isEqualTo(1); } + // Pending vRPC drained but cancelled before attaching to a real call: bookkeeping must mirror + // the OK-completion path (session returned to AFE ready list, counters restored), but the + // picker's latency stats must NOT be updated since nothing actually ran on the wire. + @Test + void testReadyInUseToIdleOnPendingCancelled() { + SessionList list = new SessionList(); + PoolStats stats = list.getStats(); + + fakeSession.state = SessionState.STARTING; + SessionHandle handle = list.newHandle(fakeSession); + + fakeSession.state = SessionState.READY; + handle.onSessionStarted(); + + SessionList.AfeHandle afe = list.getAfesWithReadySessions().get(0); + double e2eCostBefore = afe.getE2eCost(); + double transportCostBefore = afe.getTransportCost(); + + handle = list.checkoutSession(afe).get(); + handle.onPendingVRpcCancelled(); + + // Same pool bookkeeping as the OK-completion path. + assertThat(list.getAfesWithReadySessions().get(0).sessions).containsExactly(handle); + assertThat(list.getAllSessions()).containsExactly(handle); + assertThat(stats.getStartingCount()).isEqualTo(0); + assertThat(stats.getReadyCount()).isEqualTo(1); + assertThat(stats.getInUseCount()).isEqualTo(0); + assertThat(stats.getExpectedCapacity()).isEqualTo(1); + + // No latency update. + assertThat(afe.getE2eCost()).isEqualTo(e2eCostBefore); + assertThat(afe.getTransportCost()).isEqualTo(transportCostBefore); + } + @Test void testReadyIdleToSoftClosing() { SessionList list = new SessionList(); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 21419ce4bf09..84d2747449eb 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java @@ -79,15 +79,18 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.LongPredicate; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Answers; import org.mockito.ArgumentCaptor; @@ -95,6 +98,7 @@ import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; +@Timeout(30) @Nested @ExtendWith(MockitoExtension.class) public class SessionPoolImplTest { @@ -276,6 +280,48 @@ public void onClose(VRpcResult result) { } } + @Test + void pendingVRpcOnClosedPoolDoesNotLeakDeadlineMonitor() throws InterruptedException { + // Regression: PendingVRpc.start used to arm the deadline timer before the pool-state + // check, so the fast-fail "pool closed" branch leaked an armed timer that fired later + // and called listener.onClose a second time with DEADLINE_EXCEEDED. + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("close before issuing rpc") + .build()); + + CopyOnWriteArrayList closes = new CopyOnWriteArrayList<>(); + CountDownLatch firstClose = new CountDownLatch(1); + Duration deadline = Duration.ofMillis(100); + + VRpc vrpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + vrpc.start( + SessionFakeScriptedRequest.getDefaultInstance(), + VRpcCallContext.create( + Deadline.after(deadline.toMillis(), TimeUnit.MILLISECONDS), true, vrpcTracer), + new VRpc.VRpcListener() { + @Override + public void onMessage(SessionFakeScriptedResponse msg) {} + + @Override + public void onClose(VRpcResult result) { + closes.add(result); + firstClose.countDown(); + } + }); + + // The fast-fail UNAVAILABLE onClose should arrive immediately. + assertThat(firstClose.await(1, TimeUnit.SECONDS)).isTrue(); + assertThat(closes).hasSize(1); + + // Wait past the deadline. With the bug (leaked deadlineMonitor), a phantom + // onClose(DEADLINE_EXCEEDED) would arrive in this window. With the fix, no second close. + Thread.sleep(deadline.toMillis() * 5); + assertThat(closes).hasSize(1); + } + @Test void testCreateSessionDoesntPropagateDeadline() { DeadlineInterceptor deadlineInterceptor = new DeadlineInterceptor(); @@ -381,10 +427,9 @@ public void test() throws Exception { // retry-create-session site computes its delay against the real wall clock and the fake // budget clock, so it can land anywhere from sub-second to a couple of penalty intervals. // Match anything that isn't one of the two fixed cadences. - long watchdogMs = java.time.Duration.ofMinutes(5).toMillis(); + long watchdogMs = Duration.ofMinutes(5).toMillis(); long afePruneMs = SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(); - java.util.function.LongPredicate isRetrySchedule = - d -> d > 0 && d != watchdogMs && d != afePruneMs; + LongPredicate isRetrySchedule = d -> d > 0 && d != watchdogMs && d != afePruneMs; // start the pool sessionPool.start(