diff --git a/rust/crates/dataplane/src/gate.rs b/rust/crates/dataplane/src/gate.rs index 7069af34a..ffb4ffdb1 100644 --- a/rust/crates/dataplane/src/gate.rs +++ b/rust/crates/dataplane/src/gate.rs @@ -21,8 +21,10 @@ use std::sync::{Arc, Mutex}; -use session_core::fsm::{Effects, SessionEvent, SessionFsm, SessionState, TransitionError}; -use tokio::sync::mpsc; +use session_core::fsm::{ + Effects, SessionEffect, SessionEvent, SessionFsm, SessionState, TransitionError, +}; +use tokio::sync::{mpsc, watch}; /// 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 @@ -38,10 +40,15 @@ struct FsmCore { /// Registered-but-not-yet-processed messages across all producers/channels. pending: u32, /// 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)] + /// under the lock before the loop takes any terminal path or when the + /// shutdown relay observes shutdown; once false it never reopens. It is the + /// single admission authority: both [`Gate::apply_slow`] (gated loop events) + /// and [`Gate::try_commit_steady`] (engine fast path) read it under this + /// lock, so they linearize on the same seal. open: bool, + /// Count of transitions the engine committed inline on the fast path + /// (diagnostic; read by tests and metrics). + fast_transitions: u64, } /// Handle to the shared session core (FSM + admission gate) for one session. @@ -66,6 +73,7 @@ impl Gate { fsm: SessionFsm::new(), pending: 0, open: true, + fast_transitions: 0, })), } } @@ -76,27 +84,120 @@ impl Gate { 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, SessionState) { + /// Applies one event to the FSM on the slow path, under the lock. + /// + /// When `gated` is `true` the event is an ordinary data-path event and is + /// refused ([`SlowOutcome::DroppedClosed`]) if the core is already sealed + /// closed — this is the admission check that, atomically with the + /// transition, makes a post-seal `ClientCommand` behave the same whether it + /// arrives on the slow path or is refused by the fast path. Terminal + /// teardown events pass `gated = false` so sealing never blocks the FSM + /// from reaching `Closed`. On [`SlowOutcome::Applied`] the caller executes + /// the returned owned effects *outside* the lock; no guard escapes. + pub(crate) fn apply_slow(&self, event: SessionEvent, gated: bool) -> SlowOutcome { let mut inner = lock(&self.inner); + if gated && !inner.open { + return SlowOutcome::DroppedClosed; + } let result = inner.fsm.on_event(event); let post_state = inner.fsm.state(); - (result, post_state) + SlowOutcome::Applied { 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. + /// Seals the loop closed under the lock before it takes a terminal path, or + /// when the shutdown relay observes shutdown. Idempotent; once closed it + /// never reopens. The admission paths ([`Gate::apply_slow`] gated, + /// [`Gate::try_commit_steady`]) read `open` under this same lock. pub(crate) fn seal_closed(&self) { lock(&self.inner).open = false; } + /// Whether the core is still open (loop running, not terminating). The + /// loop reads this as its per-iteration shutdown fence; the admission paths + /// ([`Gate::apply_slow`], [`Gate::try_commit_steady`]) also read `open` + /// inline under the lock. + pub(crate) fn is_open(&self) -> bool { + lock(&self.inner).open + } + + /// Engine fast path: attempt to commit one steady-state transition inline, + /// skipping the `events`→loop→`cmd` round-trip. Under the single core lock + /// this checks the whole admission predicate atomically: + /// + /// * the core is `open` (not sealed by the relay or a terminal path), + /// * shutdown is not observed (sender-closed counts as shutdown), + /// * `pending == 0` (no redirect/drain/control in flight), + /// * `event` is on the whitelist (only [`SessionEvent::ClientCommand`]), + /// * and a speculative transition on a *clone* yields exactly the single + /// expected steady effect. + /// + /// Only then is the clone committed back, an in-flight [`GatePermit`] + /// registered (so `pending` holds the floor while the caller performs the + /// forward, exactly as the slow path's effect permit does), and + /// `fast_transitions` bumped. Otherwise it is [`FastPath::Fallback`] with + /// zero writes to the core. The returned permit is constructed under the + /// lock but only dropped by the caller *outside* it. + pub(crate) fn try_commit_steady( + &self, + event: SessionEvent, + expected_effect: SessionEffect, + shutdown: &watch::Receiver, + ) -> FastPath { + let mut inner = lock(&self.inner); + let closed = shutdown.has_changed().map_or(true, |_| *shutdown.borrow()); + if !inner.open + || closed + || inner.pending != 0 + || event != SessionEvent::ClientCommand + || inner.fsm.state() != SessionState::Ready + { + return FastPath::Fallback; + } + let mut speculative = inner.fsm.clone(); + match speculative.on_event(event) { + Ok(effects) if effects.len() == 1 && effects[0] == expected_effect => { + inner.fsm = speculative; + inner.pending += 1; + inner.fast_transitions += 1; + // Registered above; its `Drop` will decrement `pending`. Built + // under the lock but handed back to the caller, who drops it + // only after the forward completes and never under this mutex. + FastPath::Committed(GatePermit { + inner: Arc::clone(&self.inner), + }) + } + _ => FastPath::Fallback, + } + } + + /// Admits a quiescent loop-side operation (currently the backend probe) + /// atomically with the core: only when the core is `open`, `pending == 0`, + /// and `state_ok` holds for the current FSM state does it register an + /// operation [`GatePermit`] and return it. Holding that permit across the + /// operation's `.await` keeps `pending >= 1`, so the engine fast path falls + /// back for the duration — the operation cannot race a concurrent inline + /// commit (e.g. a probe's stale `inactive` result tearing down a command + /// the fast path committed mid-probe). Returns `None` (skip) otherwise. + pub(crate) fn try_quiescent_operation(&self, state_ok: F) -> Option + where + F: FnOnce(SessionState) -> bool, + { + let mut inner = lock(&self.inner); + if !inner.open || inner.pending != 0 || !state_ok(inner.fsm.state()) { + return None; + } + inner.pending += 1; + Some(GatePermit { + inner: Arc::clone(&self.inner), + }) + } + + /// Count of inline fast-path commits. Test/diagnostic only. + #[cfg(test)] + pub(crate) fn fast_transitions(&self) -> u64 { + lock(&self.inner).fast_transitions + } + /// 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` @@ -125,14 +226,30 @@ impl Gate { pub(crate) fn pending(&self) -> u32 { lock(&self.inner).pending } +} - /// Whether the loop is still open. Test/diagnostic only until the terminal - /// gate consumes it. - #[cfg(test)] - #[allow(dead_code)] - pub(crate) fn is_open(&self) -> bool { - lock(&self.inner).open - } +/// Outcome of a gated slow-path transition ([`Gate::apply_slow`]). +pub(crate) enum SlowOutcome { + /// The transition ran; `result` carries its effects or rejection and + /// `post_state` the FSM state afterwards. + Applied { + result: Result, + post_state: SessionState, + }, + /// The core was sealed closed and the event was gated, so it was dropped + /// with no transition — the shutdown-admission refusal. + DroppedClosed, +} + +/// Result of an engine fast-path attempt ([`Gate::try_commit_steady`]). +#[must_use = "a Committed fast path carries the in-flight permit that must be held across the forward"] +pub(crate) enum FastPath { + /// The steady transition was committed inline; the caller skips the + /// `events`+`await_effect` round-trip and holds the permit across the + /// forward, dropping it only after the command's effect is fully handled. + Committed(GatePermit), + /// The predicate did not hold; the caller runs the ordinary slow path. + Fallback, } /// RAII registration for one in-flight message. Decrements the gate counter on @@ -253,7 +370,7 @@ mod tests { type R = Result<(), Box>; #[test] - fn apply_slow_and_state_snapshot_drive_the_fsm_like_a_standalone() { + fn apply_slow_and_state_snapshot_drive_the_fsm_like_a_standalone() -> R { // 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(); @@ -269,7 +386,9 @@ mod tests { SessionEvent::BackendResponseTxnDone, ]; for event in script { - let (result, post_state) = core.apply_slow(event); + let SlowOutcome::Applied { result, post_state } = core.apply_slow(event, true) else { + return Err(format!("open core must apply gated {event:?}, not drop it").into()); + }; let expected = reference.on_event(event); assert_eq!( result, expected, @@ -286,6 +405,29 @@ mod tests { "a later snapshot still matches for {event:?}" ); } + Ok(()) + } + + #[test] + fn gated_apply_is_dropped_once_closed_but_terminal_still_applies() { + let core = Gate::new(); + core.seal_closed(); + // A gated (ordinary) event is refused after the seal... + assert!( + matches!( + core.apply_slow(SessionEvent::ClientCommand, true), + SlowOutcome::DroppedClosed + ), + "a gated event is dropped once the core is sealed closed" + ); + // ...but a terminal (ungated) event still drives the FSM. + assert!( + matches!( + core.apply_slow(SessionEvent::ConnectionAccepted, false), + SlowOutcome::Applied { .. } + ), + "an ungated terminal event still applies after the seal" + ); } #[test] @@ -300,6 +442,235 @@ mod tests { assert!(!core.is_open(), "once closed the core never reopens"); } + /// Drives a fresh core through the handshake to `Ready` (the steady state + /// where the fast path is eligible). + fn core_in_ready() -> Gate { + let core = Gate::new(); + for event in [ + SessionEvent::ConnectionAccepted, + SessionEvent::ClientHandshakeResponse, + SessionEvent::BackendGreetingReceived, + SessionEvent::BackendAuthOk, + ] { + assert!( + matches!(core.apply_slow(event, true), SlowOutcome::Applied { .. }), + "handshake event {event:?} should apply" + ); + } + assert_eq!(core.state_snapshot(), SessionState::Ready); + core + } + + #[test] + fn try_commit_steady_commits_and_matches_the_slow_transition() -> R { + let (_tx, shutdown) = watch::channel(false); + let core = core_in_ready(); + + // Reference: what the slow path would do for the same command. + let reference = core_in_ready(); + let SlowOutcome::Applied { result, post_state } = + reference.apply_slow(SessionEvent::ClientCommand, true) + else { + return Err("slow path must apply a Ready ClientCommand".into()); + }; + assert_eq!( + result, + Ok(vec![SessionEffect::ForwardCommandToBackend]), + "the slow transition emits exactly the expected steady effect" + ); + + // Fast path: commits the same transition inline, registers one in-flight + // permit, bumps the fast counter, and lands on the same FSM state. + let FastPath::Committed(permit) = core.try_commit_steady( + SessionEvent::ClientCommand, + SessionEffect::ForwardCommandToBackend, + &shutdown, + ) else { + return Err("a quiescent Ready ClientCommand must commit on the fast path".into()); + }; + assert_eq!(core.fast_transitions(), 1); + assert_eq!(core.pending(), 1, "the in-flight guard holds the floor"); + assert_eq!( + core.state_snapshot(), + post_state, + "fast commit lands on the same FSM state as the slow transition" + ); + drop(permit); + assert_eq!(core.pending(), 0, "dropping the permit frees the floor"); + Ok(()) + } + + #[test] + fn try_commit_steady_falls_back_without_touching_the_core() { + let assert_fallback = |core: &Gate, shutdown: &watch::Receiver, event, label| { + let state_before = core.state_snapshot(); + let fast_before = core.fast_transitions(); + let pending_before = core.pending(); + assert!( + matches!( + core.try_commit_steady(event, SessionEffect::ForwardCommandToBackend, shutdown), + FastPath::Fallback + ), + "{label}: must fall back" + ); + assert_eq!( + core.state_snapshot(), + state_before, + "{label}: FSM untouched" + ); + assert_eq!( + core.fast_transitions(), + fast_before, + "{label}: no fast count" + ); + assert_eq!( + core.pending(), + pending_before, + "{label}: no permit registered" + ); + }; + + let (_open_tx, open) = watch::channel(false); + + // Shutdown observed (value true). + let (_tx, closed) = watch::channel(true); + assert_fallback( + &core_in_ready(), + &closed, + SessionEvent::ClientCommand, + "shutdown-true", + ); + + // Shutdown sender dropped (has_changed -> Err -> treated as closed). + let dropped = { + let (tx, rx) = watch::channel(false); + drop(tx); + rx + }; + assert_fallback( + &core_in_ready(), + &dropped, + SessionEvent::ClientCommand, + "sender-closed", + ); + + // Core sealed closed. + let sealed = core_in_ready(); + sealed.seal_closed(); + assert_fallback(&sealed, &open, SessionEvent::ClientCommand, "sealed"); + + // Pending non-zero (an operation permit in flight). + let busy = core_in_ready(); + let op = busy.operation_permit(); + assert_fallback(&busy, &open, SessionEvent::ClientCommand, "pending!=0"); + drop(op); + + // Non-whitelisted event (Quit). + assert_fallback( + &core_in_ready(), + &open, + SessionEvent::ClientCommandQuit, + "quit", + ); + + // FSM not in Ready (fresh core in Accept). + assert_fallback( + &Gate::new(), + &open, + SessionEvent::ClientCommand, + "not-ready", + ); + } + + #[test] + fn try_commit_steady_rejects_an_effect_mismatch() { + // Defensive exact-effect guard: even in Ready with a quiescent gate, a + // commit is refused unless the speculative transition yields exactly the + // single expected effect. Passing a different expected effect models a + // caller/FSM drift and must fall back with zero writes. + let (_tx, shutdown) = watch::channel(false); + let core = core_in_ready(); + assert!( + matches!( + core.try_commit_steady( + SessionEvent::ClientCommand, + SessionEffect::ForwardResponseToClient, + &shutdown, + ), + FastPath::Fallback + ), + "a non-matching expected effect must fall back" + ); + assert_eq!(core.fast_transitions(), 0); + assert_eq!(core.pending(), 0); + assert_eq!(core.state_snapshot(), SessionState::Ready); + } + + #[test] + fn probe_permit_is_quiescent_and_blocks_the_fast_path_across_its_await() { + let (_tx, shutdown) = watch::channel(false); + let only_ready = |state| state == SessionState::Ready; + + // Not probe-safe: a fresh core in Accept yields no permit. + assert!( + Gate::new().try_quiescent_operation(only_ready).is_none(), + "a non-probe-safe state is not admitted" + ); + + // Sealed: no probe once terminating. + let sealed = core_in_ready(); + sealed.seal_closed(); + assert!( + sealed.try_quiescent_operation(only_ready).is_none(), + "a sealed core admits no probe" + ); + + // Already busy: pending != 0 yields no permit. + let busy = core_in_ready(); + let op = busy.operation_permit(); + assert!( + busy.try_quiescent_operation(only_ready).is_none(), + "a non-quiescent core (pending != 0) admits no probe" + ); + drop(op); + + // Quiescent: a permit is granted, and WHILE it is held (modeling the + // probe's `backend_active().await`) the engine fast path falls back — + // so no concurrent inline commit can turn Ready -> Command mid-probe. + let core = core_in_ready(); + let Some(probe_permit) = core.try_quiescent_operation(only_ready) else { + unreachable!("a quiescent probe-safe core admits the probe"); + }; + assert_eq!(core.pending(), 1, "the probe permit holds the floor"); + assert!( + matches!( + core.try_commit_steady( + SessionEvent::ClientCommand, + SessionEffect::ForwardCommandToBackend, + &shutdown, + ), + FastPath::Fallback + ), + "the fast path falls back while the probe permit is held" + ); + assert_eq!( + core.fast_transitions(), + 0, + "no commit happened during the probe" + ); + assert_eq!( + core.state_snapshot(), + SessionState::Ready, + "FSM untouched during the probe" + ); + drop(probe_permit); + assert_eq!( + core.pending(), + 0, + "dropping the probe permit frees the floor" + ); + } + #[tokio::test] async fn pending_visible_before_apply_and_cleared_after_drop() -> R { let gate = Gate::new(); diff --git a/rust/crates/dataplane/src/session.rs b/rust/crates/dataplane/src/session.rs index 521c77c35..762335f81 100644 --- a/rust/crates/dataplane/src/session.rs +++ b/rust/crates/dataplane/src/session.rs @@ -304,15 +304,45 @@ impl LoopEventRx { /// Parks on the shared shutdown `watch` once and fires `fired` when the /// signal becomes `true` or its sender is dropped. Ignores changes back to /// `false`, which the loop treated as a no-op wake. -async fn shutdown_relay(mut shutdown: watch::Receiver, fired: oneshot::Sender<()>) { +async fn shutdown_relay( + mut shutdown: watch::Receiver, + core: crate::gate::Gate, + started: oneshot::Sender<()>, + fired: oneshot::Sender<()>, +) { + // First observation, made before the loop is allowed to touch input: a + // shutdown that already holds (value `true`) or a sender already dropped is + // sealed here, and `started` acks that the relay has observed the initial + // state. The loop awaits `started` before its first select, so a + // pre-existing shutdown is sealed strictly before any event can be + // dispatched — restoring the synchronous precheck guarantee without the + // loop reading the watch directly. + let down_at_start = shutdown + .has_changed() + .map_or(true, |_| *shutdown.borrow_and_update()); + if down_at_start { + core.seal_closed(); + } + let _ = started.send(()); + if down_at_start { + let _ = fired.send(()); + return; + } + // Otherwise park until the signal flips to `true` or the sender drops. loop { - if *shutdown.borrow_and_update() { + if shutdown.changed().await.is_err() { break; } - if shutdown.changed().await.is_err() { + if *shutdown.borrow_and_update() { break; } } + // Seal the core *before* waking the loop, so `open == false` is published + // under the core lock strictly before the loop can act on the oneshot. This + // makes the relay the single shutdown-admission linearization point: the + // engine's fast path (reads `open`) and the loop's gated apply (reads + // `open`) both refuse once this seal lands, regardless of oneshot latency. + core.seal_closed(); let _ = fired.send(()); } @@ -458,6 +488,9 @@ impl SessionLoop { end } + // One cohesive select/dispatch loop; splitting the arms out would scatter + // the shared deadline/probe/queued-command state and obscure the ordering. + #[allow(clippy::too_many_lines)] 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)> = @@ -479,17 +512,33 @@ impl SessionLoop { // session. Awaiting `changed()` inside the select registered this // task on the channel's waiter list and unlinked it again on every // iteration, serializing all sessions on that list's mutex. Instead a - // per-session relay task parks on the shared channel once and fires a - // private oneshot; the loop polls the oneshot (no shared state) and - // keeps an exact lock-free precheck via `has_changed`. - // `mark_changed` makes the first precheck read the current value, so - // a shutdown that predates the loop is still seen before any select. - self.shutdown.mark_changed(); + // per-session relay task parks on the shared channel once, seals the + // core, and fires a private oneshot; the loop polls the oneshot (no + // shared state). + // + // S2b: the relay is the loop's SOLE shutdown-observation path — the + // loop no longer runs its own `has_changed` precheck. The relay seals + // the core (under its lock) before firing, so `open == false` is the + // single admission authority and the engine's fast path cannot commit a + // command that this loop would have dropped at shutdown. + // + // The relay acks via `started` once it has made its first observation + // (and sealed, if a shutdown already held); the loop awaits that ack + // before touching any input, so a pre-existing shutdown is sealed + // strictly before the first select — the loop's `open` fence below then + // refuses the first event. This restores the synchronous-precheck + // guarantee the direct `has_changed` read used to provide. + let (started_tx, started_rx) = oneshot::channel::<()>(); let (relay_tx, mut shutdown_fired) = oneshot::channel::<()>(); let _relay = AbortOnDrop(tokio::spawn(shutdown_relay( self.shutdown.clone(), + self.core.clone(), + started_tx, relay_tx, ))); + // The relay cannot outlive this await without having observed the + // initial state; a dropped ack (relay aborted) is treated as observed. + let _ = started_rx.await; // The wire engine has at most one unacknowledged command. Its payload // stays in the engine; retain only the classified event while a @@ -505,8 +554,13 @@ impl SessionLoop { // the children drain in the terminal cleanup, under the // single absolute budget, after the pump releases the // transport. - self.apply(SessionEvent::TeardownComplete, &mut armed_deadline, None) - .await; + self.apply( + SessionEvent::TeardownComplete, + &mut armed_deadline, + None, + false, + ) + .await; continue; } // Sync the persistent timers to the current deadlines; reset only @@ -530,6 +584,20 @@ impl SessionLoop { next_probe, ) .await; + // Shutdown fence: the relay may have sealed the core during the + // select above (its `seal_closed` is published before the oneshot + // it also fires). Once sealed, refuse to dispatch *any* action — + // this closes the seal->fire window where a `gated = false` branch + // (Deadline / BackendProbe / Control) could otherwise transition + // before the `ServerShutdown` oneshot is observed. The dequeued + // action is dropped here (its permit, if any, is returned); + // terminal close in `run` is reached through this `Shutdown` end and + // runs ungated. Only the relay seals `open` mid-loop (`seal_and_end` + // seals and returns in the same step), so `!open` here means + // shutdown. + if !self.core.is_open() { + return LoopEnd::Shutdown; + } match action { LoopAction::ServerShutdown => return self.seal_and_end(LoopEnd::Shutdown), LoopAction::SourceExhausted => return self.seal_and_end(LoopEnd::SourceExhausted), @@ -544,21 +612,31 @@ impl SessionLoop { if let Some(probe) = next_probe.as_mut() { *probe = Instant::now() + self.config.backend_check_interval; } - // 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; + // Admit the probe atomically with the core: only probe when + // it is quiescent (open, probe-safe state, `pending == 0`), + // and hold the returned operation permit across the await. + // While it is held `pending >= 1`, so the engine fast path + // falls back and cannot turn `Ready -> Command` mid-probe — + // a stale `inactive` result can no longer tear down a command + // the fast path committed during the probe. If the core is + // not quiescent the probe is skipped this round. The permit + // binding stays in scope across the `&&` await and the body, + // so it is held for the whole probe and dropped afterwards. + if let Some(_probe_permit) = self.core.try_quiescent_operation(probe_safe) + && !self.handler.backend_active().await + { + self.apply( + SessionEvent::BackendIoError, + &mut armed_deadline, + None, + false, + ) + .await; } } LoopAction::Deadline(event) => { armed_deadline = None; - self.apply(event, &mut armed_deadline, None).await; + self.apply(event, &mut armed_deadline, None, false).await; } LoopAction::Event(event, permit) => { if self.core.state_snapshot() == SessionState::RedirectPending @@ -575,7 +653,10 @@ impl SessionLoop { // stays non-empty until it is finally applied. queued_command = Some((event, permit)); } else { - self.apply(event, &mut armed_deadline, None).await; + // Ordinary engine event: gated, so a post-seal command + // is dropped under the lock (shutdown-admission refusal), + // matching the fast path refusing to commit once sealed. + self.apply(event, &mut armed_deadline, None, true).await; drop(permit); } } @@ -584,8 +665,13 @@ impl SessionLoop { SessionControl::GracefulCloseAfter(deadline) => Some(deadline), _ => None, }; - self.apply(command.session_event(), &mut armed_deadline, drain_deadline) - .await; + self.apply( + command.session_event(), + &mut armed_deadline, + drain_deadline, + false, + ) + .await; // Permit drops here, after the control's effects are applied // (any effect it spawned already holds its own cmd permit). } @@ -596,7 +682,7 @@ impl SessionLoop { // apply() enqueues all terminal redirect effects first, so the // backend swap (or failure retaining the old backend) precedes // this command's forwarding ACK. Closing instead drops the slot. - self.apply(event, &mut armed_deadline, None).await; + self.apply(event, &mut armed_deadline, None, true).await; drop(permit); } } @@ -611,19 +697,12 @@ impl SessionLoop { armed_deadline: Option<(Instant, SessionEvent)>, next_probe: Option, ) -> LoopAction { - // A shutdown that predates this call (including one set before the - // loop started) must not be lost to the relay's scheduling. The - // version check is a lock-free atomic load; the value is read only - // when the version moved. A dropped sender is a shutdown, as before. - match self.shutdown.has_changed() { - Err(_) => return LoopAction::ServerShutdown, - Ok(true) => { - if *self.shutdown.borrow_and_update() { - return LoopAction::ServerShutdown; - } - } - Ok(false) => {} - } + // S2b: no direct `has_changed` precheck here — the relay is the sole + // shutdown-observation path (it seals the core and fires `shutdown_fired` + // below). A shutdown that predates this call reaches us through the + // oneshot (the relay's first `borrow_and_update` sees it), and a + // command that races the seal is refused by the gated `apply` under the + // core lock. // opt#6: poll the caller's persistent timers instead of building a // fresh `sleep_until` here. A new timer per call armed and (on the // untaken branch) dropped a tokio timer-wheel entry on every packet, @@ -661,11 +740,18 @@ impl SessionLoop { event: SessionEvent, armed_deadline: &mut Option<(Instant, SessionEvent)>, drain_deadline: Option, + gated: bool, ) { // 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); + // (the std mutex guards only pure computation). A gated ordinary event + // that arrives after the core is sealed closed is dropped here — the + // shutdown-admission refusal that mirrors the fast path refusing to + // commit once sealed. + let (result, post_state) = match self.core.apply_slow(event, gated) { + crate::gate::SlowOutcome::Applied { result, post_state } => (result, post_state), + crate::gate::SlowOutcome::DroppedClosed => return, + }; match result { Ok(effects) => { for effect in effects { @@ -706,11 +792,13 @@ impl SessionLoop { // 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; + // Terminal events are ungated: the core is already sealed closed, so + // gating would drop them and the FSM could never reach `Closed`. if self.core.state_snapshot() != SessionState::Closed { - self.apply(close_event, &mut deadline, None).await; + self.apply(close_event, &mut deadline, None, false).await; } if self.core.state_snapshot() != SessionState::Closed { - self.apply(SessionEvent::TeardownComplete, &mut deadline, None) + self.apply(SessionEvent::TeardownComplete, &mut deadline, None, false) .await; } } @@ -776,3 +864,201 @@ enum LoopEnd { Shutdown, SourceExhausted, } + +#[cfg(test)] +mod relay_tests { + use super::shutdown_relay; + use crate::gate::Gate; + use tokio::sync::{oneshot, watch}; + + /// A shutdown that already holds at startup is sealed and acked before the + /// loop is released: `started` fires, the core is sealed, and the loop + /// oneshot fires — all before the relay returns. The seal is published + /// strictly before the wake, the single shutdown-admission linearization + /// point. + #[tokio::test] + async fn relay_seals_and_acks_a_preexisting_true_signal() { + let (_tx, rx) = watch::channel(true); + let core = Gate::new(); + assert!(core.is_open()); + let (started_tx, started_rx) = oneshot::channel::<()>(); + let (fired_tx, fired_rx) = oneshot::channel::<()>(); + shutdown_relay(rx, core.clone(), started_tx, fired_tx).await; + assert!( + started_rx.await.is_ok(), + "relay acked its first observation" + ); + assert!( + !core.is_open(), + "a pre-existing shutdown is sealed at startup" + ); + assert!(fired_rx.await.is_ok(), "relay fired the loop oneshot"); + } + + /// A sender already dropped at startup is a shutdown too, sealed and acked + /// before the loop runs (so a torn-down runtime never leaves the core + /// admitting). + #[tokio::test] + async fn relay_seals_and_acks_a_preclosed_sender() { + let (tx, rx) = watch::channel(false); + drop(tx); + let core = Gate::new(); + let (started_tx, started_rx) = oneshot::channel::<()>(); + let (fired_tx, fired_rx) = oneshot::channel::<()>(); + shutdown_relay(rx, core.clone(), started_tx, fired_tx).await; + assert!(started_rx.await.is_ok()); + assert!( + !core.is_open(), + "a pre-closed sender is observed as shutdown" + ); + assert!(fired_rx.await.is_ok()); + } + + /// When no shutdown holds at startup, the relay acks but leaves the core + /// open, then seals and fires only once the signal later flips to `true`. + #[tokio::test] + async fn relay_acks_open_then_seals_on_a_later_signal() { + let (tx, rx) = watch::channel(false); + let core = Gate::new(); + let (started_tx, started_rx) = oneshot::channel::<()>(); + let (fired_tx, fired_rx) = oneshot::channel::<()>(); + let relay = tokio::spawn(shutdown_relay(rx, core.clone(), started_tx, fired_tx)); + // The startup ack lands while the core is still open. + assert!( + started_rx.await.is_ok(), + "relay acked its first observation" + ); + assert!( + core.is_open(), + "no shutdown held at startup, so the core stays open" + ); + // A later shutdown seals and fires. + assert!(tx.send(true).is_ok(), "receiver still alive"); + assert!(fired_rx.await.is_ok(), "relay fired once shutdown arrived"); + assert!(!core.is_open(), "the later shutdown sealed the core"); + assert!(relay.await.is_ok(), "relay task joins"); + } +} + +#[cfg(test)] +mod fence_tests { + use super::{ + EffectHandler, SessionControl, SessionEnd, SessionEventSource, SessionLoop, + SessionLoopConfig, + }; + use crate::gate::{Gate, GatedReceiver, SlowOutcome}; + use session_core::fsm::{SessionEffect, SessionEvent, SessionState}; + use std::sync::{Arc, Mutex}; + use tokio::task::JoinSet; + + /// A source that hands the loop a pre-built gated channel (the engine's + /// path), so events arrive with their admission permits and the fast path + /// is eligible. `next_event` is never reached on this path. + struct GatedSource { + rx: Option>, + } + + impl SessionEventSource for GatedSource { + async fn next_event(&mut self) -> Option { + std::future::pending().await + } + fn into_gated_event_channel(self) -> Result, Self> { + match self.rx { + Some(rx) => Ok(rx), + None => Err(self), + } + } + } + + /// Records the effects the loop executes. + struct RecordingHandler { + effects: Arc>>, + } + + impl EffectHandler for RecordingHandler { + async fn execute(&mut self, effect: SessionEffect, _children: &mut JoinSet<()>) { + if let Ok(mut effects) = self.effects.lock() { + effects.push(effect); + } + } + } + + /// Drives a fresh core through the handshake to `Ready` without a handler + /// (the transitions' effects are computed and discarded). + fn core_in_ready() -> Gate { + let core = Gate::new(); + for event in [ + SessionEvent::ConnectionAccepted, + SessionEvent::ClientHandshakeResponse, + SessionEvent::BackendGreetingReceived, + SessionEvent::BackendAuthOk, + ] { + assert!(matches!( + core.apply_slow(event, true), + SlowOutcome::Applied { .. } + )); + } + assert_eq!(core.state_snapshot(), SessionState::Ready); + core + } + + /// The post-select `open` fence drops a *ready* `ClientCommand` once the + /// core is sealed, even though the shutdown oneshot has not fired. Sealing + /// the core directly (rather than via the relay) reproduces exactly that + /// window deterministically: the loop is still in its select, picks the + /// already-queued gated command, and the fence refuses it before any match + /// or effect. The command is never forwarded and its admission permit is + /// returned. + #[tokio::test(start_paused = true)] + async fn post_select_fence_drops_a_ready_command_on_a_sealed_ready_core() { + let core = core_in_ready(); + + // A gated command is in flight (permit registered)... + let (tx, rx) = crate::gate::channel::(core.clone(), 4); + assert!(tx.send(SessionEvent::ClientCommand).await.is_ok()); + assert_eq!(core.pending(), 1, "the command's permit is registered"); + + // ...and the core is sealed with the oneshot not yet fired. + core.seal_closed(); + + let effects = Arc::new(Mutex::new(Vec::new())); + let handler = RecordingHandler { + effects: Arc::clone(&effects), + }; + let (_control_tx, control_rx) = crate::gate::channel::(core.clone(), 8); + let (_shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + + let summary = SessionLoop::new( + core.clone(), + GatedSource { rx: Some(rx) }, + handler, + control_rx, + shutdown_rx, + SessionLoopConfig::default(), + ) + .run() + .await; + + assert_eq!( + summary.end, + SessionEnd::ServerShutdown, + "the fence ends the loop as a shutdown" + ); + let recorded = effects + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!( + recorded + .iter() + .filter(|effect| **effect == SessionEffect::ForwardCommandToBackend) + .count(), + 0, + "a command dropped by the fence is never forwarded" + ); + assert_eq!( + core.pending(), + 0, + "the dropped command's admission permit is returned" + ); + } +} diff --git a/rust/crates/dataplane/src/session_engine.rs b/rust/crates/dataplane/src/session_engine.rs index 58b802cfc..36954561f 100644 --- a/rust/crates/dataplane/src/session_engine.rs +++ b/rust/crates/dataplane/src/session_engine.rs @@ -703,6 +703,22 @@ pub async fn run_bound_session( .await; } +/// Whether the v5 steady-state fast path is enabled, read once from the +/// `TIPROXY_FSM_FASTPATH` environment variable. Default off: only the exact +/// values `1`/`true`/`on` (case-insensitive) enable it, so the gate plumbing is +/// behaviourally inert unless deliberately switched on for same-binary A/B. +fn fast_path_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("TIPROXY_FSM_FASTPATH").is_ok_and(|value| { + matches!( + value.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "on" + ) + }) + }) +} + #[allow(clippy::too_many_arguments, clippy::too_many_lines)] async fn run_bound_session_observed( connection: AcceptedConnection, @@ -833,6 +849,9 @@ async fn run_bound_session_observed( metrics: metrics.clone(), log_context: log_context.clone(), seat, + core: gate.clone(), + shutdown: shutdown.clone(), + fast_path_enabled: fast_path_enabled(), }; // The owner watches the same shutdown signal the loop consumes, so @@ -1475,6 +1494,17 @@ struct Engine { metrics: MetricsRecorder, log_context: SessionLogContext, seat: SessionSeat, + /// Shared session core (FSM + admission gate). S2b: the engine may commit a + /// whitelisted steady transition inline via [`crate::gate::Gate::try_commit_steady`] + /// instead of the `events`→loop→`cmd` round-trip. + core: crate::gate::Gate, + /// Engine-side view of the server shutdown signal, read under the core lock + /// by the fast-path predicate so shutdown is linearized with the commit. + shutdown: watch::Receiver, + /// Whether the steady-state fast path is enabled (`TIPROXY_FSM_FASTPATH`). + /// Default off: when false the engine always takes the slow path, so the + /// gate plumbing is behaviourally inert. + fast_path_enabled: bool, } /// Outcome of waiting for one specific FSM effect. @@ -2819,23 +2849,55 @@ impl Engine { since_connection: command_started.saturating_duration_since(self.accepted_at), traffic_before: self.backend_traffic(), }); - if self.events.send(event).await.is_err() { - return Some(WireErrorSource::Proxy); - } - if event == SessionEvent::ClientCommandQuit { - // Quit tears down: the FSM goes straight to Closing and - // the teardown effects arrive; drain them here. - if let Some(pending) = self.pending_command.take() { - self.record_command(&pending); - } - return None; - } - let _fwd_cmd_permit = match self - .await_effect(SessionEffect::ForwardCommandToBackend) - .await + // S2b steady-state fast path: for a plain `ClientCommand` with the + // gate quiescent (`open && !shutdown && pending == 0` and the FSM in + // `Ready`), commit the `Ready -> Ready` transition inline under the + // core lock and hold its in-flight permit across the forward, + // skipping the `events` -> loop -> `cmd` round-trip. Default-off; + // `Quit` and every non-whitelisted event always take the slow path. + // The permit (fast) or the awaited effect's permit (slow) is held to + // the end of this iteration, keeping `pending >= 1` across the + // forward exactly as the slow path does. + let _fwd_cmd_permit = if self.fast_path_enabled && event == SessionEvent::ClientCommand { - EffectAck::Got(permit) => permit, - EffectAck::Closing => return None, + match self.core.try_commit_steady( + SessionEvent::ClientCommand, + SessionEffect::ForwardCommandToBackend, + &self.shutdown, + ) { + crate::gate::FastPath::Committed(permit) => permit, + crate::gate::FastPath::Fallback => { + if self.events.send(event).await.is_err() { + return Some(WireErrorSource::Proxy); + } + match self + .await_effect(SessionEffect::ForwardCommandToBackend) + .await + { + EffectAck::Got(permit) => permit, + EffectAck::Closing => return None, + } + } + } + } else { + if self.events.send(event).await.is_err() { + return Some(WireErrorSource::Proxy); + } + if event == SessionEvent::ClientCommandQuit { + // Quit tears down: the FSM goes straight to Closing and + // the teardown effects arrive; drain them here. + if let Some(pending) = self.pending_command.take() { + self.record_command(&pending); + } + return None; + } + match self + .await_effect(SessionEffect::ForwardCommandToBackend) + .await + { + EffectAck::Got(permit) => permit, + EffectAck::Closing => return None, + } }; let Some(mut pending) = self.pending_command.take() else { return Some(WireErrorSource::Proxy); diff --git a/rust/crates/dataplane/tests/session_engine.rs b/rust/crates/dataplane/tests/session_engine.rs index 6b77825d2..6467e4ad2 100644 --- a/rust/crates/dataplane/tests/session_engine.rs +++ b/rust/crates/dataplane/tests/session_engine.rs @@ -3205,6 +3205,86 @@ async fn multi_result_command_forwards_every_set_and_keeps_the_session_synchroni stack.dispatch_task.abort(); } +/// Wire-trace dump for the fast-path on/off A/B (see +/// `scripts/fastpath-trace-ab.sh`). Runs a fixed steady-state command scenario +/// against the real engine + fake backend and prints, between markers, the +/// ordered command-phase wire trace: each command forwarded to the backend +/// (its bytes) and the client-visible OK/err outcome, in order. The handshake +/// (random salt / connection id) is excluded by snapshotting the backend +/// transcript after login, so the trace is deterministic across runs. +/// +/// The outer script runs this twice — `TIPROXY_FSM_FASTPATH` unset (slow path) +/// and set (fast path) — and diffs the two traces; byte-identical traces prove +/// the steady-state fast path preserves the observable action sequence. With +/// the flag set, the steady `SELECT`/`UPDATE` commands here commit on the fast +/// path (a quiescent `Ready` `ClientCommand`), the exact case the gate unit +/// tests prove commits; with it unset they take the slow path. +#[tokio::test] +async fn dump_fastpath_wire_trace() { + use std::fmt::Write as _; + let fast = std::env::var("TIPROXY_FSM_FASTPATH").unwrap_or_default(); + let stack = spawn_stack().await; + spawn_route_answer(&stack, 1, 2); + let Some(mut client) = timeout(Duration::from_secs(5), MysqlClient::connect(stack.sql_port)) + .await + .ok() + .flatten() + else { + unreachable!("session established") + }; + + // Baseline: everything forwarded so far is the handshake/auth relay + // (nondeterministic salt), which the trace must exclude. + let forwarded_before = stack + .backend_transcript + .lock() + .map_or(0, |transcript| transcript.len()); + + // Fixed steady-state command scenario: each is a quiescent `Ready` + // `ClientCommand` the fast path is eligible to commit. + let scenario = ["SELECT 1", "SELECT 2", "UPDATE t SET a = 1", "SELECT 3"]; + let mut outcomes = Vec::new(); + for sql in scenario { + let ok = timeout(Duration::from_secs(5), client.query_ok(sql)) + .await + .unwrap_or(false); + outcomes.push(ok); + } + + let forwarded: Vec> = { + let transcript = stack + .backend_transcript + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + transcript[forwarded_before..].to_vec() + }; + + println!("WIRE_TRACE_BEGIN fast_env={fast:?}"); + for (index, bytes) in forwarded.iter().enumerate() { + let mut hex = String::with_capacity(bytes.len() * 2); + for byte in bytes { + let _ = write!(hex, "{byte:02x}"); + } + println!("FWD {index} {hex}"); + } + for (index, ok) in outcomes.iter().enumerate() { + println!("RSP {index} ok={ok}"); + } + println!("WIRE_TRACE_END"); + + // The scenario itself must have run end to end (every command answered), + // independent of which path served it. + assert!(outcomes.iter().all(|ok| *ok), "every command answered OK"); + assert_eq!( + forwarded.len(), + scenario.len(), + "each command forwarded once" + ); + + client.quit().await; + stack.dispatch_task.abort(); +} + /// A complete topology generation re-applies the admission snapshot's /// health-specific policy to an already authenticated socket, without /// replacing its protocol config or disturbing the next command boundary. diff --git a/rust/crates/dataplane/tests/session_loop.rs b/rust/crates/dataplane/tests/session_loop.rs index b9506c989..182aa5037 100644 --- a/rust/crates/dataplane/tests/session_loop.rs +++ b/rust/crates/dataplane/tests/session_loop.rs @@ -900,6 +900,45 @@ async fn preexisting_shutdown_is_not_missed() { assert_eq!(count(&recorded, SessionEffect::ClassifySessionEnd), 1); } +/// S2b startup-race guard: a shutdown that already holds is sealed (via the +/// relay startup ack) before the loop admits any input, so nothing is admitted +/// from a fresh (Accept) session — not even the first queued handshake event. +/// The loop awaits the ack, then its post-select `open` fence drops the first +/// action and terminates. This is the startup guarantee the removed direct +/// precheck gave; the Ready-state fence case is covered deterministically by +/// `session::fence_tests` (which has in-crate core access). +#[tokio::test(start_paused = true)] +async fn preexisting_shutdown_before_handshake_admits_nothing() { + let (fsm_core, _control_tx, control_rx, shutdown_tx, shutdown_rx) = channels(); + let _ = shutdown_tx.send(true); + let mut events = HANDSHAKE.to_vec(); + events.push(SessionEvent::ClientCommand); + let handler = Recorder::default(); + let effects = handler.effects(); + let summary = SessionLoop::new( + fsm_core, + FakeSource::scripted(&events), + handler, + control_rx, + shutdown_rx, + SessionLoopConfig::default(), + ) + .run() + .await; + assert_eq!(summary.end, SessionEnd::ServerShutdown); + let recorded = locked(&effects); + assert_eq!( + count(&recorded, SessionEffect::ForwardCommandToBackend), + 0, + "a ready command is never forwarded under a pre-existing shutdown" + ); + assert_eq!( + count(&recorded, SessionEffect::SendProxyGreeting), + 0, + "nothing is admitted once shutdown is sealed at startup" + ); +} + /// The handshake deadline fires for a stalled pre-auth session; after the /// `BackendAuthOk` **transition itself** it is disarmed — an authenticated /// session that then goes fully idle (probe disabled, no further events) diff --git a/scripts/fastpath-trace-ab.sh b/scripts/fastpath-trace-ab.sh new file mode 100755 index 000000000..e0dea0547 --- /dev/null +++ b/scripts/fastpath-trace-ab.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Fast-path on/off wire-trace A/B (v5 S2b acceptance). +# +# Runs the `dump_fastpath_wire_trace` integration test twice in separate +# processes — once with the steady-state fast path OFF (TIPROXY_FSM_FASTPATH +# unset) and once ON (=1) — and diffs the ordered command-phase wire trace +# (commands forwarded to the backend + client-visible outcomes). Byte-identical +# traces prove the fast path preserves the observable action sequence; any +# divergence fails the gate. +# +# The `fast_env=` marker line intentionally differs between runs and is stripped +# before the diff. +set -euo pipefail + +# Repo root is the script's parent dir; the Rust workspace lives under rust/. +cd "$(dirname "$0")/../rust" + +TEST=dump_fastpath_wire_trace +OFF=$(mktemp) +ON=$(mktemp) +trap 'rm -f "$OFF" "$ON"' EXIT + +extract() { + # Keep only the trace body between the markers, dropping the marker lines + # themselves (the BEGIN line carries the differing fast_env tag). + sed -n '/WIRE_TRACE_BEGIN/,/WIRE_TRACE_END/p' \ + | grep -E '^(FWD|RSP) ' +} + +echo "== running $TEST with fast path OFF ==" +env -u TIPROXY_FSM_FASTPATH \ + cargo test -p dataplane --test session_engine -- --exact "$TEST" --nocapture \ + | extract > "$OFF" + +echo "== running $TEST with fast path ON ==" +TIPROXY_FSM_FASTPATH=1 \ + cargo test -p dataplane --test session_engine -- --exact "$TEST" --nocapture \ + | extract > "$ON" + +if [ ! -s "$OFF" ] || [ ! -s "$ON" ]; then + echo "FAIL: empty trace captured (OFF=$(wc -l <"$OFF") ON=$(wc -l <"$ON") lines)" >&2 + exit 1 +fi + +echo "== diffing traces ==" +if diff -u "$OFF" "$ON"; then + echo "PASS: fast-path on/off wire traces are byte-identical ($(wc -l <"$OFF") lines)" +else + echo "FAIL: fast-path on/off wire traces diverge" >&2 + exit 1 +fi