From 7ed5f3acb319feb3796945689e67ee315257a8c2 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Sun, 14 Jun 2026 03:00:11 +0000 Subject: [PATCH 1/6] chore: add session SynchronizationContext alongside the existing lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SessionImpl gains sessionSyncContext that serializes stream callbacks. onMessage/onClose dispatch onto it, and the per-session heartbeat tick trampolines through it. synchronized(lock) blocks remain inside the handlers — the two coexist for now. Affinity asserts added at boundary methods and every handle*. --- .../data/v2/internal/session/SessionImpl.java | 51 +++++++++++++------ 1 file changed, 35 insertions(+), 16 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 fdaa3f2ed0ad..f8220550083a 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 @@ -44,6 +44,7 @@ import com.google.protobuf.util.Durations; import io.grpc.Metadata; import io.grpc.Status; +import io.grpc.SynchronizationContext; import java.time.Clock; import java.time.Duration; import java.time.Instant; @@ -87,6 +88,9 @@ public class SessionImpl implements Session, VRpcSessionApi { private final Clock clock; private final BigtableTimer timer; + // Serializes all session-stream callbacks. Coexists with `lock` for now; `lock` will be removed + // in a follow-up commit once every handler is on the syncContext. + private final SynchronizationContext sessionSyncContext; private final SessionTracer tracer; private final DebugTagTracer debugTagTracer; @@ -177,6 +181,14 @@ public SessionImpl( this.debugTagTracer = metrics.getDebugTagTracer(); this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); this.openParamsUpdated = false; + this.sessionSyncContext = + new SynchronizationContext( + (thread, e) -> + logger.log( + Level.WARNING, + String.format( + "Uncaught exception in session SyncContext for %s", info.getLogName()), + e)); } @Override @@ -271,12 +283,12 @@ public void onBeforeSessionStart(PeerInfo peerInfo) {} @Override public void onMessage(SessionResponse message) { - dispatchResponseMessage(message); + sessionSyncContext.execute(() -> dispatchResponseMessage(message)); } @Override public void onClose(Status status, Metadata trailers) { - dispatchStreamClosed(status, trailers); + sessionSyncContext.execute(() -> dispatchStreamClosed(status, trailers)); } }, headers); @@ -411,7 +423,9 @@ public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable private void scheduleHeartbeatCheck() { heartbeatTimeout = timer.newTimeout( - this::checkHeartbeat, HEARTBEAT_CHECK_INTERVAL.toMillis(), TimeUnit.MILLISECONDS); + () -> sessionSyncContext.execute(this::checkHeartbeat), + HEARTBEAT_CHECK_INTERVAL.toMillis(), + TimeUnit.MILLISECONDS); } @GuardedBy("lock") @@ -422,31 +436,28 @@ private void cancelHeartbeatTimeout() { } } - // Runs on the wheel-timer tick thread. Takes the per-session lock to read state/nextHeartbeat - // and force-close on miss, then chains the next tick by re-scheduling. If the session is past - // WAIT_SERVER_CLOSE we drop the chain — no further checks are useful. + // Runs on sessionSyncContext (dispatched from the wheel-timer tick body). Checks if the + // heartbeat deadline has passed and force-closes on miss; otherwise re-schedules. private void checkHeartbeat() { - CloseSessionRequest missed = null; + sessionSyncContext.throwIfNotInThisSynchronizationContext(); synchronized (lock) { if (state.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { return; } if (clock.instant().isAfter(nextHeartbeat)) { - missed = MISSED_HEARTBEAT_CLOSE_REQUEST; - } else { - scheduleHeartbeatCheck(); + logger.warning( + String.format("Missed heartbeat for %s, forcing session close", info.getLogName())); + // forceClose acquires the lock again; safe because Java monitors are re-entrant. + forceClose(MISSED_HEARTBEAT_CLOSE_REQUEST); + return; } - } - if (missed != null) { - logger.warning( - String.format("Missed heartbeat for %s, forcing session close", info.getLogName())); - // forceClose acquires the lock again and performs its own state checks. - forceClose(missed); + scheduleHeartbeatCheck(); } } // region SessionStream event handlers private void dispatchResponseMessage(SessionResponse message) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); switch (message.getPayloadCase()) { case OPEN_SESSION: handleOpenSessionResponse(message.getOpenSession()); @@ -476,6 +487,7 @@ private void dispatchResponseMessage(SessionResponse message) { } private void handleOpenSessionResponse(OpenSessionResponse openSession) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); logger.fine(String.format("%s Session is ready", info.getLogName())); PeerInfo localPeerInfo; @@ -502,6 +514,7 @@ private void handleOpenSessionResponse(OpenSessionResponse openSession) { } private void handleSessionParamsResponse(SessionParametersResponse resp) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); synchronized (lock) { if (state.phase >= SessionState.CLOSING.phase) { logger.fine( @@ -525,6 +538,7 @@ private void handleSessionParamsResponse(SessionParametersResponse resp) { } private void handleVRpcResponse(VirtualRpcResponse vrpc) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); // TODO: when stream is supported this should be updated to the next expected time instead of // session life time this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); @@ -587,10 +601,12 @@ private void handleVRpcResponse(VirtualRpcResponse vrpc) { } private void handleHeartBeatResponse(HeartbeatResponse ignored) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); this.nextHeartbeat = clock.instant().plus(heartbeatInterval); } private void handleSessionRefreshConfigResponse(SessionRefreshConfig config) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); synchronized (lock) { Metadata grpcMetadata = new Metadata(); config @@ -606,6 +622,7 @@ private void handleSessionRefreshConfigResponse(SessionRefreshConfig config) { } private void handleVRpcErrorResponse(ErrorResponse error) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); // Skips the heartbeat check when there's no active vrpc on the session this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); @@ -667,6 +684,7 @@ private void handleVRpcErrorResponse(ErrorResponse error) { } private void handleGoAwayResponse(GoAwayResponse goAwayResponse) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); synchronized (lock) { if (state.phase >= SessionState.CLOSING.phase) { debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_go_away_ignored"); @@ -703,6 +721,7 @@ private void handleUnknownResponseMessage(SessionResponse message) { } private void dispatchStreamClosed(Status status, Metadata trailers) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); SessionState prevState; VRpcImpl localVRpc; From 2efdbc896e5e8dea0eab2abd56fe4620c492bc0e Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Sun, 14 Jun 2026 03:17:23 +0000 Subject: [PATCH 2/6] chore: make startRpc and cancelRpc async via sessionSyncContext MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SessionImpl.startRpc and cancelRpc now submit to sessionSyncContext rather than running synchronously on the caller. VRpcSessionApi.startRpc is void — errors flow through rpc.handleError() onto ctx.getExecutor(). VRpcImpl drops its synchronous post-startRpc error branch. --- .../data/v2/internal/session/SessionImpl.java | 104 +++++++++--------- .../data/v2/internal/session/VRpcImpl.java | 37 +++---- 2 files changed, 71 insertions(+), 70 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 f8220550083a..ba738d3fb007 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 @@ -382,41 +382,48 @@ VRpc newCall(VRpcDescriptor descriptor) { } @Override - public Status startRpc(VRpcImpl rpc, VirtualRpcRequest payload) { - synchronized (lock) { - if (currentRpc != null) { - return Status.INTERNAL.withDescription( - "Session error: RPC multiplexing is not yet supported"); - } - if (state != SessionState.READY) { - return Status.INTERNAL.withDescription( - "Session error: Session was not ready, state = " + state); - } + public void startRpc(VRpcImpl rpc, VirtualRpcRequest payload) { + sessionSyncContext.execute( + () -> { + synchronized (lock) { + if (currentRpc != null) { + rpc.handleError( + VRpcResult.createUncommitedError( + Status.INTERNAL.withDescription( + "Session error: RPC multiplexing is not yet supported"))); + return; + } + if (state != SessionState.READY) { + rpc.handleError( + VRpcResult.createUncommitedError( + Status.INTERNAL.withDescription( + "Session error: Session was not ready, state = " + state))); + return; + } - this.currentRpc = rpc; - stream.sendMessage(SessionRequest.newBuilder().setVirtualRpc(payload).build()); - // Start monitoring for heartbeat when the vRPC is started. heartbeatInterval is read - // inside the lock to avoid a race with handleSessionParamsResponse(). nextHeartbeat is - // volatile and written here without an atomicity guarantee — that is intentional; it is - // only a scheduling hint (see field comment). - this.nextHeartbeat = clock.instant().plus(heartbeatInterval); - // Arm the heartbeat check only while a vRPC is in flight. handleVRpcResponse / - // handleVRpcErrorResponse cancel it on completion; updateState cancels on shutdown. - scheduleHeartbeatCheck(); - return Status.OK; - } + this.currentRpc = rpc; + stream.sendMessage(SessionRequest.newBuilder().setVirtualRpc(payload).build()); + this.nextHeartbeat = clock.instant().plus(heartbeatInterval); + // Arm the heartbeat check only while a vRPC is in flight. handleVRpcResponse / + // handleVRpcErrorResponse cancel it on completion; updateState cancels on shutdown. + scheduleHeartbeatCheck(); + } + }); } @Override public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable cause) { - synchronized (lock) { - if (currentRpc != null && rpcId == currentRpc.rpcId) { - currentCancel = - VRpcResult.createRejectedError( - Status.CANCELLED.withDescription(message).withCause(cause)); - } - // do nothing if the rpc is already finished - } + sessionSyncContext.execute( + () -> { + synchronized (lock) { + if (currentRpc != null && rpcId == currentRpc.rpcId) { + currentCancel = + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(message).withCause(cause)); + } + // do nothing if the rpc is already finished + } + }); } @GuardedBy("lock") @@ -542,9 +549,8 @@ private void handleVRpcResponse(VirtualRpcResponse vrpc) { // TODO: when stream is supported this should be updated to the next expected time instead of // session life time this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); - VRpcImpl localRpc; - VRpcResult localCancel; - + VRpcImpl rpc; + VRpcResult cancel; boolean needsClose; synchronized (lock) { @@ -574,22 +580,22 @@ private void handleVRpcResponse(VirtualRpcResponse vrpc) { vrpc.getRpcId()); // reset state of the current rpc - localCancel = currentCancel; - currentCancel = null; - localRpc = currentRpc; + rpc = currentRpc; + cancel = currentCancel; // TODO: handle multiplexing currentRpc = null; + currentCancel = null; needsClose = (state == SessionState.CLOSING); // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. cancelHeartbeatTimeout(); } - if (localCancel != null) { - tracer.onVRpcClose(localCancel.getStatus().getCode()); - localRpc.handleError(localCancel); + if (cancel != null) { + tracer.onVRpcClose(cancel.getStatus().getCode()); + rpc.handleError(cancel); } else { tracer.onVRpcClose(Status.OK.getCode()); - localRpc.handleResponse(vrpc); + rpc.handleResponse(vrpc); } if (needsClose) { synchronized (lock) { @@ -626,9 +632,9 @@ private void handleVRpcErrorResponse(ErrorResponse error) { // Skips the heartbeat check when there's no active vrpc on the session this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); - VRpcImpl localRpc; + VRpcImpl rpc; + VRpcResult cancel; boolean needsClose; - VRpcResult localCancel; synchronized (lock) { if (state.phase > SessionState.CLOSING.phase) { @@ -658,21 +664,21 @@ private void handleVRpcErrorResponse(ErrorResponse error) { error.getRpcId()); // reset the state of the current rpc - localCancel = currentCancel; - currentCancel = null; - localRpc = currentRpc; + rpc = currentRpc; + cancel = currentCancel; currentRpc = null; + currentCancel = null; needsClose = (state == SessionState.CLOSING); // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. cancelHeartbeatTimeout(); } - if (localCancel != null) { - tracer.onVRpcClose(localCancel.getStatus().getCode()); - localRpc.handleError(localCancel); + if (cancel != null) { + tracer.onVRpcClose(cancel.getStatus().getCode()); + rpc.handleError(cancel); } else { tracer.onVRpcClose(Status.fromCodeValue(error.getStatus().getCode()).getCode()); - localRpc.handleError(VRpcResult.createServerError(error)); + rpc.handleError(VRpcResult.createServerError(error)); } if (needsClose) { synchronized (lock) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java index 95b83b8b4cb6..83228fd43456 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/session/VRpcImpl.java @@ -54,7 +54,11 @@ class VRpcImpl rpc, VirtualRpcRequest payload); + /** + * Submit the vRPC for sending. Async: errors are delivered via {@link + * VRpcImpl#handleError(VRpcResult)}, which dispatches onto {@code ctx.getExecutor()}. + */ + void startRpc(VRpcImpl rpc, VirtualRpcRequest payload); void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable cause); } @@ -122,26 +126,17 @@ public void start(ReqT req, VRpcCallContext ctx, VRpcListener listener) { .setTraceparent(ctx.getTraceParent()) .build(); ctx.getTracer().onRequestSent(peerInfo); - Status status = - session.startRpc( - this, - VirtualRpcRequest.newBuilder() - .setRpcId(rpcId) - .setMetadata(vRpcMetadata) - .setDeadline( - Durations.fromNanos( - ctx.getOperationInfo().getDeadline().timeRemaining(TimeUnit.NANOSECONDS))) - .setPayload(desc.encode(req)) - .build()); - if (!status.isOk()) { - debugTagTracer.checkPrecondition( - state.compareAndSet(State.STARTED, State.CLOSED), - "vrpc_incorrect_start_state", - "VRpc has incorrect state. Expected to be started but was %s", - state); - VRpcResult result = VRpcResult.createUncommitedError(status); - ctx.getExecutor().execute(() -> listener.onClose(result)); - } + session.startRpc( + this, + VirtualRpcRequest.newBuilder() + .setRpcId(rpcId) + .setMetadata(vRpcMetadata) + .setDeadline( + Durations.fromNanos( + ctx.getOperationInfo().getDeadline().timeRemaining(TimeUnit.NANOSECONDS))) + .setPayload(desc.encode(req)) + .build()); + // Session delivers startRpc errors asynchronously via handleError() on ctx.getExecutor(). } void handleSessionClose(VRpcResult result) { From d96b8781bb6fb0d719dd37261acc5ecd83a1a8cf Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 28 May 2026 01:37:54 +0000 Subject: [PATCH 3/6] chore: remove synchronized(lock) from SessionImpl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All session state mutations now run on sessionSyncContext, so the per-session Object lock is no longer needed. Public methods (start, close, forceClose, startRpc, cancelRpc) submit onto sessionSyncContext; nextRpcId becomes AtomicLong for the cross-thread newCall() caller. handleVRpcResponse and handleVRpcErrorResponse drop the localCancel/localRpc capture-and-recheck dance — sessionSyncContext serializes them now. Stale lock-era comments on the fields are replaced with a sessionSyncContext-ownership note. SessionImplTest.testHeartbeat polls for the now-async nextHeartbeat update. --- .../data/v2/internal/session/SessionImpl.java | 680 ++++++++---------- .../v2/internal/session/SessionImplTest.java | 10 +- 2 files changed, 304 insertions(+), 386 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 ba738d3fb007..997835f12205 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 @@ -51,10 +51,10 @@ import java.util.Locale; import java.util.Optional; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; -import javax.annotation.concurrent.GuardedBy; /** Wraps a Bidi ClientCall and layers session semantics on top. */ @VisibleForTesting @@ -78,18 +78,10 @@ public class SessionImpl implements Session, VRpcSessionApi { .setDescription("missed heartbeat") .build(); - /* - * This lock should be mostly uncontended - all access should be naturally interleaved. Contention - * can only really happen when an unsolicited gRPC control message (ie GOAWAY) arrives at the same - * time as newCall or cancel. - * TODO: Contention will increase when multiplexing is implemented. - */ - private final Object lock = new Object(); - private final Clock clock; private final BigtableTimer timer; - // Serializes all session-stream callbacks. Coexists with `lock` for now; `lock` will be removed - // in a follow-up commit once every handler is on the syncContext. + // Serializes all session state mutations. Stream callbacks and the heartbeat tick dispatch + // onto it; every handler below runs inside a syncContext task. private final SynchronizationContext sessionSyncContext; private final SessionTracer tracer; @@ -97,63 +89,46 @@ public class SessionImpl implements Session, VRpcSessionApi { private final SessionInfo info; - @GuardedBy("lock") private final SessionStream stream; - @GuardedBy("lock") - private SessionState state = SessionState.NEW; + private volatile SessionState state = SessionState.NEW; - @GuardedBy("lock") - private Instant lastStateChangedAt; + private volatile Instant lastStateChangedAt; - // Set once under lock in start(), then read freely from gRPC callbacks without the lock. - // Safe because start() is always called before any callback fires, so the write is - // visible to all subsequent readers through the happens-before chain from stream.start(). + // All fields below are owned by sessionSyncContext: every writer runs inside a + // sessionSyncContext task, and the in-class readers do too (handlers, scheduled heartbeat + // tick). They lost their volatile / lock guards when synchronized(lock) was removed; + // SyncContext serialization is what makes plain reads/writes safe. + // + // External callers of getOpenParams / isOpenParamsUpdated / getNextHeartbeat must therefore + // either run on sessionSyncContext themselves (e.g. via a Session.Listener callback, which is + // dispatched from inside the context) or accept a possibly-stale snapshot. There are no + // off-context production readers today. private Listener sessionListener; - // volatile: written under lock in handleSessionRefreshConfigResponse(); read without lock in - // getOpenParams() and isOpenParamsUpdated() so callers get a consistent (if possibly stale) - // snapshot without contending on the lock. Stale reads are acceptable for these accessors. - private volatile OpenParams openParams; + private OpenParams openParams; - private volatile boolean openParamsUpdated; + private boolean openParamsUpdated; - // closeReason is written under lock in close(), forceClose(), handleGoAwayResponse(), and - // dispatchStreamClosed(). The one read that occurs outside the lock — in dispatchStreamClosed - // after the synchronized block — runs on the same gRPC callback thread that just released the - // lock, so the lock's release-acquire edge provides the necessary visibility. A stale read is - // structurally impossible given the control flow (closeReason is always set before the lock - // is released on every path that reaches that read site). @Nullable private CloseSessionRequest closeReason = null; - @GuardedBy("lock") - private long nextRpcId = 1; + private final AtomicLong nextRpcId = new AtomicLong(1); // TODO: replace with a map when implementing multiplexing - @GuardedBy("lock") private VRpcImpl currentRpc = null; - @GuardedBy("lock") private VRpcResult currentCancel = null; - @GuardedBy("lock") private SessionParametersResponse sessionParameters = DEFAULT_SESSION_PARAMS; - // volatile: written under lock in handleSessionParamsResponse(); read without lock in - // handleHeartBeatResponse() where a stale read is acceptable — the heartbeat deadline is a - // soft scheduling hint, not a correctness invariant. startRpc() reads this inside the lock. - private volatile Duration heartbeatInterval = + private Duration heartbeatInterval = Duration.ofMillis(Durations.toMillis(sessionParameters.getKeepAlive())); - // volatile: written from multiple sites without holding the lock (startRpc, handleVRpc*, - // handleHeartBeatResponse). Stale reads are acceptable — nextHeartbeat is used only as a - // scheduling hint by the per-session heartbeat tick. - private volatile Instant nextHeartbeat; + private Instant nextHeartbeat; // Handle for the in-flight heartbeat tick (one outstanding at a time). Set under lock when the // session enters READY (handleOpenSessionResponse) and again from checkHeartbeat to chain the // next tick. Cancelled under lock from updateState when the session transitions past READY. - @GuardedBy("lock") @Nullable private BigtableTimer.Timeout heartbeatTimeout; @@ -193,16 +168,12 @@ public SessionImpl( @Override public SessionState getState() { - synchronized (lock) { - return state; - } + return state; } @Override public Instant getLastStateChange() { - synchronized (lock) { - return lastStateChangedAt; - } + return lastStateChangedAt; } @Override @@ -222,13 +193,7 @@ public Instant getNextHeartbeat() { @Override public PeerInfo getPeerInfo() { - // This lock might not be necessary, its populated once on a gRPC callback which should - // establish a happens before relationship. However access to the underlying stream is guarded - // with errorprone, so sync block is required to get around the lint. - // TODO: consider removing the sync block - synchronized (lock) { - return stream.getPeerInfo(); - } + return stream.getPeerInfo(); } @Override @@ -238,98 +203,101 @@ public String getLogName() { @Override public void forceClose(CloseSessionRequest closeReason) { - synchronized (lock) { - debugTagTracer.checkPrecondition( - state != SessionState.NEW, - "session_force_close_wrong_state", - "Tried to forceClose an unstarted session %s in state %s", - info.getLogName(), - state); - - if (state == SessionState.CLOSED) { - return; - } + sessionSyncContext.execute( + () -> { + debugTagTracer.checkPrecondition( + state != SessionState.NEW, + "session_force_close_wrong_state", + "Tried to forceClose an unstarted session %s in state %s", + info.getLogName(), + state); + + if (state == SessionState.CLOSED) { + return; + } - updateState(SessionState.WAIT_SERVER_CLOSE); - this.closeReason = closeReason; + updateState(SessionState.WAIT_SERVER_CLOSE); + this.closeReason = closeReason; - // Not sending the CloseSessionRequest because cancel() will just drop it - stream.forceClose(closeReason.getDescription(), null); - // Listeners will be notified by dispatchStreamClosed - } + // Not sending the CloseSessionRequest because cancel() will just drop it + stream.forceClose(closeReason.getDescription(), null); + // Listeners will be notified by dispatchStreamClosed + }); } @Override public void start(OpenSessionRequest req, Metadata headers, Listener sessionListener) { - synchronized (lock) { - debugTagTracer.checkPrecondition( - state == SessionState.NEW, - "session_start_wrong_state", - "Tried to start a started session, current state: %s", - state); - - logger.fine(String.format("Starting session %s", info.getLogName())); - tracer.onStart(); - - updateState(SessionState.STARTING); - openParams = OpenParams.create(headers, req); - this.sessionListener = sessionListener; - - SessionRequest wrappedReq = SessionRequest.newBuilder().setOpenSession(req).build(); - stream.start( - new SessionStream.Listener() { - @Override - public void onBeforeSessionStart(PeerInfo peerInfo) {} - - @Override - public void onMessage(SessionResponse message) { - sessionSyncContext.execute(() -> dispatchResponseMessage(message)); - } - - @Override - public void onClose(Status status, Metadata trailers) { - sessionSyncContext.execute(() -> dispatchStreamClosed(status, trailers)); - } - }, - headers); - - stream.sendMessage(wrappedReq); - } + sessionSyncContext.execute( + () -> { + debugTagTracer.checkPrecondition( + state == SessionState.NEW, + "session_start_wrong_state", + "Tried to start a started session, current state: %s", + state); + + logger.fine(String.format("Starting session %s", info.getLogName())); + tracer.onStart(); + + updateState(SessionState.STARTING); + openParams = OpenParams.create(headers, req); + this.sessionListener = sessionListener; + + SessionRequest wrappedReq = SessionRequest.newBuilder().setOpenSession(req).build(); + stream.start( + new SessionStream.Listener() { + @Override + public void onBeforeSessionStart(PeerInfo peerInfo) {} + + @Override + public void onMessage(SessionResponse message) { + sessionSyncContext.execute(() -> dispatchResponseMessage(message)); + } + + @Override + public void onClose(Status status, Metadata trailers) { + sessionSyncContext.execute(() -> dispatchStreamClosed(status, trailers)); + } + }, + headers); + + stream.sendMessage(wrappedReq); + }); } @Override public void close(CloseSessionRequest req) { logger.fine(String.format("Closing session %s for reason: %s", info.getLogName(), req)); - synchronized (lock) { - // Throw an exception because this is a bug and we dont have a listener - debugTagTracer.checkPrecondition( - state != SessionState.NEW, - "session_close_wrong_state", - "Session error: Caller tried to close session %s before starting it with the reason: %s", - info.getLogName(), - req); - - // Multiple close is a no-op - if (state.phase >= SessionState.CLOSING.phase) { - logger.fine( - String.format( - "Session error: Caller tried to close a session %s that is %s for reason: %s", - info.getLogName(), state, req)); - return; - } + sessionSyncContext.execute( + () -> { + // Throw an exception because this is a bug and we dont have a listener + debugTagTracer.checkPrecondition( + state != SessionState.NEW, + "session_close_wrong_state", + "Session error: Caller tried to close session %s before starting it with the reason:" + + " %s", + info.getLogName(), + req); + + // Multiple close is a no-op + if (state.phase >= SessionState.CLOSING.phase) { + logger.fine( + String.format( + "Session error: Caller tried to close a session %s that is %s for reason: %s", + info.getLogName(), state, req)); + return; + } - closeReason = req; - updateState(SessionState.CLOSING); + closeReason = req; + updateState(SessionState.CLOSING); - if (currentRpc == null) { - startGracefulClose(); - } - } + if (currentRpc == null) { + startGracefulClose(); + } + }); } /** Wraps the flow of closing a session. */ - @GuardedBy("lock") private void startGracefulClose() { debugTagTracer.checkPrecondition( state == SessionState.CLOSING, @@ -369,45 +337,40 @@ VRpc newCall(VRpcDescriptor descriptor) { "session_new_call_wrong_type", "wrong VRpc descriptor type"); - synchronized (lock) { - debugTagTracer.checkPrecondition( - state != SessionState.NEW, - "session_new_call_wrong_state", - "Session error: newCall called before start"); + debugTagTracer.checkPrecondition( + state != SessionState.NEW, + "session_new_call_wrong_state", + "Session error: newCall called before start"); - long rpcId = nextRpcId; - nextRpcId = Math.incrementExact(nextRpcId); - return new VRpcImpl<>(this, descriptor, rpcId, stream.getPeerInfo(), debugTagTracer); - } + long rpcId = nextRpcId.getAndIncrement(); + return new VRpcImpl<>(this, descriptor, rpcId, stream.getPeerInfo(), debugTagTracer); } @Override public void startRpc(VRpcImpl rpc, VirtualRpcRequest payload) { sessionSyncContext.execute( () -> { - synchronized (lock) { - if (currentRpc != null) { - rpc.handleError( - VRpcResult.createUncommitedError( - Status.INTERNAL.withDescription( - "Session error: RPC multiplexing is not yet supported"))); - return; - } - if (state != SessionState.READY) { - rpc.handleError( - VRpcResult.createUncommitedError( - Status.INTERNAL.withDescription( - "Session error: Session was not ready, state = " + state))); - return; - } - - this.currentRpc = rpc; - stream.sendMessage(SessionRequest.newBuilder().setVirtualRpc(payload).build()); - this.nextHeartbeat = clock.instant().plus(heartbeatInterval); - // Arm the heartbeat check only while a vRPC is in flight. handleVRpcResponse / - // handleVRpcErrorResponse cancel it on completion; updateState cancels on shutdown. - scheduleHeartbeatCheck(); + if (currentRpc != null) { + rpc.handleError( + VRpcResult.createUncommitedError( + Status.INTERNAL.withDescription( + "Session error: RPC multiplexing is not yet supported"))); + return; } + if (state != SessionState.READY) { + rpc.handleError( + VRpcResult.createUncommitedError( + Status.INTERNAL.withDescription( + "Session error: Session was not ready, state = " + state))); + return; + } + + this.currentRpc = rpc; + stream.sendMessage(SessionRequest.newBuilder().setVirtualRpc(payload).build()); + this.nextHeartbeat = clock.instant().plus(heartbeatInterval); + // Arm the heartbeat check only while a vRPC is in flight. handleVRpcResponse / + // handleVRpcErrorResponse cancel it on completion; updateState cancels on shutdown. + scheduleHeartbeatCheck(); }); } @@ -415,18 +378,15 @@ public void startRpc(VRpcImpl rpc, VirtualRpcRequest payload) { public void cancelRpc(long rpcId, @Nullable String message, @Nullable Throwable cause) { sessionSyncContext.execute( () -> { - synchronized (lock) { - if (currentRpc != null && rpcId == currentRpc.rpcId) { - currentCancel = - VRpcResult.createRejectedError( - Status.CANCELLED.withDescription(message).withCause(cause)); - } - // do nothing if the rpc is already finished + if (currentRpc != null && rpcId == currentRpc.rpcId) { + currentCancel = + VRpcResult.createRejectedError( + Status.CANCELLED.withDescription(message).withCause(cause)); } + // do nothing if the rpc is already finished }); } - @GuardedBy("lock") private void scheduleHeartbeatCheck() { heartbeatTimeout = timer.newTimeout( @@ -435,7 +395,6 @@ private void scheduleHeartbeatCheck() { TimeUnit.MILLISECONDS); } - @GuardedBy("lock") private void cancelHeartbeatTimeout() { if (heartbeatTimeout != null) { heartbeatTimeout.cancel(); @@ -447,19 +406,16 @@ private void cancelHeartbeatTimeout() { // heartbeat deadline has passed and force-closes on miss; otherwise re-schedules. private void checkHeartbeat() { sessionSyncContext.throwIfNotInThisSynchronizationContext(); - synchronized (lock) { - if (state.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { - return; - } - if (clock.instant().isAfter(nextHeartbeat)) { - logger.warning( - String.format("Missed heartbeat for %s, forcing session close", info.getLogName())); - // forceClose acquires the lock again; safe because Java monitors are re-entrant. - forceClose(MISSED_HEARTBEAT_CLOSE_REQUEST); - return; - } - scheduleHeartbeatCheck(); + if (state.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { + return; } + if (clock.instant().isAfter(nextHeartbeat)) { + logger.warning( + String.format("Missed heartbeat for %s, forcing session close", info.getLogName())); + forceClose(MISSED_HEARTBEAT_CLOSE_REQUEST); + return; + } + scheduleHeartbeatCheck(); } // region SessionStream event handlers @@ -497,50 +453,42 @@ private void handleOpenSessionResponse(OpenSessionResponse openSession) { sessionSyncContext.throwIfNotInThisSynchronizationContext(); logger.fine(String.format("%s Session is ready", info.getLogName())); - PeerInfo localPeerInfo; - - synchronized (lock) { - debugTagTracer.checkPrecondition( - state != SessionState.NEW, - "session_open_wrong_state", - "Got session open response before session started"); - debugTagTracer.checkPrecondition( - state != SessionState.CLOSED, - "session_open_wrong_state", - "Got session open response after session was closed"); + debugTagTracer.checkPrecondition( + state != SessionState.NEW, + "session_open_wrong_state", + "Got session open response before session started"); + debugTagTracer.checkPrecondition( + state != SessionState.CLOSED, + "session_open_wrong_state", + "Got session open response after session was closed"); - if (state != SessionState.STARTING) { - logger.fine(String.format("Stream was already %s when session open was received", state)); - return; - } - localPeerInfo = stream.getPeerInfo(); - updateState(SessionState.READY); + if (state != SessionState.STARTING) { + logger.fine(String.format("Stream was already %s when session open was received", state)); + return; } + PeerInfo localPeerInfo = stream.getPeerInfo(); + updateState(SessionState.READY); tracer.onOpen(localPeerInfo); sessionListener.onReady(openSession); } private void handleSessionParamsResponse(SessionParametersResponse resp) { - sessionSyncContext.throwIfNotInThisSynchronizationContext(); - synchronized (lock) { - if (state.phase >= SessionState.CLOSING.phase) { - logger.fine( - String.format("Stream was already %s when session params were received", state)); - return; - } + if (state.phase >= SessionState.CLOSING.phase) { + logger.fine(String.format("Stream was already %s when session params were received", state)); + return; + } - if (!sessionParameters.equals(resp)) { - this.sessionParameters = resp; - this.heartbeatInterval = - Duration.ofMillis(Durations.toMillis(sessionParameters.getKeepAlive())); - logger.log( - Level.CONFIG, - () -> - String.format( - "%s session params changed: %s", - info.getLogName(), - TextFormat.printer().emittingSingleLine(true).printToString(resp))); - } + if (!sessionParameters.equals(resp)) { + this.sessionParameters = resp; + this.heartbeatInterval = + Duration.ofMillis(Durations.toMillis(sessionParameters.getKeepAlive())); + logger.log( + Level.CONFIG, + () -> + String.format( + "%s session params changed: %s", + info.getLogName(), + TextFormat.printer().emittingSingleLine(true).printToString(resp))); } } @@ -549,47 +497,41 @@ private void handleVRpcResponse(VirtualRpcResponse vrpc) { // TODO: when stream is supported this should be updated to the next expected time instead of // session life time this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); - VRpcImpl rpc; - VRpcResult cancel; - boolean needsClose; - - synchronized (lock) { - if (state.phase > SessionState.CLOSING.phase) { - debugTagTracer.record( - TelemetryConfiguration.Level.WARN, "session_closed_discard_vrpc_response"); - logger.warning( - String.format( - "%s Discarding vRPC error because session is past the CLOSING phase with the" - + " reason: %s", - info.getLogName(), closeReason)); - return; - } - debugTagTracer.checkPrecondition( - state == SessionState.READY || state == SessionState.CLOSING, - "session_vrpc_response_wrong_state", - "Unexpected vRPC response when session is %s", - state); - debugTagTracer.checkPrecondition( - currentRpc != null, "session_vrpc_null", "Got vRPC response but current vRPC is unset"); - debugTagTracer.checkPrecondition( - currentRpc.rpcId == vrpc.getRpcId(), - "session_vrpc_id_mismatch", - "Got vRPC response for the wrong vRPC: expect: %s, actual: %s", - currentRpc.rpcId, - vrpc.getRpcId()); - - // reset state of the current rpc - rpc = currentRpc; - cancel = currentCancel; - // TODO: handle multiplexing - currentRpc = null; - currentCancel = null; - needsClose = (state == SessionState.CLOSING); - // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. - cancelHeartbeatTimeout(); + if (state.phase > SessionState.CLOSING.phase) { + debugTagTracer.record( + TelemetryConfiguration.Level.WARN, "session_closed_discard_vrpc_response"); + logger.warning( + String.format( + "%s Discarding vRPC error because session is past the CLOSING phase with the" + + " reason: %s", + info.getLogName(), closeReason)); + return; } + debugTagTracer.checkPrecondition( + state == SessionState.READY || state == SessionState.CLOSING, + "session_vrpc_response_wrong_state", + "Unexpected vRPC response when session is %s", + state); + debugTagTracer.checkPrecondition( + currentRpc != null, "session_vrpc_null", "Got vRPC response but current vRPC is unset"); + debugTagTracer.checkPrecondition( + currentRpc.rpcId == vrpc.getRpcId(), + "session_vrpc_id_mismatch", + "Got vRPC response for the wrong vRPC: expect: %s, actual: %s", + currentRpc.rpcId, + vrpc.getRpcId()); + + // reset state of the current rpc + VRpcImpl rpc = currentRpc; + VRpcResult cancel = currentCancel; + // TODO: handle multiplexing + currentRpc = null; + currentCancel = null; + // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. + cancelHeartbeatTimeout(); + if (cancel != null) { tracer.onVRpcClose(cancel.getStatus().getCode()); rpc.handleError(cancel); @@ -597,12 +539,8 @@ private void handleVRpcResponse(VirtualRpcResponse vrpc) { tracer.onVRpcClose(Status.OK.getCode()); rpc.handleResponse(vrpc); } - if (needsClose) { - synchronized (lock) { - if (state == SessionState.CLOSING) { - startGracefulClose(); - } - } + if (state == SessionState.CLOSING) { + startGracefulClose(); } } @@ -612,19 +550,16 @@ private void handleHeartBeatResponse(HeartbeatResponse ignored) { } private void handleSessionRefreshConfigResponse(SessionRefreshConfig config) { - sessionSyncContext.throwIfNotInThisSynchronizationContext(); - synchronized (lock) { - Metadata grpcMetadata = new Metadata(); - config - .getMetadataList() - .forEach( - entry -> - grpcMetadata.put( - Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER), - entry.getValue().toStringUtf8())); - openParams = OpenParams.create(grpcMetadata, config.getOptimizedOpenRequest()); - openParamsUpdated = true; - } + Metadata grpcMetadata = new Metadata(); + config + .getMetadataList() + .forEach( + entry -> + grpcMetadata.put( + Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER), + entry.getValue().toStringUtf8())); + openParams = OpenParams.create(grpcMetadata, config.getOptimizedOpenRequest()); + openParamsUpdated = true; } private void handleVRpcErrorResponse(ErrorResponse error) { @@ -632,46 +567,39 @@ private void handleVRpcErrorResponse(ErrorResponse error) { // Skips the heartbeat check when there's no active vrpc on the session this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); - VRpcImpl rpc; - VRpcResult cancel; - boolean needsClose; + if (state.phase > SessionState.CLOSING.phase) { + debugTagTracer.record( + TelemetryConfiguration.Level.WARN, "session_closed_discard_vrpc_response"); + logger.warning( + String.format( + "%s Discarding vRPC error because session is past the CLOSING phase with the" + + " reason: %s, error was: %s", + info.getLogName(), closeReason, error)); + return; + } - synchronized (lock) { - if (state.phase > SessionState.CLOSING.phase) { - debugTagTracer.record( - TelemetryConfiguration.Level.WARN, "session_closed_discard_vrpc_response"); - logger.warning( - String.format( - "%s Discarding vRPC error because session is past the CLOSING phase with the" - + " reason: %s, error was: %s", - info.getLogName(), closeReason, error)); - return; - } + debugTagTracer.checkPrecondition( + state == SessionState.READY || state == SessionState.CLOSING, + "session_vrpc_response_wrong_state", + "Unexpected vRPC response when session is %s", + state); - debugTagTracer.checkPrecondition( - state == SessionState.READY || state == SessionState.CLOSING, - "session_vrpc_response_wrong_state", - "Unexpected vRPC response when session is %s", - state); - - debugTagTracer.checkPrecondition( - currentRpc != null, "session_vrpc_null", "Got vRPC response but current vRPC is unset"); - debugTagTracer.checkPrecondition( - currentRpc.rpcId == error.getRpcId(), - "session_vrpc_id_mismatch", - "Got vRPC response for the wrong vRPC: expect: %s, actual: %s", - currentRpc.rpcId, - error.getRpcId()); - - // reset the state of the current rpc - rpc = currentRpc; - cancel = currentCancel; - currentRpc = null; - currentCancel = null; - needsClose = (state == SessionState.CLOSING); - // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. - cancelHeartbeatTimeout(); - } + debugTagTracer.checkPrecondition( + currentRpc != null, "session_vrpc_null", "Got vRPC response but current vRPC is unset"); + debugTagTracer.checkPrecondition( + currentRpc.rpcId == error.getRpcId(), + "session_vrpc_id_mismatch", + "Got vRPC response for the wrong vRPC: expect: %s, actual: %s", + currentRpc.rpcId, + error.getRpcId()); + + // reset the state of the current rpc + VRpcImpl rpc = currentRpc; + VRpcResult cancel = currentCancel; + currentRpc = null; + currentCancel = null; + // No active vRPC means no useful heartbeat deadline; drop the in-flight tick. + cancelHeartbeatTimeout(); if (cancel != null) { tracer.onVRpcClose(cancel.getStatus().getCode()); @@ -680,43 +608,35 @@ private void handleVRpcErrorResponse(ErrorResponse error) { tracer.onVRpcClose(Status.fromCodeValue(error.getStatus().getCode()).getCode()); rpc.handleError(VRpcResult.createServerError(error)); } - if (needsClose) { - synchronized (lock) { - if (state == SessionState.CLOSING) { - startGracefulClose(); - } - } + if (state == SessionState.CLOSING) { + startGracefulClose(); } } private void handleGoAwayResponse(GoAwayResponse goAwayResponse) { - sessionSyncContext.throwIfNotInThisSynchronizationContext(); - synchronized (lock) { - if (state.phase >= SessionState.CLOSING.phase) { - debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_go_away_ignored"); - logger.warning( - String.format( - "Session error: %s Ignoring goaway because session is %s", - info.getLogName(), state)); - return; - } + if (state.phase >= SessionState.CLOSING.phase) { + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_go_away_ignored"); + logger.warning( + String.format( + "Session error: %s Ignoring goaway because session is %s", info.getLogName(), state)); + return; + } - debugTagTracer.checkPrecondition( - state.phase >= SessionState.STARTING.phase, - "session_go_away_wrong_state", - "Unexpected goaway when session is %s", - state); + debugTagTracer.checkPrecondition( + state.phase >= SessionState.STARTING.phase, + "session_go_away_wrong_state", + "Unexpected goaway when session is %s", + state); - updateState(SessionState.CLOSING); - closeReason = - CloseSessionRequest.newBuilder() - .setReason(CloseSessionReason.CLOSE_SESSION_REASON_GOAWAY) - .setDescription( - "Server sent GO_AWAY_" + goAwayResponse.getReason().toUpperCase(Locale.ENGLISH)) - .build(); - if (currentRpc == null) { - startGracefulClose(); - } + updateState(SessionState.CLOSING); + closeReason = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_GOAWAY) + .setDescription( + "Server sent GO_AWAY_" + goAwayResponse.getReason().toUpperCase(Locale.ENGLISH)) + .build(); + if (currentRpc == null) { + startGracefulClose(); } sessionListener.onGoAway(goAwayResponse); } @@ -727,51 +647,44 @@ private void handleUnknownResponseMessage(SessionResponse message) { } private void dispatchStreamClosed(Status status, Metadata trailers) { - sessionSyncContext.throwIfNotInThisSynchronizationContext(); - SessionState prevState; - VRpcImpl localVRpc; + SessionState prevState = state; - PeerInfo localPeerInfo; - synchronized (lock) { - prevState = state; + if (!status.isOk()) { + String augmentedDescription = + Optional.ofNullable(status.getDescription()).map(d -> d + ". ").orElse("") + + "PeerInfo: " + + formatPeerInfo(getPeerInfo()); - if (!status.isOk()) { - String augmentedDescription = - Optional.ofNullable(status.getDescription()).map(d -> d + ". ").orElse("") - + "PeerInfo: " - + formatPeerInfo(getPeerInfo()); + status = status.withDescription(augmentedDescription); + } - status = status.withDescription(augmentedDescription); - } + if (state == SessionState.WAIT_SERVER_CLOSE) { + logger.fine(String.format("%s closed normally with status %s", info.getLogName(), status)); + } else { + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_abnormal_close"); + // Unexpected path + String msg = + String.format( + "Session error: %s session closed unexpectedly in state %s. Status: %s", + info.getLogName(), state, status); + logger.warning(msg); - if (state == SessionState.WAIT_SERVER_CLOSE) { - logger.fine(String.format("%s closed normally with status %s", info.getLogName(), status)); - } else { - debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_abnormal_close"); - // Unexpected path - String msg = - String.format( - "Session error: %s session closed unexpectedly in state %s. Status: %s", - info.getLogName(), state, status); - logger.warning(msg); - - if (state == SessionState.CLOSED) { - return; - } - - closeReason = - CloseSessionRequest.newBuilder() - .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) - .setDescription("Unexpected session close with status: " + status.getCode()) - .build(); + if (state == SessionState.CLOSED) { + return; } - localVRpc = currentRpc; - localPeerInfo = stream.getPeerInfo(); - currentRpc = null; - updateState(SessionState.CLOSED); + closeReason = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) + .setDescription("Unexpected session close with status: " + status.getCode()) + .build(); } + VRpcImpl localVRpc = currentRpc; + PeerInfo localPeerInfo = stream.getPeerInfo(); + currentRpc = null; + updateState(SessionState.CLOSED); + if (localVRpc != null) { try { localVRpc.handleSessionClose(VRpcResult.createRemoteTransportError(status, trailers)); @@ -790,7 +703,6 @@ private void dispatchStreamClosed(Status status, Metadata trailers) { sessionListener.onClose(prevState, status, trailers); } - @GuardedBy("lock") private void updateState(SessionState newState) { this.state = newState; this.lastStateChangedAt = clock.instant(); 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 c7739de633ad..3a7911253e1a 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 @@ -497,8 +497,14 @@ void testHeartbeat() throws Exception { VRpcCallContext.create(Deadline.after(1, TimeUnit.MINUTES), true, tracer), f); - assertThat(session.getNextHeartbeat()) - .isEqualTo(time.plus(Duration.ofMillis(keepAliveDurationMs))); + // startRpc() is now async; poll until sessionSyncContext processes it. + Instant expectedHeartbeat = time.plus(Duration.ofMillis(keepAliveDurationMs)); + Stopwatch sw = Stopwatch.createStarted(); + while (!session.getNextHeartbeat().equals(expectedHeartbeat) + && sw.elapsed(TimeUnit.SECONDS) < 5) { + Thread.sleep(10); + } + assertThat(session.getNextHeartbeat()).isEqualTo(expectedHeartbeat); assertThat(f.get()).isEqualTo(SessionFakeScriptedResponse.getDefaultInstance()); From 5312d045c029868209305a8bcff5cdc2ef5e58c4 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 4 Jun 2026 20:57:57 +0000 Subject: [PATCH 4/6] chore: abort session on uncaught exception in sessionSyncContext MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../data/v2/internal/session/SessionImpl.java | 167 ++++++++++++++++-- .../v2/internal/session/SessionImplTest.java | 143 +++++++++++++++ 2 files changed, 292 insertions(+), 18 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 997835f12205..e05c51f6dd08 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 @@ -126,11 +126,13 @@ public class SessionImpl implements Session, VRpcSessionApi { private Instant nextHeartbeat; - // Handle for the in-flight heartbeat tick (one outstanding at a time). Set under lock when the - // session enters READY (handleOpenSessionResponse) and again from checkHeartbeat to chain the - // next tick. Cancelled under lock from updateState when the session transitions past READY. - @Nullable - private BigtableTimer.Timeout heartbeatTimeout; + // Handle for the in-flight heartbeat tick (one outstanding at a time). Cancelled on terminal + // transitions so the wheel doesn't carry a no-op entry until the next fire. + @Nullable private BigtableTimer.Timeout heartbeatTimeout; + + // 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; public SessionImpl( Metrics metrics, @@ -156,14 +158,77 @@ public SessionImpl( this.debugTagTracer = metrics.getDebugTagTracer(); this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); this.openParamsUpdated = false; + // On uncaught exception, drive the session through a clean terminal-close path so the pool + // and the in-flight vRpc are always notified. notifyTerminalClose has local guards, and + // isAborting prevents recursion if the abort path itself throws. this.sessionSyncContext = - new SynchronizationContext( - (thread, e) -> - logger.log( - Level.WARNING, - String.format( - "Uncaught exception in session SyncContext for %s", info.getLogName()), - e)); + new SynchronizationContext((thread, e) -> abortFromUncaughtException(e)); + } + + private void abortFromUncaughtException(Throwable e) { + if (isAborting) { + logger.log( + Level.WARNING, + String.format( + "Session error: %s Secondary uncaught exception during abort, ignoring", + info.getLogName()), + e); + return; + } + isAborting = true; + + logger.log( + Level.SEVERE, + String.format( + "Session error: %s Uncaught exception in session SyncContext in state %s, PeerInfo:" + + " %s — aborting session", + info.getLogName(), state, formatPeerInfo(safeGetPeerInfo())), + e); + + if (state == SessionState.CLOSED) { + return; + } + + // Always overwrite closeReason: the abort is what actually happened, not whatever clean close + // we may have been attempting. Fold the prior reason into the description for forensics so + // downstream metrics (which bucket by reason) attribute this to ERROR rather than the + // interrupted close. + String prevDesc = + (closeReason != null) + ? " (was closing for: " + + closeReason.getReason() + + " — " + + closeReason.getDescription() + + ")" + : ""; + closeReason = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) + .setDescription("Uncaught exception in session SyncContext: " + e + prevDesc) + .build(); + + VRpcImpl localRpc = currentRpc; + currentRpc = null; + SessionState prevState = state; + updateState(SessionState.CLOSED); + + // Defensively tell the transport we're done. Safe on un-started streams via the try/catch. + try { + stream.forceClose("Session aborted due to uncaught exception", e); + } catch (Throwable t) { + logger.log( + Level.WARNING, + String.format( + "Session error: %s Exception while force-closing stream during abort", + info.getLogName()), + t); + } + + notifyTerminalClose( + Status.INTERNAL.withDescription("Session aborted").withCause(e), + new Metadata(), + localRpc, + prevState); } @Override @@ -681,13 +746,45 @@ private void dispatchStreamClosed(Status status, Metadata trailers) { } VRpcImpl localVRpc = currentRpc; - PeerInfo localPeerInfo = stream.getPeerInfo(); currentRpc = null; updateState(SessionState.CLOSED); - if (localVRpc != null) { + notifyTerminalClose(status, trailers, localVRpc, prevState); + } + + /** + * Fan out terminal notifications to the in-flight vRpc, tracer, and session listener with local + * guards so a throw in one notification does not suppress the others. + * + *

Caller contract: must have already transitioned to {@link SessionState#CLOSED} and captured + * and cleared {@code currentRpc}. Callers should also set {@code closeReason}; if missing we + * synthesize a fallback here rather than throw, since throwing from this fan-out aborts the + * remaining notifications and (because the state is already CLOSED) defeats the + * sessionSyncContext uncaught handler's cleanup. + */ + private void notifyTerminalClose( + Status status, + Metadata trailers, + @Nullable VRpcImpl localRpc, + SessionState prevState) { + // Should never happen — matches the synthesizer in startGracefulClose. + if (closeReason == null) { + debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_close_no_reason"); + logger.log( + Level.WARNING, + String.format( + "%s notifyTerminalClose reached without a closeReason; status=%s", + info.getLogName(), status), + new IllegalStateException("notifyTerminalClose without closeReason")); + closeReason = + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_ERROR) + .setDescription("notifyTerminalClose reached without closeReason; status=" + status) + .build(); + } + if (localRpc != null) { try { - localVRpc.handleSessionClose(VRpcResult.createRemoteTransportError(status, trailers)); + localRpc.handleSessionClose(VRpcResult.createRemoteTransportError(status, trailers)); } catch (Throwable t) { logger.log( Level.WARNING, @@ -697,10 +794,44 @@ private void dispatchStreamClosed(Status status, Metadata trailers) { info.getLogName(), status), t); } - tracer.onVRpcClose(Status.UNAVAILABLE.getCode()); + try { + tracer.onVRpcClose(Status.UNAVAILABLE.getCode()); + } catch (Throwable t) { + logger.log( + Level.WARNING, + String.format( + "Session error: %s Unhandled exception in tracer.onVRpcClose", info.getLogName()), + t); + } + } + try { + tracer.onClose(safeGetPeerInfo(), closeReason.getReason(), status); + } catch (Throwable t) { + logger.log( + Level.WARNING, + String.format("Session error: %s Unhandled exception in tracer.onClose", info.getLogName()), + t); + } + if (sessionListener != null) { + try { + sessionListener.onClose(prevState, status, trailers); + } catch (Throwable t) { + logger.log( + Level.WARNING, + String.format( + "Session error: %s Unhandled exception in sessionListener.onClose", + info.getLogName()), + t); + } + } + } + + private PeerInfo safeGetPeerInfo() { + try { + return stream.getPeerInfo(); + } catch (Throwable t) { + return SessionStream.DISCONNECTED_PEER_INFO; } - tracer.onClose(localPeerInfo, closeReason.getReason(), status); - sessionListener.onClose(prevState, status, trailers); } private void updateState(SessionState newState) { 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 3a7911253e1a..2a1d00327da4 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 @@ -672,6 +672,149 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception { assertThat(sessionListener.popUntil(Status.class)).isOk(); } + // region uncaught-exception abort behaviors + + @Test + void abortFiresWhenListenerOnReadyThrows() throws Exception { + SessionImpl session = + new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + + java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.atomic.AtomicReference capturedStatus = + new java.util.concurrent.atomic.AtomicReference<>(); + + Session.Listener throwingListener = + new Session.Listener() { + @Override + public void onReady(OpenSessionResponse msg) { + throw new RuntimeException("simulated onReady failure"); + } + + @Override + public void onGoAway(GoAwayResponse msg) {} + + @Override + public void onClose(Session.SessionState prevState, Status status, Metadata trailers) { + capturedStatus.set(status); + onCloseLatch.countDown(); + } + }; + + session.start( + OpenSessionRequest.newBuilder() + .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString()) + .build(), + new Metadata(), + throwingListener); + + // The abort path must drive the session to CLOSED and notify the listener via onClose, even + // though the original onReady threw. + assertWithMessage("listener.onClose must be invoked after onReady throws") + .that(onCloseLatch.await(5, TimeUnit.SECONDS)) + .isTrue(); + assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); + assertThat(capturedStatus.get().getCode()).isEqualTo(Status.Code.INTERNAL); + } + + @Test + void abortDoesNotHangWhenListenerOnCloseThrows() throws Exception { + SessionImpl session = + new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + + java.util.concurrent.CountDownLatch onReadyLatch = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); + + Session.Listener throwingListener = + new Session.Listener() { + @Override + public void onReady(OpenSessionResponse msg) { + onReadyLatch.countDown(); + } + + @Override + public void onGoAway(GoAwayResponse msg) {} + + @Override + public void onClose(Session.SessionState prevState, Status status, Metadata trailers) { + onCloseLatch.countDown(); + throw new RuntimeException("simulated onClose failure"); + } + }; + + session.start( + OpenSessionRequest.newBuilder() + .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString()) + .build(), + new Metadata(), + throwingListener); + + assertThat(onReadyLatch.await(5, TimeUnit.SECONDS)).isTrue(); + + // Close normally. The listener's onClose throws — the local guard inside notifyTerminalClose + // must swallow it so the SyncContext drain doesn't recurse infinitely or hang. + session.close( + CloseSessionRequest.newBuilder() + .setReason(CloseSessionReason.CLOSE_SESSION_REASON_USER) + .setDescription("test") + .build()); + + assertWithMessage("listener.onClose should be invoked exactly once during normal close") + .that(onCloseLatch.await(5, TimeUnit.SECONDS)) + .isTrue(); + + // The session should reach CLOSED state cleanly within the test timeout. + Stopwatch sw = Stopwatch.createStarted(); + while (session.getState() != Session.SessionState.CLOSED && sw.elapsed().getSeconds() < 5) { + Thread.sleep(10); + } + assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); + } + + @Test + void abortDoesNotInfiniteLoopWhenRecoveryListenerAlsoThrows() throws Exception { + SessionImpl session = + new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + + java.util.concurrent.CountDownLatch onCloseInvoked = + new java.util.concurrent.CountDownLatch(1); + + Session.Listener doublyThrowingListener = + new Session.Listener() { + @Override + public void onReady(OpenSessionResponse msg) { + throw new RuntimeException("simulated onReady failure"); + } + + @Override + public void onGoAway(GoAwayResponse msg) {} + + @Override + public void onClose(Session.SessionState prevState, Status status, Metadata trailers) { + onCloseInvoked.countDown(); + throw new RuntimeException("simulated onClose failure during abort"); + } + }; + + session.start( + OpenSessionRequest.newBuilder() + .setPayload(OpenFakeSessionRequest.getDefaultInstance().toByteString()) + .build(), + new Metadata(), + doublyThrowingListener); + + // onReady throws → abort fires → abort calls onClose, which also throws → Guard 4 swallows + // and isAborting prevents the handler from re-driving abort. The session must reach CLOSED + // without hanging (the @Timeout(30) on the class is the safety net for infinite loops). + assertThat(onCloseInvoked.await(5, TimeUnit.SECONDS)).isTrue(); + Stopwatch sw = Stopwatch.createStarted(); + while (session.getState() != Session.SessionState.CLOSED && sw.elapsed().getSeconds() < 5) { + Thread.sleep(10); + } + assertThat(session.getState()).isEqualTo(Session.SessionState.CLOSED); + } + + // endregion + // 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 { From aa09bfea1b3e854e73e841a748fde005205ac184 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 22 Jun 2026 18:57:57 +0000 Subject: [PATCH 5/6] chore: assert sessionSyncContext on all stream-callback handlers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five SessionImpl handlers reachable from dispatchResponseMessage — handleSessionParamsResponse, handleSessionRefreshConfigResponse, handleGoAwayResponse, handleUnknownResponseMessage — and the dispatchStreamClosed entry point were missing the throwIfNotInThisSynchronizationContext assertion that the other handlers already have. The assertions are strictly redundant today (every caller is either dispatchResponseMessage, which already asserts, or sessionSyncContext.execute), but the asymmetry is incidental and adding them uniformly documents the threading contract and guards against future direct callers. --- .../cloud/bigtable/data/v2/internal/session/SessionImpl.java | 5 +++++ 1 file changed, 5 insertions(+) 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 e05c51f6dd08..677651dedf67 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 @@ -538,6 +538,7 @@ private void handleOpenSessionResponse(OpenSessionResponse openSession) { } private void handleSessionParamsResponse(SessionParametersResponse resp) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); if (state.phase >= SessionState.CLOSING.phase) { logger.fine(String.format("Stream was already %s when session params were received", state)); return; @@ -615,6 +616,7 @@ private void handleHeartBeatResponse(HeartbeatResponse ignored) { } private void handleSessionRefreshConfigResponse(SessionRefreshConfig config) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); Metadata grpcMetadata = new Metadata(); config .getMetadataList() @@ -679,6 +681,7 @@ private void handleVRpcErrorResponse(ErrorResponse error) { } private void handleGoAwayResponse(GoAwayResponse goAwayResponse) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); if (state.phase >= SessionState.CLOSING.phase) { debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_go_away_ignored"); logger.warning( @@ -707,11 +710,13 @@ private void handleGoAwayResponse(GoAwayResponse goAwayResponse) { } private void handleUnknownResponseMessage(SessionResponse message) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); debugTagTracer.record(TelemetryConfiguration.Level.WARN, "session_unknown_response"); logger.warning(String.format("%s Unknown control message: %s", info.getLogName(), message)); } private void dispatchStreamClosed(Status status, Metadata trailers) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); SessionState prevState = state; if (!status.isOk()) { From f285a299a6b63f306e523c3e6792cbd330c8bef2 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Mon, 22 Jun 2026 21:19:10 +0000 Subject: [PATCH 6/6] chore: apply fmt-maven-plugin --- .../data/v2/internal/session/SessionImpl.java | 3 ++- .../data/v2/internal/session/SessionImplTest.java | 12 ++++-------- 2 files changed, 6 insertions(+), 9 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 677651dedf67..1915edf5306e 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 @@ -814,7 +814,8 @@ private void notifyTerminalClose( } catch (Throwable t) { logger.log( Level.WARNING, - String.format("Session error: %s Unhandled exception in tracer.onClose", info.getLogName()), + String.format( + "Session error: %s Unhandled exception in tracer.onClose", info.getLogName()), t); } if (sessionListener != null) { 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 2a1d00327da4..a777a366e6aa 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 @@ -676,8 +676,7 @@ void testHeartbeatScheduledOnlyDuringVRpc() throws Exception { @Test void abortFiresWhenListenerOnReadyThrows() throws Exception { - SessionImpl session = - new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); java.util.concurrent.atomic.AtomicReference capturedStatus = @@ -718,8 +717,7 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai @Test void abortDoesNotHangWhenListenerOnCloseThrows() throws Exception { - SessionImpl session = - new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); java.util.concurrent.CountDownLatch onReadyLatch = new java.util.concurrent.CountDownLatch(1); java.util.concurrent.CountDownLatch onCloseLatch = new java.util.concurrent.CountDownLatch(1); @@ -772,11 +770,9 @@ public void onClose(Session.SessionState prevState, Status status, Metadata trai @Test void abortDoesNotInfiniteLoopWhenRecoveryListenerAlsoThrows() throws Exception { - SessionImpl session = - new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); + SessionImpl session = new SessionImpl(metrics, poolInfo, 0, sessionFactory.createNew(), timer); - java.util.concurrent.CountDownLatch onCloseInvoked = - new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.CountDownLatch onCloseInvoked = new java.util.concurrent.CountDownLatch(1); Session.Listener doublyThrowingListener = new Session.Listener() {