perf(dataplane): fold session FSM into the shared core (v5 S2a, no fast path)

rust/crates/dataplane/src/gate.rs

@@ -1,12 +1,15 @@
-//! Fast-path admission gate (v5, slice S1 — plumbing only, no fast path yet).
+//! Shared session core + fast-path admission gate (v5, slices S1–S2a — no fast
+//! path yet).
 //!
 //! Every message that can change how the *next* client command is handled —
 //! engine→loop [`SessionEvent`]s, loop/owner→engine [`EngineCmd`]s, and
 //! owner→loop [`SessionControl`]s — is registered in a shared counter before it
 //! is enqueued and deregistered only after it (and every effect it spawns) has
-//! been fully processed. A later slice reads `pending == 0 && open` to decide
-//! whether the engine may inline a steady transition; S1 only installs the
-//! counter and routes all three channels through it, with no behavior change.
+//! been fully processed (S1). S2a folds the session FSM into that same locked
+//! core, so a later slice can read `pending == 0 && open` together with a
+//! speculative transition to decide whether the engine may inline a steady
+//! transition. Through S2a the loop is still the FSM's only writer (via
+//! [`Gate::apply_slow`]) and there is no behavior change.
 //!
 //! Registration is a RAII [`GatePermit`] carried alongside the message in the
 //! channel (`mpsc::Sender<(T, GatePermit)>`). The permit's `Drop` decrements
@@ -18,24 +21,33 @@
 
 use std::sync::{Arc, Mutex};
 
+use session_core::fsm::{Effects, SessionEvent, SessionFsm, SessionState, TransitionError};
 use tokio::sync::mpsc;
 
-/// Shared gate state. The mutex guards only the small counters; no `.await`,
-/// send, or effect runs while it is held.
+/// Shared session core: the session FSM together with the admission counters,
+/// behind one std `Mutex` so a later slice can evaluate the fast-path predicate
+/// (`open && pending == 0` with a speculative transition) atomically. The mutex
+/// guards only pure computation — FSM transitions and counter arithmetic —
+/// never an `.await`, channel send, or effect execution.
 #[derive(Debug)]
-struct GateInner {
+struct FsmCore {
+    /// The session state machine. In S2a the loop is still its only writer (via
+    /// [`Gate::apply_slow`]); a later slice lets the engine commit whitelisted
+    /// steady transitions under this same lock.
+    fsm: SessionFsm,
     /// Registered-but-not-yet-processed messages across all producers/channels.
     pending: u32,
-    /// Loop is running and has not entered terminal handling. (S1: always true
-    /// and not yet read; the terminal gate consumes it in a later slice.)
+    /// Loop is running and has not entered terminal handling. Sealed to `false`
+    /// under the lock before the loop takes any terminal path; once false it
+    /// never reopens. (S2a maintains it; a later slice's predicate reads it.)
     #[allow(dead_code)]
     open: bool,
 }
 
-/// Handle to the shared admission gate for one session.
+/// Handle to the shared session core (FSM + admission gate) for one session.
 #[derive(Clone, Debug)]
 pub struct Gate {
-    inner: Arc<Mutex<GateInner>>,
+    inner: Arc<Mutex<FsmCore>>,
 }
 
 impl Default for Gate {
@@ -45,17 +57,46 @@ impl Default for Gate {
 }
 
 impl Gate {
-    /// Creates a gate with zero in-flight messages and `open == true`.
+    /// Creates a core with a fresh FSM (starting at `Accept`), zero in-flight
+    /// messages, and `open == true`.
     #[must_use]
     pub fn new() -> Self {
         Self {
-            inner: Arc::new(Mutex::new(GateInner {
+            inner: Arc::new(Mutex::new(FsmCore {
+                fsm: SessionFsm::new(),
                 pending: 0,
                 open: true,
             })),
         }
     }
 
+    /// Reads the current FSM state (a `Copy` value) under the lock and releases
+    /// immediately. Never exposes the guard.
+    pub(crate) fn state_snapshot(&self) -> SessionState {
+        lock(&self.inner).fsm.state()
+    }
+
+    /// Applies one event to the FSM on the slow path, under the lock, returning
+    /// the owned effects (or the rejection) together with the post-transition
+    /// state. The caller executes the effects *outside* the lock. The lock holds
+    /// only the pure `on_event` computation; no guard escapes.
+    pub(crate) fn apply_slow(
+        &self,
+        event: SessionEvent,
+    ) -> (Result<Effects, TransitionError>, SessionState) {
+        let mut inner = lock(&self.inner);
+        let result = inner.fsm.on_event(event);
+        let post_state = inner.fsm.state();
+        (result, post_state)
+    }
+
+    /// Seals the loop closed under the lock before it takes a terminal path.
+    /// Idempotent; once closed it never reopens. A later slice's fast-path
+    /// predicate reads `open` to refuse committing once the loop is terminating.
+    pub(crate) fn seal_closed(&self) {
+        lock(&self.inner).open = false;
+    }
+
     /// Registers a standalone "operation" permit not tied to a single message.
     /// Used to bridge a multi-message atomic operation (e.g. the redirect pair
     /// `PrepareRedirect` + `Redirect`): held across both sends so `pending`
@@ -102,7 +143,7 @@ impl Gate {
 #[derive(Debug)]
 #[must_use = "a GatePermit must travel with its message; dropping it early frees the gate slot"]
 pub struct GatePermit {
-    inner: Arc<Mutex<GateInner>>,
+    inner: Arc<Mutex<FsmCore>>,
 }
 
 impl Drop for GatePermit {
@@ -211,6 +252,54 @@ mod tests {
 
     type R = Result<(), Box<dyn std::error::Error>>;
 
+    #[test]
+    fn apply_slow_and_state_snapshot_drive_the_fsm_like_a_standalone() {
+        // The relocated FSM, driven through the core, must transition and emit
+        // effects identically to a standalone `SessionFsm` on the same events.
+        let core = Gate::new();
+        let mut reference = SessionFsm::new();
+        assert_eq!(core.state_snapshot(), reference.state());
+
+        let script = [
+            SessionEvent::ConnectionAccepted,
+            SessionEvent::ClientHandshakeResponse,
+            SessionEvent::BackendGreetingReceived,
+            SessionEvent::BackendAuthOk,
+            SessionEvent::ClientCommand,
+            SessionEvent::BackendResponseTxnDone,
+        ];
+        for event in script {
+            let (result, post_state) = core.apply_slow(event);
+            let expected = reference.on_event(event);
+            assert_eq!(
+                result, expected,
+                "core effects/rejection match the standalone FSM for {event:?}"
+            );
+            assert_eq!(
+                post_state,
+                reference.state(),
+                "core post-state matches the standalone FSM for {event:?}"
+            );
+            assert_eq!(
+                core.state_snapshot(),
+                reference.state(),
+                "a later snapshot still matches for {event:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn seal_closed_is_sticky_and_idempotent() {
+        let core = Gate::new();
+        assert!(core.is_open(), "a fresh core is open");
+        core.seal_closed();
+        assert!(!core.is_open(), "seal_closed closes the core");
+        // The loop may reach several terminal checks; sealing again is a no-op
+        // and never reopens.
+        core.seal_closed();
+        assert!(!core.is_open(), "once closed the core never reopens");
+    }
+
     #[tokio::test]
     async fn pending_visible_before_apply_and_cleared_after_drop() -> R {
         let gate = Gate::new();

rust/crates/dataplane/src/session.rs

@@ -61,7 +61,7 @@
 use std::pin::Pin;
 use std::time::Duration;
 
-use session_core::fsm::{SessionEffect, SessionEvent, SessionFsm, SessionState, TransitionError};
+use session_core::fsm::{SessionEffect, SessionEvent, SessionState, TransitionError};
 use tokio::sync::{mpsc, oneshot, watch};
 use tokio::task::{JoinHandle, JoinSet};
 use tokio::time::{Instant, Sleep, sleep_until, timeout_at};
@@ -248,7 +248,11 @@ const EVENT_PUMP_CAPACITY: usize = 1;
 
 /// The single owner of one session's mutable state.
 pub struct SessionLoop<S, E> {
-    fsm: SessionFsm,
+    /// The shared session core (FSM + admission gate). The loop is the FSM's
+    /// only writer on the slow path, via [`crate::gate::Gate::apply_slow`]; a
+    /// later slice lets the engine commit whitelisted steady transitions under
+    /// the same lock.
+    core: crate::gate::Gate,
     source: Option<S>,
     handler: E,
     control: crate::gate::GatedReceiver<SessionControl>,
@@ -322,17 +326,20 @@ impl Drop for AbortOnDrop {
 }
 
 impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
-    /// Creates a session loop; the FSM starts at `Accept`.
+    /// Creates a session loop over the shared `core` (which carries the FSM,
+    /// starting at `Accept`). The loop is the FSM's only writer on the slow
+    /// path.
     #[must_use]
     pub fn new(
+        core: crate::gate::Gate,
         source: S,
         handler: E,
         control: crate::gate::GatedReceiver<SessionControl>,
         shutdown: watch::Receiver<bool>,
         config: SessionLoopConfig,
     ) -> Self {
         Self {
-            fsm: SessionFsm::new(),
+            core,
             source: Some(source),
             handler,
             control,
@@ -434,14 +441,23 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
         let cleanup = self.cleanup_within(cleanup_by).await;
         SessionSummary {
             end,
-            final_state: self.fsm.state(),
+            final_state: self.core.state_snapshot(),
             effects_executed: self.effects_executed,
             rejected_events: self.rejected_events,
             control_detached: self.control_detached,
             cleanup,
         }
     }
 
+    /// Terminal fence: seals the core closed — so a later slice's engine fast
+    /// path refuses to commit once the loop is terminating — then returns the
+    /// loop's end reason. Called at every path that leaves the steady loop,
+    /// before the FSM is still `Ready` but the session is tearing down.
+    fn seal_and_end(&self, end: LoopEnd) -> LoopEnd {
+        self.core.seal_closed();
+        end
+    }
+
     async fn event_loop(&mut self, events: &mut LoopEventRx) -> LoopEnd {
         let handshake_deadline = Instant::now() + self.config.handshake_deadline;
         let mut armed_deadline: Option<(Instant, SessionEvent)> =
@@ -480,10 +496,10 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
         // redirect owns the command boundary (Go's processLock).
         let mut queued_command = None;
         loop {
-            if self.fsm.state() == SessionState::Closed {
-                return LoopEnd::FsmClosed;
+            if self.core.state_snapshot() == SessionState::Closed {
+                return self.seal_and_end(LoopEnd::FsmClosed);
             }
-            if self.fsm.state() == SessionState::Closing {
+            if self.core.state_snapshot() == SessionState::Closing {
                 // The loop is the runtime: teardown effects have executed
                 // and their children are tracked, so seal the FSM now;
                 // the children drain in the terminal cleanup, under the
@@ -515,8 +531,8 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
                 )
                 .await;
             match action {
-                LoopAction::ServerShutdown => return LoopEnd::Shutdown,
-                LoopAction::SourceExhausted => return LoopEnd::SourceExhausted,
+                LoopAction::ServerShutdown => return self.seal_and_end(LoopEnd::Shutdown),
+                LoopAction::SourceExhausted => return self.seal_and_end(LoopEnd::SourceExhausted),
                 LoopAction::ControlDetached => {
                     // Control-v1 last-good: losing the control channel never
                     // tears down an established session. Redirect/drain
@@ -528,7 +544,14 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
                     if let Some(probe) = next_probe.as_mut() {
                         *probe = Instant::now() + self.config.backend_check_interval;
                     }
-                    if probe_safe(self.fsm.state()) && !self.handler.backend_active().await {
+                    // Snapshot the state out of the lock *before* the await: a
+                    // `self.core.state_snapshot()` sub-expression would otherwise
+                    // keep the guard alive across `backend_active().await`.
+                    // S2b: this read must be re-snapshotted after the await
+                    // before acting, once the engine can commit concurrently —
+                    // it is not commit authority.
+                    let probe_state = self.core.state_snapshot();
+                    if probe_safe(probe_state) && !self.handler.backend_active().await {
                         self.apply(SessionEvent::BackendIoError, &mut armed_deadline, None)
                             .await;
                     }
@@ -538,7 +561,7 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
                     self.apply(event, &mut armed_deadline, None).await;
                 }
                 LoopAction::Event(event, permit) => {
-                    if self.fsm.state() == SessionState::RedirectPending
+                    if self.core.state_snapshot() == SessionState::RedirectPending
                         && matches!(
                             event,
                             SessionEvent::ClientCommand | SessionEvent::ClientCommandQuit
@@ -567,7 +590,7 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
                     // (any effect it spawned already holds its own cmd permit).
                 }
             }
-            if self.fsm.state() == SessionState::Ready
+            if self.core.state_snapshot() == SessionState::Ready
                 && let Some((event, permit)) = queued_command.take()
             {
                 // apply() enqueues all terminal redirect effects first, so the
@@ -639,7 +662,11 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
         armed_deadline: &mut Option<(Instant, SessionEvent)>,
         drain_deadline: Option<Duration>,
     ) {
-        match self.fsm.on_event(event) {
+        // The transition runs under the core lock and returns owned effects plus
+        // the post-transition state; the effects then execute *outside* the lock
+        // (the std mutex guards only pure computation).
+        let (result, post_state) = self.core.apply_slow(event);
+        match result {
             Ok(effects) => {
                 for effect in effects {
                     if effect == SessionEffect::BeginDrainTimer {
@@ -654,10 +681,12 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
                 // The handshake deadline is judged on the post-transition
                 // state: the very transition into an authenticated state
                 // (`BackendAuthOk`) disarms it, with no dependency on any
-                // later event arriving.
+                // later event arriving. The post-state captured with the
+                // transition is authoritative here (the loop is the only
+                // writer of these non-steady transitions).
                 if let Some((_, pending)) = *armed_deadline
                     && pending == SessionEvent::HandshakeTimerExpired
-                    && authenticated_phase(self.fsm.state())
+                    && authenticated_phase(post_state)
                 {
                     *armed_deadline = None;
                 }
@@ -674,11 +703,13 @@ impl<S: SessionEventSource, E: EffectHandler> SessionLoop<S, E> {
     /// the machine. Rejections are tolerated (the FSM may already be
     /// closing or closed).
     async fn close_via_fsm(&mut self, close_event: SessionEvent) {
+        // Reached only from the terminal path in `run`, after the core is
+        // sealed closed; the loop is the sole writer here.
         let mut deadline = None;
-        if self.fsm.state() != SessionState::Closed {
+        if self.core.state_snapshot() != SessionState::Closed {
             self.apply(close_event, &mut deadline, None).await;
         }
-        if self.fsm.state() != SessionState::Closed {
+        if self.core.state_snapshot() != SessionState::Closed {
             self.apply(SessionEvent::TeardownComplete, &mut deadline, None)
                 .await;
         }

rust/crates/dataplane/src/session_engine.rs

@@ -840,6 +840,7 @@ async fn run_bound_session_observed(
     // fires — not only when the loop finishes its own cleanup.
     let mut owner_shutdown = shutdown.clone();
     let session_loop = SessionLoop::new(
+        gate.clone(),
         EventRx { events: event_rx },
         CmdTx {
             cmds: cmd_tx.clone(),