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 56d517f021a7..ec4cbd9a3f06 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 @@ -43,15 +43,27 @@ import io.opencensus.tags.Tags; import io.opentelemetry.sdk.OpenTelemetrySdk; import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; import java.util.Set; import java.util.WeakHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; public class Client implements AutoCloseable { + private static final Logger logger = Logger.getLogger(Client.class.getName()); + + // Per-pool drain budget during close. One full watchdog tick (5 min) plus 1 min buffer; if a + // pool can't drain in that window, something is genuinely wrong on the server side and we give + // up on it so close() returns. The watchdog interval is what makes the worst case finite. + private static final Duration POOL_DRAIN_TIMEOUT = Duration.ofMinutes(6); + public static final FeatureFlags BASE_FEATURE_FLAGS = FeatureFlags.newBuilder() .setReverseScans(false) @@ -91,6 +103,10 @@ public class Client implements AutoCloseable { private final Resource configManager; private final Set> sessionPools = Collections.newSetFromMap(new WeakHashMap<>()); + // Guarded by sessionPools' monitor: close() sets it before snapshotting the pool set, and the + // open* methods check it before adding a new pool, so a racing open cannot insert a pool that + // close() has already missed in its snapshot. + private boolean closed = false; public static Client create(ClientSettings settings) throws IOException { FeatureFlags featureFlags = @@ -218,78 +234,134 @@ public Client( @Override public void close() { - sessionPools.forEach( - pool -> - pool.close( - CloseSessionRequest.newBuilder() - .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. + List> toClose; + synchronized (sessionPools) { + if (closed) { + return; // idempotent + } + closed = true; + toClose = new ArrayList<>(sessionPools); + } + + CloseSessionRequest closeReq = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("Client closing") + .build(); + + // Phase 1: initiate graceful close on each pool. Returns immediately; sessions transition + // CLOSING → graceful CloseSessionRequest → WAIT_SERVER_CLOSE → CLOSED asynchronously. + toClose.forEach(p -> p.close(closeReq)); + + // Phase 2: wait for sessions to drain. The pool's watchdog stays alive during this wait and + // escalates anything stuck in WAIT_SERVER_CLOSE longer than its tick interval (5 min). Once + // a pool's last session reaches CLOSED, drainedFuture completes and awaitTerminated returns. + // Sequential: worst case is POOL_DRAIN_TIMEOUT * N pools, but the happy path drains in << 1s. + for (SessionPool pool : toClose) { + try { + if (!pool.awaitTerminated(POOL_DRAIN_TIMEOUT)) { + logger.warning( + "SessionPool did not drain within " + + POOL_DRAIN_TIMEOUT + + "; abandoning and continuing shutdown"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + logger.log(Level.WARNING, "Interrupted while draining SessionPool", e); + break; + } + } + + // 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(); } + // The closed check and pool insertion run under sessionPools' monitor so close() (which flips + // closed under the same monitor) cannot snapshot the pool set between our check and our insert. + // Opens are infrequent (typically once per table at app startup), so holding the monitor across + // createAndStart is acceptable. public TableAsync openTableAsync(String tableId, Permission permission) { - TableAsync tableAsync = - TableAsync.createAndStart( - featureFlags, - clientInfo, - configManager.get(), - channelPool, - defaultCallOptions, - tableId, - permission, - metrics.get(), - sessionTimer, - userCallbackExecutor.get()); - sessionPools.add(tableAsync.getSessionPool()); - return tableAsync; + synchronized (sessionPools) { + if (closed) { + throw new IllegalStateException("Client is closed"); + } + TableAsync tableAsync = + TableAsync.createAndStart( + featureFlags, + clientInfo, + configManager.get(), + channelPool, + defaultCallOptions, + tableId, + permission, + metrics.get(), + sessionTimer, + userCallbackExecutor.get()); + sessionPools.add(tableAsync.getSessionPool()); + return tableAsync; + } } public AuthorizedViewAsync openAuthorizedViewAsync( String tableId, String viewId, OpenAuthorizedViewRequest.Permission permission) { - AuthorizedViewAsync viewAsync = - AuthorizedViewAsync.createAndStart( - featureFlags, - clientInfo, - configManager.get(), - channelPool, - defaultCallOptions, - tableId, - viewId, - permission, - metrics.get(), - sessionTimer, - userCallbackExecutor.get()); - sessionPools.add(viewAsync.getSessionPool()); - return viewAsync; + synchronized (sessionPools) { + if (closed) { + throw new IllegalStateException("Client is closed"); + } + AuthorizedViewAsync viewAsync = + AuthorizedViewAsync.createAndStart( + featureFlags, + clientInfo, + configManager.get(), + channelPool, + defaultCallOptions, + tableId, + viewId, + permission, + metrics.get(), + sessionTimer, + userCallbackExecutor.get()); + sessionPools.add(viewAsync.getSessionPool()); + return viewAsync; + } } public MaterializedViewAsync openMaterializedViewAsync( String viewId, OpenMaterializedViewRequest.Permission permission) { - MaterializedViewAsync viewAsync = - MaterializedViewAsync.createAndStart( - featureFlags, - clientInfo, - configManager.get(), - channelPool, - defaultCallOptions, - viewId, - permission, - metrics.get(), - sessionTimer, - userCallbackExecutor.get()); - sessionPools.add(viewAsync.getSessionPool()); - return viewAsync; + synchronized (sessionPools) { + if (closed) { + throw new IllegalStateException("Client is closed"); + } + MaterializedViewAsync viewAsync = + MaterializedViewAsync.createAndStart( + featureFlags, + clientInfo, + configManager.get(), + channelPool, + defaultCallOptions, + viewId, + permission, + metrics.get(), + sessionTimer, + userCallbackExecutor.get()); + sessionPools.add(viewAsync.getSessionPool()); + return viewAsync; + } } public static class Resource { @@ -324,8 +396,10 @@ public T get() { } // 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) { + // so close() doesn't hang the caller on a pathological listener. Public so the compat + // ShimImpl (different package) can reuse the same shutdown semantics for the user-callback + // executor it owns. + public static void shutdownAndAwait(ExecutorService exec) { exec.shutdown(); try { if (!exec.awaitTermination(5, TimeUnit.SECONDS)) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index d1c7080828e8..24ae2e546939 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -178,7 +178,8 @@ public static Shim create( Resource.createShared(metrics), Resource.createShared(configManager), Resource.createShared(bgExecutor), - Resource.createOwned(userCallbackExecutor, userCallbackExecutor::shutdown)); + Resource.createOwned( + userCallbackExecutor, () -> Client.shutdownAndAwait(userCallbackExecutor))); return new ShimImpl(configManager, client); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java index 9102714457ae..163f79cac2ad 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/MutateRowShim.java @@ -22,28 +22,26 @@ import com.google.cloud.bigtable.data.v2.internal.api.Client; import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.compat.ShimImpl; -import com.google.cloud.bigtable.data.v2.internal.compat.Util; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; -import com.google.common.cache.LoadingCache; import io.grpc.Deadline; import java.io.IOException; import java.util.concurrent.CompletableFuture; public class MutateRowShim implements UnaryShim { - private final LoadingCache tables; - private final LoadingCache authViews; + private final SessionPoolMap tables; + private final SessionPoolMap authViews; public MutateRowShim(Client client) { tables = - Util.createSessionMap( + new SessionPoolMap<>( k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_WRITE)); authViews = - Util.createSessionMap( + new SessionPoolMap<>( k -> client.openAuthorizedViewAsync( k.getTableId(), @@ -63,9 +61,9 @@ public boolean supports(RowMutation request) { SessionPool pool; // TODO: avoid double lookup if (targetId instanceof TableId) { - pool = tables.getUnchecked((TableId) targetId).getSessionPool(); + pool = tables.get((TableId) targetId).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.getUnchecked((AuthorizedViewId) targetId).getSessionPool(); + pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); } else { return false; } @@ -83,16 +81,12 @@ public CompletableFuture call(RowMutation request, Deadline deadline) { SessionMutateRowRequest innerReq = request.toSessionProto(); if (targetId instanceof TableId) { - return tables - .getUnchecked((TableId) targetId) - .mutateRow(innerReq, deadline) - .thenApply(r -> null); + return tables.apply( + (TableId) targetId, t -> t.mutateRow(innerReq, deadline).thenApply(r -> null)); } if (targetId instanceof AuthorizedViewId) { - return authViews - .getUnchecked((AuthorizedViewId) targetId) - .mutateRow(innerReq, deadline) - .thenApply(r -> null); + return authViews.apply( + (AuthorizedViewId) targetId, v -> v.mutateRow(innerReq, deadline).thenApply(r -> null)); } CompletableFuture f = new CompletableFuture<>(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShimInner.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShimInner.java index 9dc134bd2392..85d1ab6284e7 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShimInner.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/ReadRowShimInner.java @@ -25,35 +25,33 @@ import com.google.cloud.bigtable.data.v2.internal.api.MaterializedViewAsync; import com.google.cloud.bigtable.data.v2.internal.api.TableAsync; import com.google.cloud.bigtable.data.v2.internal.compat.ShimImpl; -import com.google.cloud.bigtable.data.v2.internal.compat.Util; import com.google.cloud.bigtable.data.v2.internal.session.SessionPool; import com.google.cloud.bigtable.data.v2.models.AuthorizedViewId; import com.google.cloud.bigtable.data.v2.models.MaterializedViewId; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.TableId; import com.google.cloud.bigtable.data.v2.models.TargetId; -import com.google.common.cache.LoadingCache; import io.grpc.Deadline; import java.util.concurrent.CompletableFuture; public class ReadRowShimInner implements UnaryShim { - private final LoadingCache tables; - private final LoadingCache authViews; - private final LoadingCache matViews; + private final SessionPoolMap tables; + private final SessionPoolMap authViews; + private final SessionPoolMap matViews; public ReadRowShimInner(Client client) { tables = - Util.createSessionMap( + new SessionPoolMap<>( k -> client.openTableAsync(k.getTableId(), Permission.PERMISSION_READ)); authViews = - Util.createSessionMap( + new SessionPoolMap<>( k -> client.openAuthorizedViewAsync( k.getTableId(), k.getAuthorizedViewId(), OpenAuthorizedViewRequest.Permission.PERMISSION_READ)); matViews = - Util.createSessionMap( + new SessionPoolMap<>( k -> client.openMaterializedViewAsync( k.getMaterializedViewId(), @@ -73,11 +71,11 @@ public boolean supports(Query request) { SessionPool pool; // TODO avoid double lookup if (targetId instanceof TableId) { - pool = tables.getUnchecked((TableId) targetId).getSessionPool(); + pool = tables.get((TableId) targetId).getSessionPool(); } else if (targetId instanceof AuthorizedViewId) { - pool = authViews.getUnchecked((AuthorizedViewId) targetId).getSessionPool(); + pool = authViews.get((AuthorizedViewId) targetId).getSessionPool(); } else if (targetId instanceof MaterializedViewId) { - pool = matViews.getUnchecked((MaterializedViewId) targetId).getSessionPool(); + pool = matViews.get((MaterializedViewId) targetId).getSessionPool(); } else { return false; } @@ -95,13 +93,13 @@ public CompletableFuture call(Query query, Deadline dead SessionReadRowRequest innerReq = query.toSessionPointProto(); if (targetId instanceof TableId) { - return tables.getUnchecked((TableId) targetId).readRow(innerReq, deadline); + return tables.apply((TableId) targetId, t -> t.readRow(innerReq, deadline)); } if (targetId instanceof AuthorizedViewId) { - return authViews.getUnchecked((AuthorizedViewId) targetId).readRow(innerReq, deadline); + return authViews.apply((AuthorizedViewId) targetId, v -> v.readRow(innerReq, deadline)); } if (targetId instanceof MaterializedViewId) { - return matViews.getUnchecked((MaterializedViewId) targetId).readRow(innerReq, deadline); + return matViews.apply((MaterializedViewId) targetId, v -> v.readRow(innerReq, deadline)); } CompletableFuture f = new CompletableFuture<>(); diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java new file mode 100644 index 000000000000..a1f8257922b3 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java @@ -0,0 +1,84 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.data.v2.internal.compat.ops; + +import com.google.cloud.bigtable.data.v2.internal.compat.Util; +import com.google.common.cache.LoadingCache; +import com.google.common.util.concurrent.UncheckedExecutionException; +import java.io.Closeable; +import java.util.concurrent.CompletableFuture; +import java.util.function.Function; + +/** + * Lazily-loaded map from target IDs to per-target handles (TableAsync / AuthorizedViewAsync / + * MaterializedViewAsync, each holding a SessionPool), backed by a Guava {@link LoadingCache}. + * Centralizes the conversion from Guava's {@link UncheckedExecutionException} (which wraps any + * loader throw) into the appropriate shape for each call site: a raw {@link RuntimeException} for + * sync inspections, or a failed {@link CompletableFuture} for async ops. + * + *

The most common loader failure is {@link IllegalStateException} from {@code + * Client.openTableAsync} after {@code Client.close}; without unwrapping, callers see {@code + * UncheckedExecutionException} instead of the documented exception type and async callers see a + * thrown exception instead of a failed future. + */ +public class SessionPoolMap { + private final LoadingCache cache; + + public SessionPoolMap(Function loader) { + this.cache = Util.createSessionMap(loader); + } + + /** Returns the cached value, loading on miss. Loader throws surface as the original cause. */ + public V get(K key) { + try { + return cache.getUnchecked(key); + } catch (UncheckedExecutionException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + if (cause instanceof RuntimeException) throw (RuntimeException) cause; + if (cause instanceof Error) throw (Error) cause; + throw new RuntimeException(cause); + } + } + + /** + * Looks up the cached value and applies {@code op}. Both the lookup (e.g. {@link + * IllegalStateException} from a closed Client) and a synchronous throw from {@code op.apply} + * (e.g. NPE on malformed input, RejectedExecutionException during shutdown) are converted to a + * failed future so callers consistently observe failures through the future surface. + */ + public CompletableFuture apply(K key, Function> op) { + V v; + try { + v = get(key); + } catch (Exception e) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(e); + return f; + } + try { + return op.apply(v); + } catch (Throwable t) { + CompletableFuture f = new CompletableFuture<>(); + f.completeExceptionally(t); + return f; + } + } + + /** Evicts all entries, triggering Closeable.close on each via the cache's removal listener. */ + public void invalidateAll() { + cache.invalidateAll(); + } +} 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 0c9d2887d87a..614f45eb2eb4 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 @@ -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,41 @@ 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(); } 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 67d7209a7cf4..a11a1e5cbbd1 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 @@ -93,17 +93,26 @@ public void start(ReqT req, VRpcListener listener) { 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). + // start. 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()); - } + // Queueing the registration onto the op executor is what makes the closed-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. + // + // 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) { + grpcContext.addListener(cancellationListener, MoreExecutors.directExecutor()); + } + }); } @Override diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java index 48930ece77e1..e3edd34fd606 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/BigtableTimer.java @@ -47,14 +47,33 @@ 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}. + * + *

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. + * + *

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(); + } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java index 6d66b06ff09d..16ffd0127b10 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/NettyWheelTimer.java +++ b/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 stopHooks = ConcurrentHashMap.newKeySet(); + private volatile boolean stopped = false; + public NettyWheelTimer(String name, Executor dispatcher) { this.dispatcher = dispatcher; this.delegate = @@ -49,11 +61,38 @@ public NettyWheelTimer(String name, Executor dispatcher) { @Override public Timeout newTimeout(Runnable task, long delay, TimeUnit 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 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(); } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPool.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPool.java index 0b8cd1cdaea6..9ff1d6ffe9f8 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPool.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPool.java @@ -20,6 +20,7 @@ import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc; import com.google.protobuf.Message; import io.grpc.Metadata; +import java.time.Duration; public interface SessionPool { void start(OpenReqT openReq, Metadata md); @@ -29,6 +30,12 @@ VRpc newCall( void close(CloseSessionRequest req); + /** + * Blocks until all sessions in this pool have terminated, or the timeout elapses. Must be called + * after {@link #close} to be meaningful. Returns true if drained, false on timeout. + */ + boolean awaitTerminated(Duration timeout) throws InterruptedException; + SessionPoolInfo getInfo(); int getConsecutiveUnimplementedFailures(); 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 2dfd73f1ba7f..04f31e33c335 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 @@ -64,7 +64,10 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -133,6 +136,11 @@ private enum PoolState { @GuardedBy("this") private boolean closed = false; + // Completed when this pool has been close()d AND every session has reached the CLOSED terminal + // state. Drives Client.close()'s drain barrier so that listener.onClose tasks finish queueing + // onto userCallbackExecutor before that executor is shut down. + private final CompletableFuture drainedFuture = new CompletableFuture<>(); + @GuardedBy("this") private BigtableTimer.Timeout retryCreateSessionFuture = null; @@ -293,8 +301,13 @@ public void close(CloseSessionRequest req) { retryCreateSessionFuture.cancel(); retryCreateSessionFuture = null; } - watchdog.close(); + // Watchdog stays alive past close() so it can escalate any session that lingers in + // WAIT_SERVER_CLOSE during shutdown. awaitTerminated() takes ownership of closing it. sessions.close(req); + // If the pool had no sessions, drainedFuture would never be completed by onSessionClose. + if (sessions.getAllSessions().isEmpty()) { + drainedFuture.complete(null); + } } // cancelWithResult trampolines through ctx.getExecutor() — required because the public @@ -308,6 +321,25 @@ public void close(CloseSessionRequest req) { } } + @Override + public boolean awaitTerminated(Duration timeout) throws InterruptedException { + try { + drainedFuture.get(timeout.toNanos(), TimeUnit.NANOSECONDS); + return true; + } catch (TimeoutException e) { + return false; + } catch (ExecutionException e) { + // drainedFuture is only completed via .complete(null), never .completeExceptionally — + // a CancellationException would still be wrapped here. Treat as a bug. + debugTagTracer.record( + TelemetryConfiguration.Level.ERROR, "session_pool_drained_future_failed"); + throw new IllegalStateException("drainedFuture failed unexpectedly", e); + } finally { + // Close the watchdog on the way out — drained or timed out, its job is done. + watchdog.close(); + } + } + @Override public synchronized void start(OpenReqT openReq, Metadata md) { Preconditions.checkState(poolState == PoolState.NEW); @@ -515,6 +547,10 @@ private void onSessionClose( // If the pool is closed then there is nothing else to do // dont need to create a replacement session and pending vRpcs get cleaned up in close() if (poolState == PoolState.CLOSED) { + // Signal awaitTerminated() once the last session has drained. + if (sessions.getAllSessions().isEmpty()) { + drainedFuture.complete(null); + } return; } 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 e27f20c809d1..eb2e0f78c8a0 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 @@ -107,6 +107,41 @@ void tearDown() { executor.shutdownNow(); } + @Test + public void openAfterCloseThrows() { + client.close(); + + IllegalStateException tableEx = + assertThrows( + IllegalStateException.class, + () -> client.openTableAsync("fake-table", OpenTableRequest.Permission.PERMISSION_READ)); + assertThat(tableEx).hasMessageThat().contains("closed"); + + IllegalStateException viewEx = + assertThrows( + IllegalStateException.class, + () -> + client.openAuthorizedViewAsync( + "fake-table", + "fake-view", + OpenAuthorizedViewRequest.Permission.PERMISSION_READ)); + assertThat(viewEx).hasMessageThat().contains("closed"); + + IllegalStateException mvEx = + assertThrows( + IllegalStateException.class, + () -> + client.openMaterializedViewAsync( + "fake-view", OpenMaterializedViewRequest.Permission.PERMISSION_READ)); + assertThat(mvEx).hasMessageThat().contains("closed"); + } + + @Test + public void closeIsIdempotent() { + client.close(); + client.close(); // must not throw or hang + } + @Test public void testRequestFails() { TableAsync table = 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 c3e4adaa2141..60c706a7af9d 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 @@ -35,6 +35,7 @@ import com.google.protobuf.Message; import io.grpc.Deadline; import io.grpc.Metadata; +import java.time.Duration; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import javax.annotation.Nullable; @@ -178,6 +179,11 @@ public void start(OpenTableRequest openReq, Metadata md) {} @Override public void close(CloseSessionRequest req) {} + @Override + public boolean awaitTerminated(Duration timeout) { + return true; + } + @Override public SessionPoolInfo getInfo() { return SessionPoolInfo.create(clientInfo, VRpcDescriptor.TABLE_SESSION, "fake-pool"); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java new file mode 100644 index 000000000000..f6e22c7ed392 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java @@ -0,0 +1,119 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.data.v2.internal.compat.ops; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.Closeable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +class SessionPoolMapTest { + + // Minimal Closeable value type — tracks how many times close() was called so invalidateAll's + // removal-listener wiring is observable from tests. + private static final class CountingHandle implements Closeable { + final AtomicInteger closeCount = new AtomicInteger(); + + @Override + public void close() { + closeCount.incrementAndGet(); + } + } + + @Test + void get_unwrapsLoaderRuntimeException() { + // Production trigger: Client.openTableAsync throws IllegalStateException after Client.close. + // The Guava LoadingCache wraps it in UncheckedExecutionException; SessionPoolMap.get must + // surface the original IllegalStateException so callers can pattern-match on it. + SessionPoolMap map = + new SessionPoolMap<>( + key -> { + throw new IllegalStateException("Client is closed"); + }); + + IllegalStateException thrown = assertThrows(IllegalStateException.class, () -> map.get("k")); + assertThat(thrown).hasMessageThat().isEqualTo("Client is closed"); + } + + @Test + void apply_convertsLoaderThrowToFailedFuture() throws Exception { + // The async surface must not throw — callers wire onFailure handlers on the returned future. + SessionPoolMap map = + new SessionPoolMap<>( + key -> { + throw new IllegalStateException("Client is closed"); + }); + + CompletableFuture result = + map.apply("k", v -> CompletableFuture.completedFuture("unreached")); + + assertThat(result.isCompletedExceptionally()).isTrue(); + ExecutionException ee = + assertThrows(ExecutionException.class, () -> result.get(1, TimeUnit.SECONDS)); + assertThat(ee).hasCauseThat().isInstanceOf(IllegalStateException.class); + assertThat(ee).hasCauseThat().hasMessageThat().isEqualTo("Client is closed"); + } + + @Test + void apply_convertsOpSyncThrowToFailedFuture() throws Exception { + // op.apply may throw synchronously (NPE on malformed input, REE during shutdown, any + // RuntimeException from the wrapped call chain). The async surface must not propagate + // the throw — same contract as the loader-throw path above. + CountingHandle handle = new CountingHandle(); + SessionPoolMap map = new SessionPoolMap<>(key -> handle); + + CompletableFuture result = + map.apply( + "k", + v -> { + throw new IllegalStateException("op blew up"); + }); + + assertThat(result.isCompletedExceptionally()).isTrue(); + ExecutionException ee = + assertThrows(ExecutionException.class, () -> result.get(1, TimeUnit.SECONDS)); + assertThat(ee).hasCauseThat().isInstanceOf(IllegalStateException.class); + assertThat(ee).hasCauseThat().hasMessageThat().isEqualTo("op blew up"); + } + + @Test + void apply_happyPathInvokesOp() throws Exception { + CountingHandle handle = new CountingHandle(); + SessionPoolMap map = new SessionPoolMap<>(key -> handle); + + CompletableFuture result = map.apply("k", v -> CompletableFuture.completedFuture("ok")); + + assertThat(result.get(1, TimeUnit.SECONDS)).isEqualTo("ok"); + } + + @Test + void invalidateAll_closesCachedValues() { + CountingHandle handle = new CountingHandle(); + SessionPoolMap map = new SessionPoolMap<>(key -> handle); + + // Populate the cache so invalidateAll has something to evict. + map.get("k"); + assertThat(handle.closeCount.get()).isEqualTo(0); + + map.invalidateAll(); + assertThat(handle.closeCount.get()).isEqualTo(1); + } +} 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 new file mode 100644 index 000000000000..a69bed4c0462 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java @@ -0,0 +1,155 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.bigtable.data.v2.internal.middleware; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.cloud.bigtable.data.v2.internal.csm.NoopMetrics; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcCallContext; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcListener; +import com.google.cloud.bigtable.data.v2.internal.middleware.VRpc.VRpcResult; +import com.google.common.util.concurrent.MoreExecutors; +import io.grpc.Context; +import io.grpc.Deadline; +import io.grpc.Status; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import javax.annotation.Nullable; +import org.junit.jupiter.api.Test; + +class VOperationImplTest { + + @Test + void grpcContextCancelPropagatesToChain() throws InterruptedException { + // Normal-path sanity: addListener fires (chain.start doesn't queue an onClose), then a + // later grpcContext.cancel routes through cancellationListener -> opExecutor -> chain.cancel. + // This exercises the new runInline-based addListener path introduced to close the TOCTOU + // between wrapped.closed and addListener. + Context.CancellableContext grpcContext = Context.current().withCancellation(); + AtomicInteger chainCancelCount = new AtomicInteger(); + CountDownLatch cancelLatch = new CountDownLatch(1); + + VRpc chain = + new VRpc() { + @Override + public void start(String req, VRpcCallContext ctx, VRpcListener listener) { + // Hold open — do not call listener.onClose. Cancellation must drive termination. + } + + @Override + public void cancel(@Nullable String msg, @Nullable Throwable cause) { + chainCancelCount.incrementAndGet(); + cancelLatch.countDown(); + } + + @Override + public void requestNext() {} + }; + + VOperationImpl op = + new VOperationImpl<>( + chain, + grpcContext, + MoreExecutors.directExecutor(), + NoopMetrics.NoopVrpcTracer.INSTANCE, + Deadline.after(10, TimeUnit.SECONDS), + true); + + op.start( + "req", + new VRpcListener() { + @Override + public void onMessage(String msg) {} + + @Override + public void onClose(VRpcResult result) {} + }); + + grpcContext.cancel(Status.CANCELLED.withDescription("test").asException()); + + assertThat(cancelLatch.await(2, TimeUnit.SECONDS)).isTrue(); + assertThat(chainCancelCount.get()).isEqualTo(1); + } + + @Test + void asyncOnCloseFromChainDoesNotPropagateLaterContextCancel() throws InterruptedException { + // Regression for the wrapped.closed 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, + // 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(); + AtomicInteger chainCancelCount = new AtomicInteger(); + AtomicReference userClose = new AtomicReference<>(); + CountDownLatch onCloseLatch = new CountDownLatch(1); + + VRpc chain = + new VRpc() { + @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")))); + } + + @Override + public void cancel(@Nullable String msg, @Nullable Throwable cause) { + chainCancelCount.incrementAndGet(); + } + + @Override + public void requestNext() {} + }; + + VOperationImpl op = + new VOperationImpl<>( + chain, + grpcContext, + MoreExecutors.directExecutor(), + NoopMetrics.NoopVrpcTracer.INSTANCE, + Deadline.after(10, TimeUnit.SECONDS), + true); + + op.start( + "req", + new VRpcListener() { + @Override + public void onMessage(String msg) {} + + @Override + public void onClose(VRpcResult result) { + userClose.set(result); + onCloseLatch.countDown(); + } + }); + + assertThat(onCloseLatch.await(2, TimeUnit.SECONDS)).isTrue(); + assertThat(userClose.get().getStatus().getCode()).isEqualTo(Status.UNAVAILABLE.getCode()); + + grpcContext.cancel(Status.CANCELLED.withDescription("test").asException()); + Thread.sleep(50); // give any leaked listener a chance to fire + + // No chain.cancel — the cancellationListener was correctly skipped because the chain had + // already reached its terminal state via the queued onClose. + assertThat(chainCancelCount.get()).isEqualTo(0); + } +} 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 7aef572c3e7b..59d9b6e3e702 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 @@ -843,6 +843,11 @@ public boolean isCancelled() { }; } + @Override + public Registration onStop(Runnable hook) { + return delegate.onStop(hook); + } + @Override public void stop() { delegate.stop(); 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 84d2747449eb..64d9bfbb9ebb 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 @@ -73,6 +73,7 @@ import io.grpc.ServerCall; import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; +import io.grpc.Status; import java.io.IOException; import java.time.Duration; import java.time.Instant; @@ -99,7 +100,6 @@ import org.mockito.junit.jupiter.MockitoExtension; @Timeout(30) -@Nested @ExtendWith(MockitoExtension.class) public class SessionPoolImplTest { private static final ClientInfo CLIENT_INFO = @@ -170,12 +170,15 @@ void setUp() throws IOException { } @AfterEach - void tearDown() { + void tearDown() throws InterruptedException { sessionPool.close( CloseSessionRequest.newBuilder() .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) .setDescription("close session") .build()); + // Wait for sessions to drain so the watchdog can be closed before testTimer.stop() races + // with its self-reschedule loop. + sessionPool.awaitTerminated(Duration.ofSeconds(10)); channelPool.close(); // channel gets shutdown in channelPool.close() server.shutdownNow(); @@ -280,6 +283,126 @@ public void onClose(VRpcResult result) { } } + @Test + void awaitTerminatedReturnsTrueWhenPoolIsEmpty() throws InterruptedException { + // A pool that was never started has no sessions; close() should complete drainedFuture + // immediately and awaitTerminated should return true without blocking. + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("empty pool") + .build()); + assertThat(sessionPool.awaitTerminated(Duration.ofMillis(100))).isTrue(); + } + + @Test + void awaitTerminatedReturnsTrueAfterSessionsDrain() + throws InterruptedException, ExecutionException, TimeoutException { + // Start a real session, issue + complete a vRPC so the session is fully open, then close + // the pool and verify awaitTerminated returns true (sessions cleanly drained). + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + VRpc vrpc = + sessionPool.newCall(FakeDescriptor.SCRIPTED); + UnaryResponseFuture resultFuture = new UnaryResponseFuture<>(); + vrpc.start( + SessionFakeScriptedRequest.getDefaultInstance(), + VRpcCallContext.create(Deadline.after(10, TimeUnit.SECONDS), true, vrpcTracer), + resultFuture); + resultFuture.get(10, TimeUnit.SECONDS); + + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("after drain") + .build()); + + assertThat(sessionPool.awaitTerminated(Duration.ofSeconds(10))).isTrue(); + } + + @Test + void closeCancelsInFlightVRpcAndDrains() throws InterruptedException { + // Park the session-open response in the client interceptor so the vRPC cannot drain to a + // ready session and sits in pendingRpcs. Then close the pool and verify (a) the queued + // vRPC's listener gets a CANCELLED onClose with the "SessionPool closed:" description from + // SessionPoolImpl.close()'s cancellation loop, (b) awaitTerminated returns true (the + // OPENING session is driven to CLOSED by sessions.close(req), not parked until the + // watchdog tick), and (c) no second onClose arrives once the delayed responses finally + // flush past the interceptor. + SessionPool testSessionPool = null; + try (ChannelPool delayedPool = + new SingleChannelPool( + Suppliers.ofInstance( + Grpc.newChannelBuilderForAddress( + "localhost", server.getPort(), InsecureChannelCredentials.create()) + .intercept(new DelayedClientInterceptor(Duration.ofMillis(100))) + .build()))) { + + delayedPool.start(); + + testSessionPool = + new SessionPoolImpl<>( + metrics, + FeatureFlags.getDefaultInstance(), + CLIENT_INFO, + configManager, + delayedPool, + CallOptions.DEFAULT, + FakeDescriptor.FAKE_SESSION, + "fake-pool-in-flight", + testTimer); + + testSessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + VRpc vrpc = + testSessionPool.newCall(FakeDescriptor.SCRIPTED); + CopyOnWriteArrayList closes = new CopyOnWriteArrayList<>(); + CountDownLatch firstClose = new CountDownLatch(1); + Duration deadline = Duration.ofSeconds(10); + 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(); + } + }); + + // vRPC is now queued in pendingRpcs behind the parked session-open. Close the pool. + testSessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("close while in flight") + .build()); + + // Cancellation hops through the vRPC's per-op SerializingExecutor — not the gRPC + // listener thread — so it should arrive promptly even while the interceptor is asleep. + assertThat(firstClose.await(2, TimeUnit.SECONDS)).isTrue(); + assertThat(closes).hasSize(1); + VRpcResult result = closes.get(0); + assertThat(result).status().code().isEqualTo(Status.Code.CANCELLED); + assertThat(result).status().description().contains("SessionPool closed:"); + + // OPENING session drains via sessions.close(req) once the delayed onMessages flush. + assertThat(testSessionPool.awaitTerminated(Duration.ofSeconds(10))).isTrue(); + + // Regression guard: the delayed session-open responses arriving post-cancellation must + // not produce a phantom second onClose on the cancelled vRPC's listener. + Thread.sleep(200); + assertThat(closes).hasSize(1); + } finally { + if (testSessionPool != null) { + testSessionPool.close(CloseSessionRequest.getDefaultInstance()); + } + } + } + @Test void pendingVRpcOnClosedPoolDoesNotLeakDeadlineMonitor() throws InterruptedException { // Regression: PendingVRpc.start used to arm the deadline timer before the pool-state diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java index 3582e66e1971..790d997f4295 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/WatchdogTest.java +++ b/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() {} }