Skip to content

Commit 5312d04

Browse files
igorbernstein2igorbernstein2
committed
chore: abort session on uncaught exception in sessionSyncContext
Split terminal close into notifyTerminalClose (per-target try/catch fan-out) and abortFromUncaughtException (global handler). Uncaught syncContext exceptions always set closeReason to ERROR — the prior reason is folded into the description so tracer/metrics correctly attribute aborts. notifyTerminalClose synthesizes a fallback closeReason if missing: every caller sets it today (forceClose, startGracefulClose, dispatchStreamClosed, abortFromUncaughtException), but a future writer who forgets would NPE inside the fan-out — and the throw escapes to the syncContext uncaught handler, which early-returns on the already-CLOSED state and silently skips the remaining cleanup. The synthesizer mirrors startGracefulClose: log a warning with an IllegalStateException for stack-trace observability, then build a fallback CloseSessionRequest so the rest of the fan-out runs. Adds three regression tests (listener.onReady throws, onClose throws, both throw).
1 parent d96b878 commit 5312d04

2 files changed

Lines changed: 292 additions & 18 deletions

File tree

  • java-bigtable/google-cloud-bigtable/src

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImpl.java‎

Lines changed: 149 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -126,11 +126,13 @@ public class SessionImpl implements Session, VRpcSessionApi {
126126

127127
private Instant nextHeartbeat;
128128

129-
// Handle for the in-flight heartbeat tick (one outstanding at a time). Set under lock when the
130-
// session enters READY (handleOpenSessionResponse) and again from checkHeartbeat to chain the
131-
// next tick. Cancelled under lock from updateState when the session transitions past READY.
132-
@Nullable
133-
private BigtableTimer.Timeout heartbeatTimeout;
129+
// Handle for the in-flight heartbeat tick (one outstanding at a time). Cancelled on terminal
130+
// transitions so the wheel doesn't carry a no-op entry until the next fire.
131+
@Nullable private BigtableTimer.Timeout heartbeatTimeout;
132+
133+
// Set by the global SyncContext handler when an uncaught exception triggers an abort. Read on
134+
// re-entry to break out instead of looping. Only accessed inside sessionSyncContext.
135+
private boolean isAborting = false;
134136

135137
public SessionImpl(
136138
Metrics metrics,
@@ -156,14 +158,77 @@ public SessionImpl(
156158
this.debugTagTracer = metrics.getDebugTagTracer();
157159
this.nextHeartbeat = clock.instant().plus(FUTURE_TIME);
158160
this.openParamsUpdated = false;
161+
// On uncaught exception, drive the session through a clean terminal-close path so the pool
162+
// and the in-flight vRpc are always notified. notifyTerminalClose has local guards, and
163+
// isAborting prevents recursion if the abort path itself throws.
159164
this.sessionSyncContext =
160-
new SynchronizationContext(
161-
(thread, e) ->
162-
logger.log(
163-
Level.WARNING,
164-
String.format(
165-
"Uncaught exception in session SyncContext for %s", info.getLogName()),
166-
e));
165+
new SynchronizationContext((thread, e) -> abortFromUncaughtException(e));
166+
}
167+
168+
private void abortFromUncaughtException(Throwable e) {
169+
if (isAborting) {
170+
logger.log(
171+
Level.WARNING,
172+
String.format(
173+
"Session error: %s Secondary uncaught exception during abort, ignoring",
174+
info.getLogName()),
175+
e);
176+
return;
177+
}
178+
isAborting = true;
179+
180+
logger.log(
181+
Level.SEVERE,
182+
String.format(
183+
"Session error: %s Uncaught exception in session SyncContext in state %s, PeerInfo:"
184+
+ " %s — aborting session",
185+
info.getLogName(), state, formatPeerInfo(safeGetPeerInfo())),
186+
e);
187+
188+
if (state == SessionState.CLOSED) {
189+
return;
190+
}
191+
192+
// Always overwrite closeReason: the abort is what actually happened, not whatever clean close
193+
// we may have been attempting. Fold the prior reason into the description for forensics so
194+
// downstream metrics (which bucket by reason) attribute this to ERROR rather than the
195+
// interrupted close.
196+
String prevDesc =
197+
(closeReason != null)
198+
? " (was closing for: "
199+
+ closeReason.getReason()
200+
+ " — "
201+
+ closeReason.getDescription()
202+
+ ")"
203+
: "";
204+
closeReason =
205+
CloseSessionRequest.newBuilder()
206+
.setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
207+
.setDescription("Uncaught exception in session SyncContext: " + e + prevDesc)
208+
.build();
209+
210+
VRpcImpl<?, ?, ?> localRpc = currentRpc;
211+
currentRpc = null;
212+
SessionState prevState = state;
213+
updateState(SessionState.CLOSED);
214+
215+
// Defensively tell the transport we're done. Safe on un-started streams via the try/catch.
216+
try {
217+
stream.forceClose("Session aborted due to uncaught exception", e);
218+
} catch (Throwable t) {
219+
logger.log(
220+
Level.WARNING,
221+
String.format(
222+
"Session error: %s Exception while force-closing stream during abort",
223+
info.getLogName()),
224+
t);
225+
}
226+
227+
notifyTerminalClose(
228+
Status.INTERNAL.withDescription("Session aborted").withCause(e),
229+
new Metadata(),
230+
localRpc,
231+
prevState);
167232
}
168233

169234
@Override
@@ -681,13 +746,45 @@ private void dispatchStreamClosed(Status status, Metadata trailers) {
681746
}
682747

683748
VRpcImpl<?, ?, ?> localVRpc = currentRpc;
684-
PeerInfo localPeerInfo = stream.getPeerInfo();
685749
currentRpc = null;
686750
updateState(SessionState.CLOSED);
687751

688-
if (localVRpc != null) {
752+
notifyTerminalClose(status, trailers, localVRpc, prevState);
753+
}
754+
755+
/**
756+
* Fan out terminal notifications to the in-flight vRpc, tracer, and session listener with local
757+
* guards so a throw in one notification does not suppress the others.
758+
*
759+
* <p>Caller contract: must have already transitioned to {@link SessionState#CLOSED} and captured
760+
* and cleared {@code currentRpc}. Callers should also set {@code closeReason}; if missing we
761+
* synthesize a fallback here rather than throw, since throwing from this fan-out aborts the
762+
* remaining notifications and (because the state is already CLOSED) defeats the
763+
* sessionSyncContext uncaught handler's cleanup.
764+
*/
765+
private void notifyTerminalClose(
766+
Status status,
767+
Metadata trailers,
768+
@Nullable VRpcImpl<?, ?, ?> localRpc,
769+
SessionState prevState) {
770+
// Should never happen — matches the synthesizer in startGracefulClose.
771+
if (closeReason == null) {
772+
debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_close_no_reason");
773+
logger.log(
774+
Level.WARNING,
775+
String.format(
776+
"%s notifyTerminalClose reached without a closeReason; status=%s",
777+
info.getLogName(), status),
778+
new IllegalStateException("notifyTerminalClose without closeReason"));
779+
closeReason =
780+
CloseSessionRequest.newBuilder()
781+
.setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR)
782+
.setDescription("notifyTerminalClose reached without closeReason; status=" + status)
783+
.build();
784+
}
785+
if (localRpc != null) {
689786
try {
690-
localVRpc.handleSessionClose(VRpcResult.createRemoteTransportError(status, trailers));
787+
localRpc.handleSessionClose(VRpcResult.createRemoteTransportError(status, trailers));
691788
} catch (Throwable t) {
692789
logger.log(
693790
Level.WARNING,
@@ -697,10 +794,44 @@ private void dispatchStreamClosed(Status status, Metadata trailers) {
697794
info.getLogName(), status),
698795
t);
699796
}
700-
tracer.onVRpcClose(Status.UNAVAILABLE.getCode());
797+
try {
798+
tracer.onVRpcClose(Status.UNAVAILABLE.getCode());
799+
} catch (Throwable t) {
800+
logger.log(
801+
Level.WARNING,
802+
String.format(
803+
"Session error: %s Unhandled exception in tracer.onVRpcClose", info.getLogName()),
804+
t);
805+
}
806+
}
807+
try {
808+
tracer.onClose(safeGetPeerInfo(), closeReason.getReason(), status);
809+
} catch (Throwable t) {
810+
logger.log(
811+
Level.WARNING,
812+
String.format("Session error: %s Unhandled exception in tracer.onClose", info.getLogName()),
813+
t);
814+
}
815+
if (sessionListener != null) {
816+
try {
817+
sessionListener.onClose(prevState, status, trailers);
818+
} catch (Throwable t) {
819+
logger.log(
820+
Level.WARNING,
821+
String.format(
822+
"Session error: %s Unhandled exception in sessionListener.onClose",
823+
info.getLogName()),
824+
t);
825+
}
826+
}
827+
}
828+
829+
private PeerInfo safeGetPeerInfo() {
830+
try {
831+
return stream.getPeerInfo();
832+
} catch (Throwable t) {
833+
return SessionStream.DISCONNECTED_PEER_INFO;
701834
}
702-
tracer.onClose(localPeerInfo, closeReason.getReason(), status);
703-
sessionListener.onClose(prevState, status, trailers);
704835
}
705836

706837
private void updateState(SessionState newState) {

‎java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/session/SessionImplTest.java‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -672,6 +672,149 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception {
672672
assertThat(sessionListener.popUntil(Status.class)).isOk();
673673
}
674674

675+
// region uncaught-exception abort behaviors
676+
677+
@Test
678+
void abortFiresWhenListenerOnReadyThrows() throws Exception {
679+
SessionImpl session =
680+
new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer);
681+
682+
java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1);
683+
java.util.concurrent.atomic.AtomicReference<Status> capturedStatus =
684+
new java.util.concurrent.atomic.AtomicReference<>();
685+
686+
Session.Listener throwingListener =
687+
new Session.Listener() {
688+
@Override
689+
public void onReady(OpenSessionResponse msg) {
690+
throw new RuntimeException("simulated onReady failure");
691+
}
692+
693+
@Override
694+
public void onGoAway(GoAwayResponse msg) {}
695+
696+
@Override
697+
public void onClose(Session.SessionState prevState, Status status, Metadata trailers) {
698+
capturedStatus.set(status);
699+
onCloseLatch.countDown();
700+
}
701+
};
702+
703+
session.start(
704+
OpenSessionRequest.newBuilder()
705+
.setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString())
706+
.build(),
707+
new Metadata(),
708+
throwingListener);
709+
710+
// The abort path must drive the session to CLOSED and notify the listener via onClose, even
711+
// though the original onReady threw.
712+
assertWithMessage("listener.onClose must be invoked after onReady throws")
713+
.that(onCloseLatch.await(5, TimeUnit.SECONDS))
714+
.isTrue();
715+
assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED);
716+
assertThat(capturedStatus.get().getCode()).isEqualTo(Status.Code.INTERNAL);
717+
}
718+
719+
@Test
720+
void abortDoesNotHangWhenListenerOnCloseThrows() throws Exception {
721+
SessionImpl session =
722+
new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer);
723+
724+
java.util.concurrent.CountDownLatch onReadyLatch = new java.util.concurrent.CountDownLatch(1);
725+
java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1);
726+
727+
Session.Listener throwingListener =
728+
new Session.Listener() {
729+
@Override
730+
public void onReady(OpenSessionResponse msg) {
731+
onReadyLatch.countDown();
732+
}
733+
734+
@Override
735+
public void onGoAway(GoAwayResponse msg) {}
736+
737+
@Override
738+
public void onClose(Session.SessionState prevState, Status status, Metadata trailers) {
739+
onCloseLatch.countDown();
740+
throw new RuntimeException("simulated onClose failure");
741+
}
742+
};
743+
744+
session.start(
745+
OpenSessionRequest.newBuilder()
746+
.setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString())
747+
.build(),
748+
new Metadata(),
749+
throwingListener);
750+
751+
assertThat(onReadyLatch.await(5, TimeUnit.SECONDS)).isTrue();
752+
753+
// Close normally. The listener's onClose throws — the local guard inside notifyTerminalClose
754+
// must swallow it so the SyncContext drain doesn't recurse infinitely or hang.
755+
session.close(
756+
CloseSessionRequest.newBuilder()
757+
.setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER)
758+
.setDescription("test")
759+
.build());
760+
761+
assertWithMessage("listener.onClose should be invoked exactly once during normal close")
762+
.that(onCloseLatch.await(5, TimeUnit.SECONDS))
763+
.isTrue();
764+
765+
// The session should reach CLOSED state cleanly within the test timeout.
766+
Stopwatch sw = Stopwatch.createStarted();
767+
while (session.getState() != Session.SessionState.CLOSED && sw.elapsed().getSeconds() < 5) {
768+
Thread.sleep(10);
769+
}
770+
assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED);
771+
}
772+
773+
@Test
774+
void abortDoesNotInfiniteLoopWhenRecoveryListenerAlsoThrows() throws Exception {
775+
SessionImpl session =
776+
new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer);
777+
778+
java.util.concurrent.CountDownLatch onCloseInvoked =
779+
new java.util.concurrent.CountDownLatch(1);
780+
781+
Session.Listener doublyThrowingListener =
782+
new Session.Listener() {
783+
@Override
784+
public void onReady(OpenSessionResponse msg) {
785+
throw new RuntimeException("simulated onReady failure");
786+
}
787+
788+
@Override
789+
public void onGoAway(GoAwayResponse msg) {}
790+
791+
@Override
792+
public void onClose(Session.SessionState prevState, Status status, Metadata trailers) {
793+
onCloseInvoked.countDown();
794+
throw new RuntimeException("simulated onClose failure during abort");
795+
}
796+
};
797+
798+
session.start(
799+
OpenSessionRequest.newBuilder()
800+
.setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString())
801+
.build(),
802+
new Metadata(),
803+
doublyThrowingListener);
804+
805+
// onReady throws → abort fires → abort calls onClose, which also throws → Guard 4 swallows
806+
// and isAborting prevents the handler from re-driving abort. The session must reach CLOSED
807+
// without hanging (the @Timeout(30) on the class is the safety net for infinite loops).
808+
assertThat(onCloseInvoked.await(5, TimeUnit.SECONDS)).isTrue();
809+
Stopwatch sw = Stopwatch.createStarted();
810+
while (session.getState() != Session.SessionState.CLOSED && sw.elapsed().getSeconds() < 5) {
811+
Thread.sleep(10);
812+
}
813+
assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED);
814+
}
815+
816+
// endregion
817+
675818
// Wraps a real BigtableTimer and counts newTimeout / cancel calls. Used to assert that the
676819
// heartbeat tick is only armed while a vRPC is in flight.
677820
private static final class CountingBigtableTimer implements BigtableTimer {

0 commit comments

Comments
 (0)