From 3a1200c917614c49e9bbc8ffae8e4d51edffc159 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 24 Jul 2026 16:23:21 +0000 Subject: [PATCH] fix(bigtable): fix session creation leaks --- .../data/v2/internal/session/SessionImpl.java | 63 ++++++ .../data/v2/internal/session/SessionList.java | 18 ++ .../v2/internal/session/SessionPoolImpl.java | 80 ++++++- .../v2/internal/session/SessionImplTest.java | 105 +++++++++ .../internal/session/SessionPoolImplTest.java | 207 +++++++++++++++++- .../internal/session/fake/SessionHandler.java | 24 ++ .../test/proto/google/bigtable/v2/fake.proto | 12 + 7 files changed, 494 insertions(+), 15 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java index b88fa6e4c84f..18ca95a67631 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java @@ -78,6 +78,21 @@ public class SessionImpl implements Session, VRpcSessionApi { .setDescription("missed heartbeat") .build(); + @VisibleForTesting + // Upper bound on how long a session may remain STARTING before we give up on the open handshake + // and force-close it. Without this, a stream that connects but never delivers the OpenSession + // response (or a GoAway) leaves the session wedged in STARTING forever, holding the + // session-creation-budget slot it reserved and never firing the terminal callback that would + // release it. Force-closing routes STARTING -> WAIT_SERVER_CLOSE -> onSessionClose, which frees + // the budget. + static final Duration OPEN_SESSION_TIMEOUT = Duration.ofSeconds(90); + + private static final CloseSessionRequest OPEN_TIMEOUT_CLOSE_REQUEST = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) + .setDescription("session open timeout") + .build(); + private final Clock clock; private final BigtableTimer timer; // Serializes all session state mutations. Stream callbacks and the heartbeat tick dispatch @@ -130,6 +145,10 @@ public class SessionImpl implements Session, VRpcSessionApi { // transitions so the wheel doesn't carry a no-op entry until the next fire. @Nullable private BigtableTimer.Timeout heartbeatTimeout; + // Handle for the open-handshake deadline (armed while STARTING). Cancelled the moment the + // session leaves STARTING so a session that opens normally never trips it. + @Nullable private BigtableTimer.Timeout openTimeout; + // Set by the global SyncContext handler when an uncaught exception triggers an abort. Read on // re-entry to break out instead of looping. Only accessed inside sessionSyncContext. private boolean isAborting = false; @@ -310,6 +329,9 @@ public void start(OpenSessionRequest req, Metadata headers, Listener sessionList tracer.onStart(); updateState(SessionState.STARTING); + // Bound the open handshake: if we're still STARTING when this fires, force-close so the + // reserved session-creation-budget slot is released instead of leaking forever. + scheduleOpenTimeoutCheck(); openParams = OpenParams.create(headers, req); SessionRequest wrappedReq = SessionRequest.newBuilder().setOpenSession(req).build(); @@ -457,6 +479,41 @@ public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable }); } + private void scheduleOpenTimeoutCheck() { + openTimeout = + timer.newTimeout( + this::checkOpenTimeout, + sessionSyncContext, + OPEN_SESSION_TIMEOUT.toMillis(), + TimeUnit.MILLISECONDS); + } + + private void cancelOpenTimeout() { + if (openTimeout != null) { + openTimeout.cancel(); + openTimeout = null; + } + } + + // Runs on sessionSyncContext (dispatched from the wheel-timer tick body). If the session is + // still STARTING, the open handshake never completed; force-close so the reserved + // session-creation-budget slot is released via the STARTING -> WAIT_SERVER_CLOSE -> + // onSessionClose + // path. Any transition out of STARTING cancels this timer, so reaching here in another state is a + // benign race (tick already dispatched) and is ignored. + private void checkOpenTimeout() { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + if (state != SessionState.STARTING) { + return; + } + logger.warning( + String.format( + "Session %s did not complete the open handshake within %s, forcing session close", + info.getLogName(), OPEN_SESSION_TIMEOUT)); + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_open_timeout"); + forceClose(OPEN_TIMEOUT_CLOSE_REQUEST); + } + private void scheduleHeartbeatCheck() { heartbeatTimeout = timer.newTimeout( @@ -849,6 +906,12 @@ private PeerInfo safeGetPeerInfo() { private void updateState(SessionState newState) { this.state = newState; this.lastStateChangedAt = clock.instant(); + // The open deadline only applies while STARTING. Any other transition (READY on success, or a + // terminal state) means the handshake resolved, so drop the pending tick. NEW -> STARTING keeps + // it, since start() arms the timer right after this call. + if (newState != SessionState.STARTING) { + cancelOpenTimeout(); + } // Once we're past READY, no further heartbeat checks are useful: checkHeartbeat short-circuits // on state.phase >= WAIT_SERVER_CLOSE. Cancel any pending tick to keep the wheel clean during // session churn. diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java index 3f8d996916b2..8021e2f78199 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionList.java @@ -77,6 +77,24 @@ SessionHandle newHandle(Session session) { return h; } + /** + * Unwinds a handle produced by {@link #newHandle} whose session never started, i.e. a synchronous + * failure in {@code SessionPoolImpl.createSession} between {@code newHandle} and {@code + * session.start}. Such a session stays in {@link SessionState#NEW}, so it never receives a + * terminal {@link SessionHandle#onSessionClosed} callback (which additionally rejects NEW + * sessions). This reverses exactly the bookkeeping {@code newHandle} performed so the handle + * doesn't strand {@link #allSessions} (blocking pool drain) or skew the pool stats. + */ + void removeUnstartedHandle(SessionHandle handle) { + if (allSessions.remove(handle)) { + poolStats.startingCount--; + if (handle.inExpectedCount) { + poolStats.expectedCapacity--; + handle.inExpectedCount = false; + } + } + } + /** Get {@link PoolStats} */ PoolStats getStats() { return poolStats; 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 7046804e52a5..d2c6892d800a 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 @@ -175,6 +175,18 @@ private enum PoolState { @GuardedBy("poolLock") private final SessionCreationBudget budget; + // Handles that reserved a session-creation budget slot in createSession() and have not yet + // released it. A slot is released exactly once: as a success when the session reaches READY + // (onSessionReady), or as a failure when the session terminates without ever becoming READY + // (onSessionClose). Tracking the reservation on the handle -- instead of inferring it from the + // close-time prevState -- is what makes the release exactly-once: a session that goes + // STARTING -> WAIT_SERVER_CLOSE (e.g. a server GO_AWAY before the open handshake completes) + // closes with prevState == WAIT_SERVER_CLOSE, which the abnormal-close branch skips, so the + // prevState == STARTING check alone would leak the slot forever. SessionHandle uses identity + // equality, so a plain HashSet keys on the handle instance. + @GuardedBy("poolLock") + private final Set sessionsHoldingBudget = new HashSet<>(); + private final ClientConfigurationManager configManager; private final ClientConfigurationManager.ListenerHandle configListenerHandle; @@ -455,35 +467,65 @@ private void createSession(OpenParams openParams) { // Explicit create session streams in a detached context // We don't want to propagate the rpc deadline nor the trace context Context prevContext = Context.ROOT.attach(); + // Tracks the handle once registered so the catch below can tell whether the reservation was + // ever bound to a handle (and therefore what cleanup is required). + SessionHandle handle = null; try { try (Scope ignored = io.opentelemetry.context.Context.root().makeCurrent()) { + // Build the metadata before registering the handle so the only step between newHandle and + // session.start (whose own failures are handled by the session's abort path once the + // listener is published) is the non-throwing sessionsHoldingBudget.add. + Metadata localMd = new Metadata(); + localMd.merge(openParams.metadata()); + SessionStream stream = factory.createNew(); Session session = new SessionImpl(metrics, info, sessionNum++, stream, timer); - SessionHandle handle = sessions.newHandle(session); + handle = sessions.newHandle(session); + // Bind the budget reservation made by tryReserveSession() above to this handle so it is + // released exactly once when the session becomes READY or terminates. + sessionsHoldingBudget.add(handle); - Metadata localMd = new Metadata(); - localMd.merge(openParams.metadata()); + final SessionHandle startedHandle = handle; session.start( openParams.request(), localMd, new Listener() { @Override public void onReady(OpenSessionResponse msg) { - SessionPoolImpl.this.onSessionReady(handle, msg); + SessionPoolImpl.this.onSessionReady(startedHandle, msg); } @Override public void onGoAway(GoAwayResponse msg) { - SessionPoolImpl.this.onSessionGoAway(handle, msg); + SessionPoolImpl.this.onSessionGoAway(startedHandle, msg); } @Override public void onClose(SessionState prevState, Status status, Metadata trailers) { - SessionPoolImpl.this.onSessionClose(handle, prevState, status, trailers); + SessionPoolImpl.this.onSessionClose(startedHandle, prevState, status, trailers); } }); } + } catch (RuntimeException | Error e) { + // A synchronous failure here (e.g. factory.createNew, the SessionImpl constructor, or + // metadata merge) means no terminal session callback will ever run for this reservation, so + // the budget slot and any partially-registered handle would leak for the life of the pool. + // Release the slot as a failure and unwind the handle ourselves. Note session.start's own + // failures do NOT land here: it publishes the listener first and routes exceptions through + // the session's abort path, which fires onClose -> onSessionClose (the normal release path). + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_create_failed"); + logger.log(Level.WARNING, "Failed to create session, releasing reserved budget", e); + // If a handle was registered it owns the reservation via the set; otherwise the reservation + // was never bound to a handle and must be released directly. + if (handle == null || sessionsHoldingBudget.remove(handle)) { + budget.onSessionCreationFailure(); + } + if (handle != null) { + sessions.removeUnstartedHandle(handle); + } + // Let the pool recover instead of running permanently short a session. + maybeScheduleCreateSessionRetry(); } finally { Context.ROOT.detach(prevContext); } @@ -540,7 +582,11 @@ private void onSessionReady(SessionHandle handle, OpenSessionResponse ignored) { } handle.onSessionStarted(); - budget.onSessionCreationSuccess(); + // Release the reservation as a success. Guarded by the set so we release exactly once even if + // onSessionReady were ever delivered more than once. + if (sessionsHoldingBudget.remove(handle)) { + budget.onSessionCreationSuccess(); + } // handle pending rpcs tryDrainPendingRpcs(); @@ -613,6 +659,21 @@ private void onSessionClose( poolLock.lock(); try { + // Release the budget reservation FIRST, before any code below that can throw. This is the + // session's terminal callback, so there is no later backstop: if the reservation is not + // released here it leaks forever. In particular handle.onSessionClosed(prevState) below can + // throw (e.g. IllegalStateException on an unexpected NEW / double-CLOSED prevState), which + // would otherwise strand the slot -- the same leak this fix exists to prevent. + // + // Release as a failure if this session still holds a slot, i.e. it terminated without ever + // reaching READY; sessions that reached READY were already removed from the set in + // onSessionReady, so this never double-releases. Keying off the reservation instead of the + // close-time prevState is also what covers the STARTING -> WAIT_SERVER_CLOSE (GO_AWAY) path + // that the abnormal-close branch below skips. + if (sessionsHoldingBudget.remove(handle)) { + budget.onSessionCreationFailure(); + } + logger.fine( String.format("Removing closed session from pool %s", handle.getSession().getLogName())); @@ -642,9 +703,8 @@ private void onSessionClose( toBeClosed = popClosableRpcs(); } - if (prevState == SessionState.STARTING) { - budget.onSessionCreationFailure(); - } + // Budget release for STARTING-phase closes is handled above via sessionsHoldingBudget, + // which also covers the STARTING -> WAIT_SERVER_CLOSE (GO_AWAY) path this branch skips. // TODO: backoff creating a new session when consecutive failures > max? if (poolSizer.handleSessionClose(StatusProto.fromStatusAndTrailers(status, trailers))) { 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 0b0006135cff..326f971af88a 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java @@ -70,6 +70,9 @@ import java.io.IOException; import java.time.Duration; import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; @@ -596,6 +599,12 @@ void testHeartbeatNotScheduledWithoutVRpc() throws Exception { assertThat(sessionListener.popUntil(OpenSessionResponse.class)) .isInstanceOf(OpenSessionResponse.class); + // The open-handshake deadline is armed while STARTING and cancelled once READY (before onReady, + // so it's already settled here). Reset the counters to scope the assertion below to the + // heartbeat lifecycle only. + counting.scheduleCount.set(0); + counting.cancelCount.set(0); + // After session is READY with no vRPC, no Timeout should ever have been scheduled. Wait a // bit so that any background tick (none expected) would have shown up. Thread.sleep(50); @@ -640,6 +649,11 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception { assertThat(sessionListener.popUntil(OpenSessionResponse.class)) .isInstanceOf(OpenSessionResponse.class); + // The open-handshake deadline is armed while STARTING and cancelled once READY. Reset the + // counters so the assertions below track only the heartbeat lifecycle. + counting.scheduleCount.set(0); + counting.cancelCount.set(0); + assertThat(counting.scheduleCount.get()).isEqualTo(0); VRpc rpc = @@ -811,8 +825,99 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); } + // Regression test: a session that connects but never completes the open handshake must not stay + // wedged in STARTING forever. Without the open-handshake deadline the session would never fire a + // terminal callback, so the pool's reserved session-creation-budget slot would leak. Arming + // OPEN_SESSION_TIMEOUT force-closes such a session, routing STARTING -> WAIT_SERVER_CLOSE and + // driving the listener to a terminal Status (which releases the budget in the pool). + @Test + void sessionOpenTimeoutForcesClose() throws Exception { + // Capture the scheduled open-timeout tick so we can fire it deterministically instead of + // waiting the real OPEN_SESSION_TIMEOUT. + CapturingBigtableTimer capturing = new CapturingBigtableTimer(timer); + SessionImpl session = + new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), capturing); + + FakeSessionListener sessionListener = new FakeSessionListener(); + session.start( + OpenSessionRequest.newBuilder() + .setPayload( + OpenFakeSessionRequest.newBuilder().setHangBeforeOpen(true).build().toByteString()) + .build(), + new Metadata(), + sessionListener); + + // start() runs on the sessionSyncContext; wait until it has armed the open-timeout tick and the + // session is STARTING. The server never responds, so it stays STARTING. + Stopwatch sw = Stopwatch.createStarted(); + while ((capturing.tasks.isEmpty() || session.getState() != Session.SessionState.STARTING) + && sw.elapsed(TimeUnit.SECONDS) < 5) { + Thread.sleep(10); + } + assertThat(session.getState()).isEqualTo(Session.SessionState.STARTING); + assertWithMessage("open-timeout tick should have been armed while STARTING") + .that(capturing.tasks) + .isNotEmpty(); + + // Fire the open-timeout tick on its executor (the sessionSyncContext), simulating the deadline + // elapsing while still STARTING. + capturing.executors.get(0).execute(capturing.tasks.get(0)); + + // The session must force-close and drive the listener to a terminal Status. + assertWithMessage("terminal status should be delivered after open timeout") + .that(sessionListener.popUntil(Status.class)) + .isNotNull(); + sw.reset().start(); + while (session.getState() != Session.SessionState.WAIT_SERVER_CLOSE + && session.getState() != Session.SessionState.CLOSED + && sw.elapsed(TimeUnit.SECONDS) < 5) { + Thread.sleep(10); + } + assertThat(session.getState()) + .isAnyOf(Session.SessionState.WAIT_SERVER_CLOSE, Session.SessionState.CLOSED); + } + // endregion + // Captures newTimeout tasks/executors so a scheduled tick can be fired deterministically instead + // of waiting out the real delay. + private static final class CapturingBigtableTimer implements BigtableTimer { + private final BigtableTimer delegate; + final List tasks = Collections.synchronizedList(new ArrayList<>()); + final List executors = Collections.synchronizedList(new ArrayList<>()); + + CapturingBigtableTimer(BigtableTimer delegate) { + this.delegate = delegate; + } + + @Override + public Timeout newTimeout(Runnable task, Executor executor, long delay, TimeUnit unit) { + tasks.add(task); + executors.add(executor); + return new Timeout() { + @Override + public boolean cancel() { + return true; + } + + @Override + public boolean isCancelled() { + return false; + } + }; + } + + @Override + public Registration onStop(Runnable hook) { + return delegate.onStop(hook); + } + + @Override + public void stop() { + delegate.stop(); + } + } + // Wraps a real BigtableTimer and counts newTimeout / cancel calls. Used to assert that the // heartbeat tick is only armed while a vRPC is in flight. private static final class CountingBigtableTimer implements BigtableTimer { 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 9e7c5d542ba9..04292f8e3133 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 @@ -555,13 +555,16 @@ void tearDown() { @Test public void test() throws Exception { ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); - // Filter out watchdog (5 min, exact) and AFE prune (10 min, exact) ticks. The - // retry-create-session site computes its delay against the real wall clock and the fake - // budget clock, so it can land anywhere from sub-second to a couple of penalty intervals. - // Match anything that isn't one of the two fixed cadences. + // Filter out watchdog (5 min, exact), AFE prune (10 min, exact), and per-session open-timeout + // (OPEN_SESSION_TIMEOUT, exact) ticks. The retry-create-session site computes its delay + // against the real wall clock and the fake budget clock, so it can land anywhere from + // sub-second to a couple of penalty intervals. Match anything that isn't one of the fixed + // cadences. long watchdogMs = Duration.ofMinutes(5).toMillis(); long afePruneMs = SessionList.SESSION_LIST_PRUNE_INTERVAL.toMillis(); - LongPredicate isRetrySchedule = d -> d > 0 && d != watchdogMs && d != afePruneMs; + long openTimeoutMs = SessionImpl.OPEN_SESSION_TIMEOUT.toMillis(); + LongPredicate isRetrySchedule = + d -> d > 0 && d != watchdogMs && d != afePruneMs && d != openTimeoutMs; // start the pool sessionPool.start( @@ -616,6 +619,200 @@ public void test() throws Exception { } } + @Nested + class GoAwayBeforeReadyBudgetRelease { + + private FakeClock fakeClock; + private BigtableTimer mockTimer; + private FakeSessionService fakeService; + private Server server; + private ChannelPool channelPool; + private SessionPoolImpl sessionPool; + private SessionCreationBudget budget; + + private final Duration penalty = Duration.ofMinutes(1); + + @BeforeEach + void setUp() throws Exception { + fakeClock = new FakeClock(Instant.now()); + // Mock timer so the pool's watchdog / AFE-prune / create-session-retry ticks never fire on + // their own -- the budget stays quiescent after the session lifecycle completes, leaving the + // test thread as the only actor that touches it. + mockTimer = mock(BigtableTimer.class, Mockito.RETURNS_DEEP_STUBS); + + // Budget of 1 so exactly one session goes STARTING; any extra create attempts fail to reserve + // and just schedule a (no-op) retry on the mock timer. + budget = new SessionCreationBudget(1, penalty, fakeClock); + + fakeService = new FakeSessionService(executor); + server = FakeServiceBuilder.create(fakeService).intercept(new PeerInfoInterceptor()).start(); + + channelPool = + new SingleChannelPool( + Suppliers.ofInstance( + Grpc.newChannelBuilderForAddress( + "localhost", server.getPort(), InsecureChannelCredentials.create()) + .build())); + channelPool.start(); + + sessionPool = + new SessionPoolImpl<>( + metrics, + FeatureFlags.getDefaultInstance(), + CLIENT_INFO, + configManager, + channelPool, + CallOptions.DEFAULT, + FakeDescriptor.FAKE_SESSION, + "fake-pool", + mockTimer, + MoreExecutors.directExecutor(), + budget); + } + + @AfterEach + void tearDown() { + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("close session") + .build()); + channelPool.close(); + server.shutdownNow(); + } + + @Test + void goAwayWhileStartingReleasesBudget() throws Exception { + // Drive a single session that receives a GO_AWAY before the open handshake, so it transitions + // STARTING -> WAIT_SERVER_CLOSE and closes with prevState == WAIT_SERVER_CLOSE. The + // abnormal-close branch (prevState != WAIT_SERVER_CLOSE) is skipped, so before the fix the + // STARTING budget slot -- which was released only on the prevState == STARTING path -- leaked + // permanently. The fix keys the release off the reservation set, releasing it here as a + // creation failure. + sessionPool.start( + OpenFakeSessionRequest.newBuilder().setGoAwayBeforeOpen(true).build(), new Metadata()); + + // Poll for the leaked slot to become reclaimable. onSessionClose(WAIT_SERVER_CLOSE) records a + // creation failure (fix only); advancing past the penalty lets tryReserveSession drain it and + // succeed. SessionCreationBudget has no internal locking -- the pool guards it with poolLock + // -- + // so hold that same lock here so this doesn't race the session-close callback that mutates + // the + // budget. With the bug no slot is ever released and this stays false until the deadline. + ReentrantLock poolLock = extractPoolLock(sessionPool); + boolean recovered = false; + long deadlineMs = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadlineMs) { + poolLock.lock(); + try { + fakeClock.increment(penalty.plusMillis(1)); + if (budget.tryReserveSession()) { + recovered = true; + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + assertThat(recovered).isTrue(); + + // Sanity check that the client actually walked the GO_AWAY-before-open path: on receiving the + // GO_AWAY while STARTING it moves to WAIT_SERVER_CLOSE and sends a CloseSession. + boolean sentClose = + fakeService.getSessionRequests().stream() + .anyMatch(r -> r.getPayloadCase() == SessionRequest.PayloadCase.CLOSE_SESSION); + assertThat(sentClose).isTrue(); + } + } + + @Nested + class CreateSessionFailureBudgetRelease { + + private FakeClock fakeClock; + private BigtableTimer mockTimer; + private ChannelPool throwingChannelPool; + private SessionPoolImpl sessionPool; + private SessionCreationBudget budget; + + private final Duration penalty = Duration.ofMinutes(1); + + @BeforeEach + void setUp() { + fakeClock = new FakeClock(Instant.now()); + // Mock timer so the pool's watchdog / AFE-prune / create-session-retry ticks never fire on + // their own, leaving the test thread as the only actor touching the budget. + mockTimer = mock(BigtableTimer.class, Mockito.RETURNS_DEEP_STUBS); + + // Budget of 1 so a single create attempt reserves the only slot. + budget = new SessionCreationBudget(1, penalty, fakeClock); + + // newStream throws synchronously, so factory.createNew() fails inside createSession before a + // SessionImpl is constructed and before any listener is wired -- the failure mode Claim 2 + // addresses. With no session, no terminal callback ever runs to release the reservation. + throwingChannelPool = mock(ChannelPool.class); + when(throwingChannelPool.newStream(any(), any())) + .thenThrow(new RuntimeException("boom: newStream failed")); + + sessionPool = + new SessionPoolImpl<>( + metrics, + FeatureFlags.getDefaultInstance(), + CLIENT_INFO, + configManager, + throwingChannelPool, + CallOptions.DEFAULT, + FakeDescriptor.FAKE_SESSION, + "fake-pool", + mockTimer, + MoreExecutors.directExecutor(), + budget); + } + + @AfterEach + void tearDown() { + sessionPool.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionRequest.CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("close session") + .build()); + } + + @Test + void createSessionSynchronousFailureReleasesBudget() throws Exception { + // Starting the pool drives createSession, whose factory.createNew() throws synchronously. No + // session is constructed and no listener is wired, so no terminal callback will ever run. + // Before the fix the reserved budget slot leaked permanently; the fix releases it as a + // creation failure in createSession's catch block (and unwinds any partial handle). + sessionPool.start(OpenFakeSessionRequest.getDefaultInstance(), new Metadata()); + + // Poll for the slot to become reclaimable: draining the failure penalty must let a fresh + // reservation succeed. SessionCreationBudget has no internal locking -- the pool guards it + // with + // poolLock -- so hold that same lock here so this doesn't race any callback. With the bug no + // slot is ever released and this stays false until the deadline. + ReentrantLock poolLock = extractPoolLock(sessionPool); + boolean recovered = false; + long deadlineMs = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadlineMs) { + poolLock.lock(); + try { + fakeClock.increment(penalty.plusMillis(1)); + if (budget.tryReserveSession()) { + recovered = true; + break; + } + } finally { + poolLock.unlock(); + } + Thread.sleep(20); + } + + assertThat(recovered).isTrue(); + } + } + @Test public void refreshConfigTest() throws Exception { OpenSessionRequest refreshRequest = diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/SessionHandler.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/SessionHandler.java index f1e4b4055931..138fd153bf27 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/SessionHandler.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/fake/SessionHandler.java @@ -141,6 +141,30 @@ public State onMessage(SessionRequest sReq) { throw new RuntimeException(e); } + if (oReq.getHangBeforeOpen()) { + // Never send any handshake response, leaving the client wedged in STARTING so the + // open-handshake deadline (OPEN_SESSION_TIMEOUT) can fire. DoneState ignores all further + // messages; the client's eventual forceClose cancels the stream. + return DoneState.INSTANCE; + } + + if (oReq.getGoAwayBeforeOpen()) { + // Send a GO_AWAY before the SessionParameters / OpenSession handshake, so the client + // receives it while the session is still STARTING. The client will transition to + // WAIT_SERVER_CLOSE, send a CloseSession, and half-close; RunningState terminates the + // stream successfully on either of those. goAwayDelay is left large so RunningState's own + // lifecycle GO_AWAY timer never fires during the test. + handler.writeResponse( + SessionResponse.newBuilder() + .setGoAway( + GoAwayResponse.newBuilder() + .setReason("session_expiry") + .setDescription("go away before open handshake") + .build()) + .build()); + return new RunningState(handler, Duration.ofMinutes(1), oReq); + } + if (oReq.hasStreamError()) { Status status = Status.fromCodeValue(oReq.getStreamError().getStatus().getCode()) diff --git a/java-bigtable/google-cloud-bigtable/src/test/proto/google/bigtable/v2/fake.proto b/java-bigtable/google-cloud-bigtable/src/test/proto/google/bigtable/v2/fake.proto index bf2c5eec5a05..6c2dd7e0331a 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/proto/google/bigtable/v2/fake.proto +++ b/java-bigtable/google-cloud-bigtable/src/test/proto/google/bigtable/v2/fake.proto @@ -37,6 +37,18 @@ message OpenFakeSessionRequest { StreamError stream_error = 3; SessionParametersResponse session_params = 4; + // If set, the server sends a GO_AWAY immediately without ever sending the + // SessionParameters / OpenSession handshake responses, i.e. the session receives a + // GO_AWAY while it is still STARTING. Used to exercise the STARTING -> WAIT_SERVER_CLOSE + // session-creation-budget release path. + bool go_away_before_open = 9; + + // If set, the server accepts the stream but never sends any handshake response (no + // SessionParameters, no OpenSession, no GoAway), leaving the client wedged in STARTING. Used to + // exercise the open-handshake deadline (OPEN_SESSION_TIMEOUT) that force-closes a session that + // never finishes opening, so its reserved session-creation-budget slot is released. + bool hang_before_open = 10; + // session steady phase message Action { google.protobuf.Duration delay = 1;