From 520e54995fad25e74c5432a50401145d0ef2f3a4 Mon Sep 17 00:00:00 2001 From: bb7133 Date: Mon, 5 Oct 2026 06:34:38 -0700 Subject: [PATCH 1/4] dataplane: steady-state fast path at the command point (v5 S2b, default-off) First real fast path: for a plain ClientCommand with the gate quiescent, the engine commits the Ready->Ready transition inline under the core lock and forwards directly, skipping the events->loop->cmd round-trip. Gated by TIPROXY_FSM_FASTPATH (default off), so with the flag unset the engine always takes the slow path and behaviour is unchanged. gate.rs: - FsmCore gains `fast_transitions`; `open` is now the single admission authority, read under the one lock by both admission paths. - `try_commit_steady(event, expected_effect, shutdown)`: under the lock, commits only when `open && !shutdown && pending == 0 && event == ClientCommand && FSM Ready` and a speculative clone yields exactly the single expected effect. On commit it writes the clone back, registers an in-flight GatePermit (so `pending` holds the floor across the forward, exactly as the slow path's effect permit does), bumps `fast_transitions`, and returns FastPath::Committed(permit). Otherwise Fallback with zero core writes. shutdown is read inside the lock (sender-closed counts as shutdown) so it linearizes with the commit. - `apply_slow(event, gated)` -> SlowOutcome: a gated ordinary event is DroppedClosed once the core is sealed (the slow-path shutdown-admission refusal, atomic with the transition); terminal events pass gated=false. session.rs: - The shutdown relay is now the loop's SOLE shutdown observer: it seals the core (open=false) under the lock BEFORE firing the loop oneshot, so "open sealed" and "loop drops the command" are one linearization point. The loop's direct has_changed precheck and mark_changed are removed. - `apply` takes `gated`: LoopAction::Event (engine ordinary events) is gated=true so a post-seal ClientCommand is dropped under the lock, matching the fast path refusing to commit; Control/Deadline/Probe and close_via_fsm's terminal events are gated=false. session_engine.rs: - Engine gains `core`, `shutdown`, and `fast_path_enabled` (read once from the env). The command point tries `try_commit_steady` for ClientCommand when enabled; Committed holds the in-flight permit across the forward (same scope as the slow path's await_effect permit), Fallback and every non-whitelisted event (incl. Quit) take the original slow path. Tests: gate-level commit + each fallback condition (shutdown value / sender-closed / sealed / pending!=0 / Quit / not-Ready) + effect-mismatch guard + the fast commit lands on the SAME FSM state as the slow transition (on/off FSM equivalence); gated apply dropped once sealed while terminal still applies; relay seals the core before firing on both a true signal and a dropped sender. All suites green with the flag off (behaviour-neutral); fmt + clippy --all-targets clean. Co-Authored-By: Claude Opus 4.8 --- rust/crates/dataplane/src/gate.rs | 337 ++++++++++++++++++-- rust/crates/dataplane/src/session.rs | 143 +++++++-- rust/crates/dataplane/src/session_engine.rs | 94 +++++- 3 files changed, 499 insertions(+), 75 deletions(-) diff --git a/rust/crates/dataplane/src/gate.rs b/rust/crates/dataplane/src/gate.rs index 7069af34a..98f28e200 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,99 @@ 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 + /// production admission paths read `open` inline under the lock + /// ([`Gate::apply_slow`], [`Gate::try_commit_steady`]); this accessor is for + /// tests/diagnostics. + #[cfg(test)] + 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, + } + } + + /// 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 +205,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 +349,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 +365,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 +384,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 +421,170 @@ 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); + } + #[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..49b3b2885 100644 --- a/rust/crates/dataplane/src/session.rs +++ b/rust/crates/dataplane/src/session.rs @@ -304,8 +304,15 @@ 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, + fired: oneshot::Sender<()>, +) { loop { + // `borrow_and_update` reads the current value (so a shutdown that + // predates this task is seen on the first poll); a dropped sender is a + // shutdown too. if *shutdown.borrow_and_update() { break; } @@ -313,6 +320,12 @@ async fn shutdown_relay(mut shutdown: watch::Receiver, fired: oneshot::Sen 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 +471,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,15 +495,21 @@ 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's + // first `borrow_and_update` sees a pre-existing shutdown, so no + // `mark_changed` precheck is needed. let (relay_tx, mut shutdown_fired) = oneshot::channel::<()>(); let _relay = AbortOnDrop(tokio::spawn(shutdown_relay( self.shutdown.clone(), + self.core.clone(), relay_tx, ))); @@ -505,8 +527,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 @@ -552,13 +579,18 @@ impl SessionLoop { // 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; + 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 +607,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 +619,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 +636,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 +651,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 +694,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 +746,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 +818,38 @@ enum LoopEnd { Shutdown, SourceExhausted, } + +#[cfg(test)] +mod relay_tests { + use super::shutdown_relay; + use crate::gate::Gate; + use tokio::sync::{oneshot, watch}; + + /// The shutdown relay must seal the core *before* firing the loop oneshot, + /// so `open == false` is published under the core lock strictly before the + /// loop (or the engine's fast path) can observe the wake — the single + /// shutdown-admission linearization point. + #[tokio::test] + async fn relay_seals_core_before_firing_on_a_true_signal() { + let (_tx, rx) = watch::channel(true); + let core = Gate::new(); + assert!(core.is_open()); + let (fired_tx, fired_rx) = oneshot::channel::<()>(); + shutdown_relay(rx, core.clone(), fired_tx).await; + assert!(!core.is_open(), "relay sealed the core on a true signal"); + assert!(fired_rx.await.is_ok(), "relay fired the loop oneshot"); + } + + /// A dropped shutdown sender is a shutdown too: the relay still seals and + /// fires (so a torn-down runtime never leaves the core admitting). + #[tokio::test] + async fn relay_seals_core_when_sender_is_dropped() { + let (tx, rx) = watch::channel(false); + drop(tx); + let core = Gate::new(); + let (fired_tx, fired_rx) = oneshot::channel::<()>(); + shutdown_relay(rx, core.clone(), fired_tx).await; + assert!(!core.is_open(), "sender-closed is observed as shutdown"); + assert!(fired_rx.await.is_ok()); + } +} 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); From 125b612c70064cf33e9f8d86c4f08331fa0010df Mon Sep 17 00:00:00 2001 From: bb7133 Date: Mon, 5 Oct 2026 06:56:07 -0700 Subject: [PATCH 2/4] dataplane: close two S2b shutdown-admission gaps (relay startup ack + loop open fence) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodexM5's S2b review found two windows where shutdown could be bypassed: 1. Relay startup race. `event_loop` entered its select immediately after spawning the relay, so a shutdown that already held (value `true` or a dropped sender) could be missed if an input event was already ready — the parent could dispatch before the relay sealed. shutdown_relay now makes its first observation up front (sealing the core if shutdown already holds, treating a closed sender as shutdown), acks via a new `started` oneshot, and only then fires `ServerShutdown`. `event_loop` awaits that ack before touching any input, restoring the synchronous precheck guarantee the removed direct `has_changed` read provided. 2. Seal->fire window. There was no loop-side `open` fence between the relay sealing the core and the loop observing the oneshot, and Deadline/BackendProbe/Control apply with `gated = false`, so they could still transition in that window. A post-select fence now returns `LoopEnd::Shutdown` whenever `!core.is_open()`, before any match/effect, dropping the dequeued action (and returning its permit). Only the relay seals `open` mid-loop, so `!open` there means shutdown; terminal close in `run` is still reached ungated. `Gate::is_open` is back to a non-test `pub(crate)` accessor (the loop fence reads it in production). Tests: relay unit tests now cover the startup ack and seal ordering for a pre-existing true signal, a pre-closed sender, and the ack-open-then-seal path on a later signal; a new loop test asserts that under a pre-existing shutdown a ready ClientCommand is never forwarded (no ForwardCommandToBackend / SendProxyGreeting). All suites green with the flag off; fmt + clippy clean. Co-Authored-By: Claude Opus 4.8 --- rust/crates/dataplane/src/gate.rs | 7 +- rust/crates/dataplane/src/session.rs | 121 ++++++++++++++++---- rust/crates/dataplane/tests/session_loop.rs | 37 ++++++ 3 files changed, 141 insertions(+), 24 deletions(-) diff --git a/rust/crates/dataplane/src/gate.rs b/rust/crates/dataplane/src/gate.rs index 98f28e200..3734410f6 100644 --- a/rust/crates/dataplane/src/gate.rs +++ b/rust/crates/dataplane/src/gate.rs @@ -113,10 +113,9 @@ impl Gate { } /// Whether the core is still open (loop running, not terminating). The - /// production admission paths read `open` inline under the lock - /// ([`Gate::apply_slow`], [`Gate::try_commit_steady`]); this accessor is for - /// tests/diagnostics. - #[cfg(test)] + /// 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 } diff --git a/rust/crates/dataplane/src/session.rs b/rust/crates/dataplane/src/session.rs index 49b3b2885..42825d8dc 100644 --- a/rust/crates/dataplane/src/session.rs +++ b/rust/crates/dataplane/src/session.rs @@ -307,16 +307,33 @@ impl LoopEventRx { 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 { - // `borrow_and_update` reads the current value (so a shutdown that - // predates this task is seen on the first poll); a dropped sender is a - // shutdown too. - if *shutdown.borrow_and_update() { + if shutdown.changed().await.is_err() { break; } - if shutdown.changed().await.is_err() { + if *shutdown.borrow_and_update() { break; } } @@ -503,15 +520,25 @@ impl SessionLoop { // 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's - // first `borrow_and_update` sees a pre-existing shutdown, so no - // `mark_changed` precheck is needed. + // 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 @@ -557,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), @@ -825,31 +866,71 @@ mod relay_tests { use crate::gate::Gate; use tokio::sync::{oneshot, watch}; - /// The shutdown relay must seal the core *before* firing the loop oneshot, - /// so `open == false` is published under the core lock strictly before the - /// loop (or the engine's fast path) can observe the wake — the single - /// shutdown-admission linearization point. + /// 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_core_before_firing_on_a_true_signal() { + 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(), fired_tx).await; - assert!(!core.is_open(), "relay sealed the core on a true signal"); + 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 dropped shutdown sender is a shutdown too: the relay still seals and - /// fires (so a torn-down runtime never leaves the core admitting). + /// 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_core_when_sender_is_dropped() { + 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(), fired_tx).await; - assert!(!core.is_open(), "sender-closed is observed as shutdown"); + 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"); + } } diff --git a/rust/crates/dataplane/tests/session_loop.rs b/rust/crates/dataplane/tests/session_loop.rs index b9506c989..fd892ed6c 100644 --- a/rust/crates/dataplane/tests/session_loop.rs +++ b/rust/crates/dataplane/tests/session_loop.rs @@ -900,6 +900,43 @@ async fn preexisting_shutdown_is_not_missed() { assert_eq!(count(&recorded, SessionEffect::ClassifySessionEnd), 1); } +/// S2b startup-race guard: a shutdown that already holds must be sealed (via the +/// relay startup ack) before the loop admits any input, so even a +/// `ClientCommand` already queued in the source is never forwarded — the loop +/// awaits the ack, then its post-select `open` fence drops the first action and +/// terminates. This is the guarantee the removed direct precheck used to give. +#[tokio::test(start_paused = true)] +async fn preexisting_shutdown_refuses_a_ready_command() { + 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) From 22a3274b953c872ca2649e2b541beaa2032e94ea Mon Sep 17 00:00:00 2001 From: bb7133 Date: Mon, 5 Oct 2026 07:06:27 -0700 Subject: [PATCH 3/4] dataplane: close the S2b stale-probe race + add the on/off wire-trace A/B MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third S2b shutdown-adjacent gap (CodexM5): the backend probe snapshotted the FSM state before `backend_active().await` and applied `BackendIoError` afterwards without re-checking. A fast commit could turn `Ready -> Command` (holding a pending permit) during that await, so a stale `inactive` result would tear down an in-flight command. Fix — admit the probe atomically with the core: - gate.rs: `try_quiescent_operation(state_ok)` registers and returns an operation `GatePermit` only when the core is `open`, `pending == 0`, and `state_ok(state)` holds; otherwise `None`. - session.rs: the `BackendProbe` branch takes that permit (with `probe_safe`) and holds it across `backend_active().await`. While held `pending >= 1`, so `try_commit_steady` falls back and cannot commit a command mid-probe; a non-quiescent core simply skips the probe this round. The stale pre-await snapshot and its S2b TODO are gone. Wire-trace on/off A/B (the S2b acceptance item, CodexM5 chose this form): - tests/session_engine.rs `dump_fastpath_wire_trace`: drives a fixed steady-state scenario (SELECT/UPDATE x4) through the real engine + fake backend and prints the ordered command-phase trace (bytes forwarded to the backend + client outcomes) between markers, excluding the nondeterministic handshake. Reads `TIPROXY_FSM_FASTPATH` via the normal path. - scripts/fastpath-trace-ab.sh runs it twice in separate processes (flag unset vs `1`), strips the differing `fast_env` marker, and diffs. Verified locally: the two traces are byte-identical (fast path preserves the observable action sequence). Test: gate-level probe-permit coverage — not admitted when not probe-safe / sealed / `pending != 0`; admitted when quiescent, and while the permit is held `try_commit_steady` falls back (the probe-await x fast-commit interleaving) with the FSM untouched. All suites green with the flag off; `fmt --check` + `clippy --all-targets` clean. Co-Authored-By: Claude Opus 4.8 --- rust/crates/dataplane/src/gate.rs | 87 +++++++++++++++++++ rust/crates/dataplane/src/session.rs | 21 +++-- rust/crates/dataplane/tests/session_engine.rs | 80 +++++++++++++++++ scripts/fastpath-trace-ab.sh | 51 +++++++++++ 4 files changed, 231 insertions(+), 8 deletions(-) create mode 100755 scripts/fastpath-trace-ab.sh diff --git a/rust/crates/dataplane/src/gate.rs b/rust/crates/dataplane/src/gate.rs index 3734410f6..ffb4ffdb1 100644 --- a/rust/crates/dataplane/src/gate.rs +++ b/rust/crates/dataplane/src/gate.rs @@ -170,6 +170,28 @@ impl Gate { } } + /// 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 { @@ -584,6 +606,71 @@ mod tests { 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 42825d8dc..e8506d9dc 100644 --- a/rust/crates/dataplane/src/session.rs +++ b/rust/crates/dataplane/src/session.rs @@ -612,14 +612,19 @@ 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 { + // 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, 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/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 From 34dcd33b823e225592ce483b1535ec182ca473c5 Mon Sep 17 00:00:00 2001 From: bb7133 Date: Mon, 5 Oct 2026 07:33:03 -0700 Subject: [PATCH 4/4] dataplane: strengthen S2b shutdown-fence tests (Ready-state + sealed-window) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodexM5's S2b review flagged two test-coverage gaps: 1. `preexisting_shutdown_refuses_a_ready_command` started from a fresh (Accept) session with the handshake queued first, so the pre-existing shutdown sealed the core before the first select and the first handshake event — not the ClientCommand — was what the fence dropped. Renamed to `preexisting_shutdown_before_handshake_admits_nothing` to describe what it actually proves (the startup-ack path admits nothing from Accept). 2. There was no deterministic test of the post-select `open` fence dropping a *ready* action in the window where the core is sealed but the shutdown oneshot has not fired. The relay/watch path can't hit that window deterministically from an external test, and `tests/` can't reach the `pub(crate)` core. New in-crate `session::fence_tests`: a minimal gated source + recording handler drive a `SessionLoop` over a core that was driven to `Ready` (`apply_slow`) and then `seal_closed()` — directly reproducing the sealed-but-oneshot-not-fired window. A gated `ClientCommand` is already in flight (its permit registered); the loop picks it in the first select and the fence refuses it before any match/effect. Asserts the command is never forwarded (zero `ForwardCommandToBackend`), the loop ends as `ServerShutdown`, and the dropped command's admission permit is returned (`pending` back to 0). Seal-before-fire ordering stays covered by the relay unit tests. All suites green with the flag off; fmt + clippy --all-targets clean. Co-Authored-By: Claude Opus 4.8 --- rust/crates/dataplane/src/session.rs | 123 ++++++++++++++++++++ rust/crates/dataplane/tests/session_loop.rs | 14 ++- 2 files changed, 131 insertions(+), 6 deletions(-) diff --git a/rust/crates/dataplane/src/session.rs b/rust/crates/dataplane/src/session.rs index e8506d9dc..762335f81 100644 --- a/rust/crates/dataplane/src/session.rs +++ b/rust/crates/dataplane/src/session.rs @@ -939,3 +939,126 @@ mod relay_tests { 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/tests/session_loop.rs b/rust/crates/dataplane/tests/session_loop.rs index fd892ed6c..182aa5037 100644 --- a/rust/crates/dataplane/tests/session_loop.rs +++ b/rust/crates/dataplane/tests/session_loop.rs @@ -900,13 +900,15 @@ async fn preexisting_shutdown_is_not_missed() { assert_eq!(count(&recorded, SessionEffect::ClassifySessionEnd), 1); } -/// S2b startup-race guard: a shutdown that already holds must be sealed (via the -/// relay startup ack) before the loop admits any input, so even a -/// `ClientCommand` already queued in the source is never forwarded — the loop -/// awaits the ack, then its post-select `open` fence drops the first action and -/// terminates. This is the guarantee the removed direct precheck used to give. +/// 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_refuses_a_ready_command() { +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();