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..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 @@ -44,16 +44,17 @@ 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; 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 @@ -77,81 +78,61 @@ 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 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; private final DebugTagTracer debugTagTracer; 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). Cancelled on terminal + // transitions so the wheel doesn't carry a no-op entry until the next fire. + @Nullable private BigtableTimer.Timeout heartbeatTimeout; - // 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; + // 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, @@ -177,20 +158,87 @@ 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) -> 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 public SessionState getState() { - synchronized (lock) { - return state; - } + return state; } @Override public Instant getLastStateChange() { - synchronized (lock) { - return lastStateChangedAt; - } + return lastStateChangedAt; } @Override @@ -210,13 +258,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 @@ -226,98 +268,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; - } - - 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 - } + 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; + + // 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) { - dispatchResponseMessage(message); - } - - @Override - public void onClose(Status status, Metadata trailers) { - 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; - } - - closeReason = req; - updateState(SessionState.CLOSING); - - if (currentRpc == null) { - startGracefulClose(); - } - } + 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); + + if (currentRpc == null) { + startGracefulClose(); + } + }); } /** Wraps the flow of closing a session. */ - @GuardedBy("lock") private void startGracefulClose() { debugTagTracer.checkPrecondition( state == SessionState.CLOSING, @@ -357,64 +402,64 @@ 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 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); - } - - 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; - } + public void startRpc(VRpcImpl rpc, VirtualRpcRequest payload) { + sessionSyncContext.execute( + () -> { + 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(); + }); } @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( + () -> { + 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( - this::checkHeartbeat, HEARTBEAT_CHECK_INTERVAL.toMillis(), TimeUnit.MILLISECONDS); + () -> sessionSyncContext.execute(this::checkHeartbeat), + HEARTBEAT_CHECK_INTERVAL.toMillis(), + TimeUnit.MILLISECONDS); } - @GuardedBy("lock") private void cancelHeartbeatTimeout() { if (heartbeatTimeout != null) { heartbeatTimeout.cancel(); @@ -422,31 +467,25 @@ 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; - synchronized (lock) { - if (state.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { - return; - } - if (clock.instant().isAfter(nextHeartbeat)) { - missed = MISSED_HEARTBEAT_CLOSE_REQUEST; - } else { - scheduleHeartbeatCheck(); - } + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + if (state.phase >= SessionState.WAIT_SERVER_CLOSE.phase) { + return; } - if (missed != null) { + if (clock.instant().isAfter(nextHeartbeat)) { 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); + forceClose(MISSED_HEARTBEAT_CLOSE_REQUEST); + return; } + scheduleHeartbeatCheck(); } // region SessionStream event handlers private void dispatchResponseMessage(SessionResponse message) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); switch (message.getPayloadCase()) { case OPEN_SESSION: handleOpenSessionResponse(message.getOpenSession()); @@ -476,280 +515,281 @@ private void dispatchResponseMessage(SessionResponse message) { } 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) { - synchronized (lock) { - if (state.phase >= SessionState.CLOSING.phase) { - logger.fine( - String.format("Stream was already %s when session params were received", state)); - return; - } + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + 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))); } } 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); - VRpcImpl localRpc; - VRpcResult localCancel; - - 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 - localCancel = currentCancel; - currentCancel = null; - localRpc = currentRpc; - // TODO: handle multiplexing - currentRpc = 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; } - if (localCancel != null) { - tracer.onVRpcClose(localCancel.getStatus().getCode()); - localRpc.handleError(localCancel); + 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); } else { tracer.onVRpcClose(Status.OK.getCode()); - localRpc.handleResponse(vrpc); + rpc.handleResponse(vrpc); } - if (needsClose) { - synchronized (lock) { - if (state == SessionState.CLOSING) { - startGracefulClose(); - } - } + if (state == SessionState.CLOSING) { + startGracefulClose(); } } private void handleHeartBeatResponse(HeartbeatResponse ignored) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); this.nextHeartbeat = clock.instant().plus(heartbeatInterval); } private void handleSessionRefreshConfigResponse(SessionRefreshConfig config) { - 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; - } + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + 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) { + sessionSyncContext.throwIfNotInThisSynchronizationContext(); // Skips the heartbeat check when there's no active vrpc on the session this.nextHeartbeat = clock.instant().plus(FUTURE_TIME); - VRpcImpl localRpc; - boolean needsClose; - VRpcResult localCancel; - - 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( - 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 - localCancel = currentCancel; - currentCancel = null; - localRpc = currentRpc; - currentRpc = 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, error was: %s", + info.getLogName(), closeReason, error)); + return; } - if (localCancel != null) { - tracer.onVRpcClose(localCancel.getStatus().getCode()); - localRpc.handleError(localCancel); + 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 + 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()); + 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) { - if (state == SessionState.CLOSING) { - startGracefulClose(); - } - } + if (state == SessionState.CLOSING) { + startGracefulClose(); } } private void handleGoAwayResponse(GoAwayResponse goAwayResponse) { - 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; - } + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + 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); } 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) { - SessionState prevState; - VRpcImpl localVRpc; + sessionSyncContext.throwIfNotInThisSynchronizationContext(); + 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(); } - if (localVRpc != null) { + VRpcImpl localVRpc = currentRpc; + currentRpc = null; + updateState(SessionState.CLOSED); + + 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, @@ -759,13 +799,47 @@ 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); } - @GuardedBy("lock") private void updateState(SessionState newState) { this.state = newState; this.lastStateChangedAt = clock.instant(); 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) { 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..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 @@ -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()); @@ -666,6 +672,145 @@ 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 {