From 2564867ccf0e777856b1d12e63df5e992b77afa6 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 15 Jun 2026 19:19:40 +0000 Subject: [PATCH 1/9] fix: drain SessionPools before tearing down userCallbackExecutor on Client.close MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Client.close shut down userCallbackExecutor before draining the SessionPools that depend on it, so late listener.onClose tasks from in-flight RPCs arrived after backing was dead and got RejectedExecutionException — silently stranding the user's terminal callbacks. The earlier fix sprinkled inline-drain fallbacks inside OpExecutor; restructure shutdown instead so the race can't happen. SessionPool gains awaitTerminated(Duration), backed by a CompletableFuture SessionPoolImpl completes from onSessionClose once the pool is CLOSED and the last session has drained. close() no longer kills the watchdog — awaitTerminated takes ownership of that, so the watchdog stays alive during shutdown and can escalate any session stuck in WAIT_SERVER_CLOSE longer than its tick interval (5 min) via forceClose. Client.close becomes three explicit phases: (1) initiate graceful close on each pool, (2) awaitTerminated on each with a 6-minute per-pool budget (one full watchdog tick plus buffer), (3) tear down userCallbackExecutor / channelPool / timers in the existing order, now safely because all listener.onClose tasks are queued or drained before backing dies. VOperationImpl.start queues grpcContext.addListener through the op executor via runInline. Without this, an async-queued onClose from chain.start (PendingVRpc pool-closed fast-fail, VRpcImpl deadline-exceeded short-circuit) could drain between the CleanupListener.closed read and addListener: CleanupListener.onClose would call removeListener as a no-op pre-registration, and the caller would then register a listener with nothing to remove it. The leak is per-RPC and permanent until grpcContext cancels — for long-lived application contexts it accumulates indefinitely. FIFO ordering through the op executor makes the closed-check sound: any onClose chain.start enqueued drains first, so the check is accurate by the time we evaluate it. Add a `closed` AtomicBoolean + checkNotClosed() guard on the three openers so concurrent opens during shutdown can't create pools the close path won't see. close() is now idempotent via CAS on that flag. Tests: - ClientTest#openAfterCloseThrows / closeIsIdempotent - SessionPoolImplTest awaitTerminated* coverage - VOperationImplTest covering async onClose / context cancel ordering - SessionPoolImplTest tearDown now calls awaitTerminated so the watchdog is closed before testTimer.stop races its self-reschedule - FakeSessionPool in TableBaseTest gains a no-op awaitTerminated stub Drive-by: remove the spurious @Nested annotation from SessionPoolImplTest's top-level class. @Nested is meaningful only on non-static inner classes; on the outer class it caused Surefire to mis-attribute test counts. --- .../bigtable/data/v2/internal/api/Client.java | 80 +++++++-- .../internal/middleware/VOperationImpl.java | 29 ++-- .../data/v2/internal/session/SessionPool.java | 7 + .../v2/internal/session/SessionPoolImpl.java | 36 +++- .../data/v2/internal/api/ClientTest.java | 35 ++++ .../data/v2/internal/api/TableBaseTest.java | 6 + .../middleware/VOperationImplTest.java | 155 ++++++++++++++++++ .../internal/session/SessionPoolImplTest.java | 43 ++++- 8 files changed, 367 insertions(+), 24 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/middleware/VOperationImplTest.java 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..3226b52dfa2a 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,28 @@ 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.concurrent.atomic.AtomicBoolean; +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 +104,10 @@ public class Client implements AutoCloseable { private final Resource configManager; private final Set> sessionPools = Collections.newSetFromMap(new WeakHashMap<>()); + // Set true at the start of close(); guards openTableAsync / openAuthorizedViewAsync / + // openMaterializedViewAsync so concurrent opens during shutdown don't create pools the close + // path won't see. + private final AtomicBoolean closed = new AtomicBoolean(false); public static Client create(ClientSettings settings) throws IOException { FeatureFlags featureFlags = @@ -218,17 +235,49 @@ 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. + if (!closed.compareAndSet(false, true)) { + return; // idempotent + } + + List> toClose; + synchronized (sessionPools) { + 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. By this point all listener.onClose tasks for in-flight + // RPCs are queued on their op executors (which run on userCallbackExecutor), and no new + // session responses are coming since every session is CLOSED. The 5s await inside + // userCallbackExecutor.close() is therefore just a guard for tasks in flight — it should + // return immediately in the typical case. userCallbackExecutor.close(); metrics.close(); channelPool.close(); @@ -238,7 +287,14 @@ public void close() { backgroundExecutor.close(); } + private void checkNotClosed() { + if (closed.get()) { + throw new IllegalStateException("Client is closed"); + } + } + public TableAsync openTableAsync(String tableId, Permission permission) { + checkNotClosed(); TableAsync tableAsync = TableAsync.createAndStart( featureFlags, @@ -257,6 +313,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) { public AuthorizedViewAsync openAuthorizedViewAsync( String tableId, String viewId, OpenAuthorizedViewRequest.Permission permission) { + checkNotClosed(); AuthorizedViewAsync viewAsync = AuthorizedViewAsync.createAndStart( featureFlags, @@ -276,6 +333,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync( public MaterializedViewAsync openMaterializedViewAsync( String viewId, OpenMaterializedViewRequest.Permission permission) { + checkNotClosed(); MaterializedViewAsync viewAsync = MaterializedViewAsync.createAndStart( featureFlags, 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/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..79b15214695f 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,23 @@ 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. + 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 +545,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/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/SessionPoolImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionPoolImplTest.java index 84d2747449eb..185e1b32305e 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 @@ -99,7 +99,6 @@ import org.mockito.junit.jupiter.MockitoExtension; @Timeout(30) -@Nested @ExtendWith(MockitoExtension.class) public class SessionPoolImplTest { private static final ClientInfo CLIENT_INFO = @@ -170,12 +169,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 +282,43 @@ 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 pendingVRpcOnClosedPoolDoesNotLeakDeadlineMonitor() throws InterruptedException { // Regression: PendingVRpc.start used to arm the deadline timer before the pool-state From 8ed890598360d6453f337ba30ea675fbfa69deac Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 15 Jun 2026 22:24:07 +0000 Subject: [PATCH 2/9] fix: deliver terminal onClose to Scheduled retries pending at Client.close Scheduled RetryingVRpcs hold no session reference, so a long-delay retry (server-driven RetryInfo.retryDelay) outlives Phase 2 drain and is silently discarded when sessionTimer.stop() runs in Phase 3. The user's listener never fires. Add an onStop hook primitive to BigtableTimer. Scheduled.onStart registers a hook on entry and unregisters on every exit path (normal fire, cancel, hook fire). NettyWheelTimer.stop() runs every hook synchronously before discarding pending wheel timeouts; hooks trampoline back through the op executor to drive Scheduled to a CANCELLED Done. Reorder Client.close Phase 3 so sessionTimer.stop() runs before userCallbackExecutor.close(), giving the hook-fired onClose tasks a live op-executor backing to land on. Also replace Scheduled.onStart's dead RejectedExecutionException catch with IllegalStateException, matching BigtableTimer.stop()'s documented post-condition. --- .../bigtable/data/v2/internal/api/Client.java | 19 ++++--- .../v2/internal/middleware/RetryingVRpc.java | 50 ++++++++++++++++--- .../v2/internal/session/BigtableTimer.java | 22 +++++++- .../v2/internal/session/NettyWheelTimer.java | 42 +++++++++++++++- .../v2/internal/session/SessionImplTest.java | 5 ++ .../v2/internal/session/WatchdogTest.java | 5 ++ 6 files changed, 126 insertions(+), 17 deletions(-) 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 3226b52dfa2a..3e7da51fd906 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 @@ -273,17 +273,22 @@ public void close() { } } - // Phase 3: tear down infrastructure. By this point all listener.onClose tasks for in-flight - // RPCs are queued on their op executors (which run on userCallbackExecutor), and no new - // session responses are coming since every session is CLOSED. The 5s await inside - // userCallbackExecutor.close() is therefore just a guard for tasks in flight — it should - // return immediately in the typical case. + // Phase 3: tear down infrastructure. + // + // sessionTimer.stop() runs FIRST so its onStop hooks can drive any pending Scheduled retries + // to a terminal Done — that delivery hops through op executor → userCallbackExecutor, both + // of which must still be alive at this moment. + // + // userCallbackExecutor.close() next, with a 5s drain to catch the listener.onClose tasks + // queued by both the session drain (Phase 2) and the just-fired retry shutdowns. + // + // backgroundExecutor must close last because it's the timer's dispatcher and the op + // executor's chain ultimately runs ScheduledExecutorService tasks here. + sessionTimer.stop(); userCallbackExecutor.close(); metrics.close(); channelPool.close(); configManager.close(); - // Stop the timer before tearing down backgroundExecutor (the timer's dispatcher). - sessionTimer.stop(); backgroundExecutor.close(); } 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..277637ac41c8 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,38 @@ public void onStart() { } } + private void onTimerFired() { + unregisterStopHook(); + onStateChange(new Idle()); + } + + // Invoked from BigtableTimer.stop on the close thread. Trampoline back to the op executor so + // currentState reads and onStateChange are still single-threaded with the rest of the chain. + private void onTimerStopping() { + context.getExecutor().execute(() -> { + if (currentState != Scheduled.this) { + return; // already transitioned out via normal fire or cancel + } + onStateChange( + new Done( + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription( + "Client closing while retry pending")))); + }); + } + + private void unregisterStopHook() { + if (stopHook != null) { + stopHook.unregister(); + stopHook = null; + } + } + @Override public void onCancel(String reason, Throwable throwable) { - // future can be null if schedule throws an exception that's not RejectedExecutionException. - // In which case sync context uncaught exception handler will be called, which calls cancel on - // the current state before transition into done state. + unregisterStopHook(); + // future can be null if schedule throws and we end up here via the op-executor uncaught + // path. if (future != null && !future.isCancelled()) { future.cancel(); } 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..6731355a63e6 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,34 @@ public interface BigtableTimer { /** * Releases the tick thread and discards any pending timeouts. Idempotent. After {@code stop()}, - * subsequent calls to {@link #newTimeout} throw {@link IllegalStateException}. + * subsequent calls to {@link #newTimeout} or {@link #onStop} throw {@link + * IllegalStateException}. + * + *

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..442b4efd5f43 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,39 @@ public NettyWheelTimer(String name, Executor dispatcher) { @Override public Timeout newTimeout(Runnable task, long delay, TimeUnit unit) { - return new TimeoutHandle(delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit)); + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + return new TimeoutHandle( + delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit)); + } + + @Override + public Registration onStop(Runnable hook) { + if (stopped) { + throw new IllegalStateException("timer stopped"); + } + stopHooks.add(hook); + return () -> stopHooks.remove(hook); } @Override public void stop() { + if (stopped) { + return; + } + stopped = true; + // Snapshot then clear so hooks can no longer be unregistered while we iterate (and so a hook + // that re-enters onStop sees stopped=true and fails fast). + Set 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/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/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() {} } From 02f497d7fb29640946286755b3e4d3dd111e3813 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 15 Jun 2026 22:29:47 +0000 Subject: [PATCH 3/9] fix: serialize open*/close so racing opens cannot create orphan pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit openTableAsync / openAuthorizedViewAsync / openMaterializedViewAsync read 'closed' lock-free, then constructed the pool, then inserted it into sessionPools — close() could CAS closed=true and snapshot sessionPools in between, leaving the new pool orphaned: never closed, its callbacks landing on shut-down executors. Hold sessionPools' monitor across the closed check, the construction, and the insert. Opens are infrequent (typically once per table at app startup) so the monitor cost is negligible. Move close()'s closed flip inside the same monitor too. With every access now under the lock, downgrade 'closed' from AtomicBoolean to a plain boolean — the CAS provided no value over a plain read+write under the lock. --- .../bigtable/data/v2/internal/api/Client.java | 130 ++++++++++-------- 1 file changed, 70 insertions(+), 60 deletions(-) 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 3e7da51fd906..c99a396ca439 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 @@ -104,10 +104,10 @@ public class Client implements AutoCloseable { private final Resource configManager; private final Set> sessionPools = Collections.newSetFromMap(new WeakHashMap<>()); - // Set true at the start of close(); guards openTableAsync / openAuthorizedViewAsync / - // openMaterializedViewAsync so concurrent opens during shutdown don't create pools the close - // path won't see. - private final AtomicBoolean closed = new AtomicBoolean(false); + // 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 = @@ -235,12 +235,12 @@ public Client( @Override public void close() { - if (!closed.compareAndSet(false, true)) { - return; // idempotent - } - List> toClose; synchronized (sessionPools) { + if (closed) { + return; // idempotent + } + closed = true; toClose = new ArrayList<>(sessionPools); } @@ -292,67 +292,77 @@ public void close() { backgroundExecutor.close(); } - private void checkNotClosed() { - if (closed.get()) { - throw new IllegalStateException("Client is closed"); - } - } - + // 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) { - checkNotClosed(); - 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) { - checkNotClosed(); - 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) { - checkNotClosed(); - 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 { From dc544528818f98309d5843c47108679da3ce5cbc Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 15 Jun 2026 20:23:59 +0000 Subject: [PATCH 4/9] fix: ShimImpl uses shutdownAndAwait for userCallbackExecutor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ShimImpl registered its userCallbackExecutor with a bare shutdown() closer, while Client.create's path uses shutdownAndAwait (which gives in-flight listener.onClose tasks a 5-second drain window before shutdownNow). On ShimImpl close, queued callbacks were abandoned mid-flight — fine for quiescent shutdowns but a regression for fast-close patterns (test boundaries, dynamic config reloads) where in-flight callbacks have not yet drained. Promote Client.shutdownAndAwait from private-static to public-static so ShimImpl (different package) can reuse the same shutdown semantics, and update ShimImpl to call it. --- .../google/cloud/bigtable/data/v2/internal/api/Client.java | 6 ++++-- .../cloud/bigtable/data/v2/internal/compat/ShimImpl.java | 3 ++- 2 files changed, 6 insertions(+), 3 deletions(-) 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 c99a396ca439..ee2360774ef1 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 @@ -397,8 +397,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); } From 77d93db3a902cf7dfb332eb483751b1921c445ed Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 16 Jun 2026 00:40:31 +0000 Subject: [PATCH 5/9] fix: lift shim loader throws to failed futures via SessionPoolMap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ReadRowShim and MutateRowShim cache per-target handles in a Guava LoadingCache whose loader calls Client.openTableAsync (and friends). After the Client-close hardening, that loader throws IllegalStateException post-close — and getUnchecked wraps it in UncheckedExecutionException, so callers see an unchecked exception thrown out of readRow / mutateRow instead of a failed CompletableFuture as the surface contract promises. Introduce SessionPoolMap, a small wrapper over the existing Util.createSessionMap cache that owns the conversion: get() unwraps UncheckedExecutionException to surface the original cause, and apply() converts a loader throw into a failed CompletableFuture for the async call paths. Replace the inline LoadingCache fields in both shim ops files. Covered by a new focused unit test. --- .../v2/internal/compat/ops/MutateRowShim.java | 26 ++--- .../internal/compat/ops/ReadRowShimInner.java | 26 +++-- .../internal/compat/ops/SessionPoolMap.java | 77 +++++++++++++++ .../compat/ops/SessionPoolMapTest.java | 98 +++++++++++++++++++ 4 files changed, 197 insertions(+), 30 deletions(-) create mode 100644 java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMap.java create mode 100644 java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/SessionPoolMapTest.java 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..d7022ea4938c --- /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,77 @@ +/* + * 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}. If the lookup throws (e.g. {@link + * IllegalStateException} from a closed Client), the throw is 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; + } + return op.apply(v); + } + + /** 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/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..9b3156d14f50 --- /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,98 @@ +/* + * 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_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); + } +} From 97db090800e746758ad1531aab06ac5211197ad7 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 23 Jun 2026 17:56:38 +0000 Subject: [PATCH 6/9] test: cover pool close while a vRPC is queued in pendingRpcs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing awaitTerminated tests only close an empty pool, close after the vRPC has completed, or issue the vRPC against an already-closed pool — none of them exercise SessionPoolImpl.close()'s pendingRpcs cancellation loop with a non-empty toCancel list. Add closeCancelsInFlightVRpcAndDrains: park the session-open response in a DelayedClientInterceptor so the vRPC sits in pendingRpcs, then close the pool and verify the listener gets a single CANCELLED onClose with the "SessionPool closed:" description, awaitTerminated returns true via the OPENING session draining through sessions.close(req), and no phantom second onClose arrives once the delayed responses flush. --- .../internal/session/SessionPoolImplTest.java | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) 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 185e1b32305e..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; @@ -319,6 +320,89 @@ void awaitTerminatedReturnsTrueAfterSessionsDrain() 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 From 8f643adc74c33e1d7f541967ed1734c376eb2f19 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 23 Jun 2026 18:02:53 +0000 Subject: [PATCH 7/9] chore: record debug tag before throwing on unexpected drainedFuture failure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The catch branch in awaitTerminated is reachable only if drainedFuture is ever completed exceptionally or cancelled — neither happens today, but if it ever does we should leave a telemetry breadcrumb instead of just an uncaught IllegalStateException for the caller to puzzle over. Mirror the precondition-failure pattern used elsewhere in this file (record at ERROR, then throw). --- .../bigtable/data/v2/internal/session/SessionPoolImpl.java | 2 ++ 1 file changed, 2 insertions(+) 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 79b15214695f..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 @@ -331,6 +331,8 @@ public boolean awaitTerminated(Duration timeout) throws InterruptedException { } 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. From 3508eb32dbc5e7f16a9012a74d929aa69ce1130b Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 30 Jun 2026 16:03:32 +0000 Subject: [PATCH 8/9] chore: apply fmt-maven-plugin --- .../bigtable/data/v2/internal/api/Client.java | 1 - .../v2/internal/middleware/RetryingVRpc.java | 23 +++++++++++-------- .../v2/internal/session/BigtableTimer.java | 3 +-- .../v2/internal/session/NettyWheelTimer.java | 3 +-- .../compat/ops/SessionPoolMapTest.java | 3 +-- 5 files changed, 16 insertions(+), 17 deletions(-) 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 ee2360774ef1..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 @@ -53,7 +53,6 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Level; import java.util.logging.Logger; 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 277637ac41c8..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 @@ -331,16 +331,19 @@ private void onTimerFired() { // 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")))); - }); + 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() { 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 6731355a63e6..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,8 +47,7 @@ public interface BigtableTimer { /** * Releases the tick thread and discards any pending timeouts. Idempotent. After {@code stop()}, - * subsequent calls to {@link #newTimeout} or {@link #onStop} 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 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 442b4efd5f43..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 @@ -64,8 +64,7 @@ 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)); + return new TimeoutHandle(delegate.newTimeout(ignored -> dispatcher.execute(task), delay, unit)); } @Override 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 index 9b3156d14f50..279437e36973 100644 --- 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 @@ -77,8 +77,7 @@ void apply_happyPathInvokesOp() throws Exception { CountingHandle handle = new CountingHandle(); SessionPoolMap map = new SessionPoolMap<>(key -> handle); - CompletableFuture result = - map.apply("k", v -> CompletableFuture.completedFuture("ok")); + CompletableFuture result = map.apply("k", v -> CompletableFuture.completedFuture("ok")); assertThat(result.get(1, TimeUnit.SECONDS)).isEqualTo("ok"); } From 2a3aefbc5caaa9414783badec62ffcca1cbbd3a0 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Tue, 30 Jun 2026 17:43:03 +0000 Subject: [PATCH 9/9] fix: route SessionPoolMap.apply op throws to failed futures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `apply` already converted lookup throws to failed futures, but a sync throw from `op.apply(v)` itself (NPE on malformed input, REE during shutdown, any RuntimeException from the wrapped call chain) propagated to the caller — contradicting the docstring's contract that "callers consistently observe failures through the future surface." Wrap the op invocation in try/catch with the same pattern as the lookup arm. --- .../internal/compat/ops/SessionPoolMap.java | 15 +++++++++---- .../compat/ops/SessionPoolMapTest.java | 22 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) 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 index d7022ea4938c..a1f8257922b3 100644 --- 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 @@ -54,9 +54,10 @@ public V get(K key) { } /** - * Looks up the cached value and applies {@code op}. If the lookup throws (e.g. {@link - * IllegalStateException} from a closed Client), the throw is converted to a failed future so - * callers consistently observe failures through the future surface. + * 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; @@ -67,7 +68,13 @@ public CompletableFuture apply(K key, Function> o f.completeExceptionally(e); return f; } - return op.apply(v); + 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. */ 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 index 279437e36973..f6e22c7ed392 100644 --- 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 @@ -72,6 +72,28 @@ void apply_convertsLoaderThrowToFailedFuture() throws Exception { 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();