Skip to main content

subc_daemon/
forwarding.rs

1use std::{
2    collections::{BTreeMap, HashMap, HashSet},
3    error::Error,
4    fmt,
5    sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard},
6    time::Duration,
7};
8
9use subc_control::{ClientControlResponse, RouteCloseReason};
10use subc_protocol::{
11    manifest::Concurrency,
12    session::{LiveRoot, ModuleControlResponse, ModuleControlResponseToModule},
13    ErrorBody, Flags, FrameType, Principal, Priority,
14};
15use tokio::sync::{oneshot, Semaphore};
16use tokio::time::Instant;
17use tracing::{debug, info, warn};
18
19use crate::{
20    control::{RouteBindBreakers, RouteBindConcurrency},
21    observability::DaemonCounters,
22    registry::ConnectionId,
23    router::FrameSink,
24    scopes::{BoundScope, ScopeDrain, ScopeTag, ScopeTagChange},
25    Frame, ProjectRootId,
26};
27
28/// Default per-channel request-credit window for modules that schedule internally.
29const DEFAULT_MODULE_MANAGED_WINDOW: usize = 32;
30
31/// High per-channel cap for stateless modules; this is an OOM guard, not scheduling policy.
32const STATELESS_PARALLEL_WINDOW: usize = 1024;
33
34/// A stopped probe cycle cannot retain its last unanswered correlation forever.
35/// Active endpoints replace the tombstone on their next serial health probe;
36/// this backstop covers endpoints that stop probing altogether.
37const HEALTH_PROBE_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
38
39/// Module connection identity used in forwarding keys.
40///
41/// The generation is bumped every time a module connection is registered so a future restart cannot
42/// accidentally reuse a stale `(connection_id, route_channel)` binding.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44pub struct ModuleEndpointId {
45    pub connection_id: ConnectionId,
46    pub generation: u64,
47}
48
49/// Client-local route key. A route channel is unique only within one client connection.
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
51pub(crate) struct ClientRouteKey {
52    pub connection_id: ConnectionId,
53    pub channel: u16,
54}
55
56/// Module-local route key. A route channel is unique only within one live module endpoint.
57#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
58pub(crate) struct ModuleRouteKey {
59    pub endpoint: ModuleEndpointId,
60    pub channel: u16,
61}
62
63#[derive(Debug)]
64pub(crate) struct RouteBinding {
65    pub client_connection_id: ConnectionId,
66    pub client_sink: FrameSink,
67    pub client_negotiated_ver: u8,
68    pub client_channel: u16,
69    pub client_epoch: u32,
70    pub module_id: String,
71    pub module_endpoint: ModuleEndpointId,
72    pub module_sink: FrameSink,
73    pub module_negotiated_ver: u8,
74    pub module_channel: u16,
75    pub module_epoch: u32,
76    pub principal: Principal,
77    pub project_root: Option<ProjectRootId>,
78    pub bound_at: Instant,
79    pub flow: Arc<ChannelFlow>,
80    /// The scope the route was admitted under, with the tag its bind was
81    /// stamped at. A scope change finds the routes to close by it.
82    pub scope: Option<BoundScope>,
83}
84
85#[derive(Debug, Clone)]
86pub(crate) enum DataRoute {
87    Client(DataRouteState),
88    Module(DataRouteState),
89}
90
91#[derive(Debug, Clone)]
92pub(crate) enum DataRouteState {
93    Bound(Arc<RouteBinding>),
94    Reserved,
95    EpochMismatch,
96    Absent,
97}
98
99/// Which kind of peer a route GOODBYE is being delivered to. This decides what
100/// happens when the GOODBYE cannot be enqueued (egress full/closed):
101/// - `Client`: escalate to closing that client connection (a socket close is a
102///   stronger teardown signal, and a full client egress means it is the slow
103///   client we would drop anyway).
104/// - `Module`: never close; deliver late instead (see
105///   [`send_module_route_goodbye`]), and drop only if the module still has no
106///   room after [`LATE_MODULE_GOODBYE_DEADLINE`]. A client-disconnect notifies the
107///   SHARED module that one client's route is gone; closing the module on its
108///   egress backpressure would tear down every co-tenant client (the exact
109///   cross-tenant blast radius this never-close rule exists to prevent — observed when a
110///   flooding dead client filled BOTH its own and the module's egress, so its
111///   route-gone GOODBYE to the module failed and closed the shared connection).
112///   subc has already removed the route from its forwarding state and drops
113///   stale module frames for the released channel (see router.rs), so subc's
114///   own routing is correct. The residual: under SUSTAINED module-egress
115///   backpressure a module-targeted route-gone notification can be lost, which a
116///   module using it for client-refcounting (e.g. AFT's session accounting)
117///   would miss. This is INTENTIONALLY ACCEPTED, not a gap. A consuming module
118///   must bound stale bindings with its own idle-activity reaper (last-touched
119///   TTL) independent of route-gone signals — AFT does exactly this, so a lost
120///   GOODBYE degrades to "the
121///   binding stays warm until its idle TTL" (bounded wasted resources), never an
122///   unbounded leak; disk-durable replay is unaffected. A dedicated reliable
123///   module control lane was evaluated and deliberately NOT
124///   built: it would add starvation-avoidance machinery to the thin core for a
125///   bounded warm-resource window that is not a correctness issue. never-close
126///   is the invariant that matters here.
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
128pub(crate) enum GoodbyeTargetKind {
129    Client,
130    Module,
131}
132
133#[derive(Debug, Clone)]
134pub(crate) struct GoodbyeTarget {
135    pub connection_id: ConnectionId,
136    pub sink: FrameSink,
137    pub negotiated_ver: u8,
138    pub channel: u16,
139    pub epoch: u32,
140    pub kind: GoodbyeTargetKind,
141    /// The module on the other end of the route: for a module-targeted relay
142    /// the receiving module (attributing a dropped relay), for a client target
143    /// the module whose route is going away (naming it when an undeliverable
144    /// relay closes the client).
145    pub module_id: Option<String>,
146}
147
148/// The frame a client connection's egress queue refused, for the diagnosis
149/// logged when that refusal closes the connection.
150#[derive(Debug, Clone, Copy)]
151pub(crate) struct UndeliveredFrame<'a> {
152    /// The module on the other end of the route the frame belonged to, when known.
153    pub module_id: Option<&'a str>,
154    /// The refusing connection's sink, read for what its queue held.
155    pub sink: &'a FrameSink,
156}
157
158/// How a route's principal appears in logs: `direct`, or `reserved:<module>`.
159pub(crate) fn principal_label(principal: &Principal) -> String {
160    match principal {
161        Principal::Reserved { module_id } => format!("reserved:{module_id}"),
162        Principal::Direct => "direct".to_string(),
163        other => format!("{other:?}"),
164    }
165}
166
167/// The distinct principals of the routes bound on one client connection,
168/// sorted and comma-joined, or `none` when it has no bound route.
169fn connection_principals_locked(inner: &ForwardingInner, connection_id: ConnectionId) -> String {
170    let labels = inner
171        .client_to_module
172        .iter()
173        .filter(|(key, _)| key.connection_id == connection_id)
174        .map(|(_, route)| principal_label(&route.principal))
175        .collect::<std::collections::BTreeSet<_>>();
176    if labels.is_empty() {
177        "none".to_string()
178    } else {
179        labels.into_iter().collect::<Vec<_>>().join(",")
180    }
181}
182
183impl GoodbyeTarget {
184    /// True only when an undeliverable GOODBYE should escalate to closing the
185    /// target connection. Never escalate for module recipients.
186    pub(crate) fn close_on_delivery_failure(&self) -> bool {
187        matches!(self.kind, GoodbyeTargetKind::Client)
188    }
189}
190
191/// How long a route GOODBYE that a module's egress queue refused keeps waiting
192/// for room before it is given up and counted as dropped.
193///
194/// A module that stops reading for a moment (a GC pause, a long synchronous
195/// handler, a reconnect herd it is still working through) should still learn
196/// that the route is gone, because a module that never hears the GOODBYE keeps
197/// the route and keeps sending on it for as long as its connection lives. The
198/// value is the default module drain timeout: the daemon already treats that as
199/// the longest it is reasonable to wait on a module that is busy but alive, and
200/// a module that cannot free 21 bytes of egress in that time is wedged rather
201/// than slow. Resolving each module's own configured drain timeout here would
202/// need a registry lookup on a path that holds only the module's sink.
203pub(crate) const LATE_MODULE_GOODBYE_DEADLINE: Duration = crate::supervise::DEFAULT_DRAIN_TIMEOUT;
204
205/// Send a route GOODBYE to a module without ever closing the module's shared
206/// connection, which would sever every other route it serves.
207///
208/// The frame is enqueued at once when the module's egress queue has room. When
209/// it does not, a detached task waits for room (bounded by
210/// [`LATE_MODULE_GOODBYE_DEADLINE`]) and sends it then. Arriving late is safe:
211/// the daemon has already released the route, and module SDKs tear a route
212/// down only when both the channel AND the epoch match the route installed on
213/// that channel, so a late GOODBYE for a channel since reused at a newer epoch
214/// is ignored. The GOODBYE is counted in `goodbye_relay_module_dropped` only
215/// when the connection is closed or the deadline passes.
216pub(crate) fn send_module_route_goodbye(
217    counters: &DaemonCounters,
218    sink: &FrameSink,
219    frame: Frame,
220    module_id: Option<&str>,
221    context: &'static str,
222) {
223    let channel = frame.header.channel;
224    let epoch = frame.header.epoch;
225    let Err(err) = sink.try_send(frame.clone()) else {
226        return;
227    };
228    // Waiting is pointless once the module's writer is gone, and impossible
229    // outside a Tokio runtime (cleanup that runs from a destructor at shutdown).
230    let runtime = match tokio::runtime::Handle::try_current() {
231        Ok(runtime) if !sink.is_closed() => runtime,
232        _ => {
233            counters.increment_goodbye_relay_module_dropped(module_id);
234            warn!(
235            module_id = module_id.unwrap_or("unknown"),
236            route_channel = channel,
237            route_epoch = epoch,
238            error = %err,
239            context,
240            "route GOODBYE to module dropped: module connection is closed; not closing shared module connection"
241            );
242            return;
243        }
244    };
245    debug!(
246        module_id = module_id.unwrap_or("unknown"),
247        route_channel = channel,
248        route_epoch = epoch,
249        error = %err,
250        context,
251        "module egress queue refused route GOODBYE; delivering it once the module frees room"
252    );
253    let counters = counters.clone();
254    let sink = sink.clone();
255    let module_id = module_id.map(str::to_string);
256    runtime.spawn(async move {
257        let outcome = tokio::time::timeout(LATE_MODULE_GOODBYE_DEADLINE, sink.send(frame)).await;
258        let why = match outcome {
259            Ok(Ok(())) => {
260                debug!(
261                    module_id = module_id.as_deref().unwrap_or("unknown"),
262                    route_channel = channel,
263                    route_epoch = epoch,
264                    context,
265                    "late route GOODBYE delivered to module"
266                );
267                return;
268            }
269            Ok(Err(err)) => err.to_string(),
270            Err(_) => format!(
271                "module egress queue had no room within {LATE_MODULE_GOODBYE_DEADLINE:?}"
272            ),
273        };
274        counters.increment_goodbye_relay_module_dropped(module_id.as_deref());
275        warn!(
276            module_id = module_id.as_deref().unwrap_or("unknown"),
277            route_channel = channel,
278            route_epoch = epoch,
279            error = %why,
280            context,
281            "route GOODBYE to module dropped under backpressure; not closing shared module connection"
282        );
283    });
284}
285
286/// One route currently served by a module endpoint.
287///
288/// `goodbye_target` is deliberately retained alongside the census projection so
289/// a draining caller can address exactly the same route set that this read
290/// reports, without a second forwarding-table pass.
291#[derive(Debug, Clone)]
292pub(crate) struct EndpointRoute {
293    pub goodbye_target: GoodbyeTarget,
294    pub principal: Principal,
295    pub bound_at: Instant,
296    pub draining: bool,
297    /// WHY the endpoint is draining, when it is. Carried per-route so the
298    /// census can answer "closing because of what" without a second lookup;
299    /// `None` exactly when `draining` is false (one source: the drain map).
300    pub drain_reason: Option<RouteCloseReason>,
301}
302
303/// One live route a scope change closed: both ends to send GOODBYE to, and the
304/// reason for the client's `route.closed`.
305#[derive(Debug, Clone)]
306pub(crate) struct ScopeDrainedRoute {
307    pub reason: RouteCloseReason,
308    pub scope: BoundScope,
309    pub module_id: String,
310    pub client: GoodbyeTarget,
311    pub module: GoodbyeTarget,
312}
313
314#[derive(Debug)]
315pub(crate) struct PendingRouteBindRelay {
316    pub endpoint: ModuleEndpointId,
317    pub module_sink: FrameSink,
318    pub negotiated_ver: u8,
319    pub client_channel: u16,
320    pub client_epoch: u32,
321    pub module_channel: u16,
322    pub module_epoch: u32,
323    pub corr: u64,
324    pub receiver: oneshot::Receiver<RouteBindRelayOutcome>,
325}
326
327#[derive(Debug, Clone)]
328pub(crate) struct ModuleDrainTarget {
329    pub endpoint: ModuleEndpointId,
330    pub sink: FrameSink,
331    pub negotiated_ver: u8,
332    pub abandoned_bindings: Vec<GoodbyeTarget>,
333    pub excluded_subscriptions: u32,
334}
335
336/// One registered module connection, as [`ForwardingTable::module_connections`]
337/// reports it: enough to send it a frame and then close it.
338#[cfg(unix)]
339#[derive(Debug, Clone)]
340pub(crate) struct ModuleConnectionTarget {
341    pub module_id: String,
342    pub endpoint: ModuleEndpointId,
343    pub sink: FrameSink,
344    pub negotiated_ver: u8,
345}
346
347#[derive(Debug, Clone)]
348pub(crate) enum RouteBindRelayOutcome {
349    Accepted,
350    Rejected(ErrorBody),
351    ModuleGone(String),
352}
353
354/// What [`ForwardingTable::cutover_candidate`] did.
355#[derive(Debug, Clone, Copy, PartialEq, Eq)]
356pub(crate) struct ForwardingCutover {
357    /// The endpoint now in the active slot (the former candidate).
358    pub promoted: ModuleEndpointId,
359    /// The endpoint demoted out of the active slot, to be drained by endpoint.
360    /// `None` when the incumbent's connection was already gone.
361    pub incumbent: Option<ModuleEndpointId>,
362}
363
364/// What tearing down one connection released.
365#[derive(Debug)]
366pub(crate) struct ConnectionCleanup {
367    /// Routes whose other end must be sent GOODBYE.
368    pub released: Vec<GoodbyeTarget>,
369    /// Pending route.bind relays to the closed module that were aborted: opens
370    /// that had not been bound yet. Zero when a client connection closed.
371    pub abandoned_relays: u32,
372}
373
374#[derive(Debug, Clone)]
375pub(crate) struct PendingRelayCompletion {
376    pub settled: bool,
377    pub abandoned: Option<GoodbyeTarget>,
378}
379
380#[derive(Debug)]
381pub(crate) struct PendingModuleControlRpc {
382    pub endpoint: ModuleEndpointId,
383    pub module_sink: FrameSink,
384    pub negotiated_ver: u8,
385    pub corr: u64,
386    pub receiver: oneshot::Receiver<ModuleControlRpcOutcome>,
387}
388
389#[derive(Debug, Clone)]
390pub(crate) enum ModuleControlRpcOutcome {
391    Response(ModuleControlResponse),
392    Rejected(ErrorBody),
393    ModuleGone(String),
394    MalformedResponse(String),
395    UnexpectedOp { expected: String, actual: String },
396    DeadlineElapsed,
397}
398
399#[derive(Debug, Clone, PartialEq, Eq)]
400pub(crate) enum ModuleControlRpcCompletion {
401    Unknown,
402    Settled,
403    LateHealthAnswer {
404        module_id: String,
405        latency: Duration,
406    },
407}
408
409#[derive(Debug)]
410struct PendingModuleControlRpcEntry {
411    expected_op: String,
412    deadline: Instant,
413    health_probe_started_at: Option<Instant>,
414    sender: oneshot::Sender<ModuleControlRpcOutcome>,
415}
416
417#[derive(Debug)]
418struct HealthProbeTombstone {
419    expected_op: String,
420    module_id: String,
421    probe_started_at: Instant,
422    expires_at: Instant,
423}
424
425#[derive(Debug, Clone)]
426struct RouteReservation {
427    client_key: ClientRouteKey,
428    module_key: ModuleRouteKey,
429    client_epoch: u32,
430    module_epoch: u32,
431    project_root: Option<ProjectRootId>,
432}
433
434#[derive(Debug)]
435struct PendingRouteBindRelayEntry {
436    reservation: RouteReservation,
437    client_sink: FrameSink,
438    client_negotiated_ver: u8,
439    client_permit: crate::router::EgressPermit,
440    route_open_frame: Frame,
441    principal: Principal,
442    /// The scope admitted at `route.open`, with the tag captured then. Commit
443    /// compares it with the published tag.
444    scope: Option<BoundScope>,
445    deadline: Instant,
446    /// Set once `route.open` has finished putting the route.bind on the
447    /// module's egress queue. A path that drops this entry WITHOUT a module
448    /// reply (client teardown, module drain, the open's own abort) sends the
449    /// module a GOODBYE for the bind only when this is set, because only then
450    /// is the GOODBYE guaranteed to queue behind the bind. When it is not set,
451    /// the removal is recorded in `relays_released_before_enqueue` instead and
452    /// `route.open` sends that GOODBYE itself once its send has completed.
453    ///
454    /// A module reply does not consult it: an ack proves the module received
455    /// the bind, so the reply path decides about the GOODBYE on its own.
456    relay_enqueued: bool,
457    sender: oneshot::Sender<RouteBindRelayOutcome>,
458}
459
460/// What `mark_route_bind_relay_enqueued` found once `route.open` had put the
461/// route.bind on the module's queue.
462#[derive(Debug, Clone, Copy, PartialEq, Eq)]
463pub(crate) enum RouteBindRelayMark {
464    /// The relay is still pending; later removals handle their own GOODBYE.
465    Marked,
466    /// The module already answered and the reply path settled the relay:
467    /// it committed the route, or it refused it and sent any GOODBYE the
468    /// module needed. `route.open` must send nothing.
469    SettledByReply,
470    /// The relay was dropped without a module reply while the bind was still
471    /// being queued, so nothing has told the module to drop the binding it
472    /// is about to create. `route.open` must send that GOODBYE.
473    ReleasedWithoutReply,
474}
475
476#[derive(Debug, Clone)]
477pub(crate) enum RouteRelease {
478    Removed(GoodbyeTarget),
479    Stale,
480    Absent,
481}
482
483#[derive(Debug, Clone)]
484pub(crate) enum RoutePollSnapshot {
485    Bound {
486        module_id: String,
487        status: Option<String>,
488    },
489    Absent,
490}
491
492#[derive(Debug, Clone)]
493struct ModuleConnection {
494    endpoint: ModuleEndpointId,
495    sink: FrameSink,
496    negotiated_ver: u8,
497    concurrency: Concurrency,
498}
499
500#[derive(Debug, Default)]
501struct ForwardingInner {
502    operator_confirms: Arc<crate::operator_confirm::OperatorConfirms>,
503    daemon_draining: bool,
504    /// Whether forwarding reached quiescence during daemon shutdown, for each
505    /// active endpoint present when shutdown began. A pending wait is false
506    /// because it has not proved quiescence. Endpoint generations prevent a
507    /// result from describing a replacement connection with the same module id.
508    daemon_drain_results: HashMap<ModuleEndpointId, bool>,
509    /// The ACTIVE slot: the one endpoint per module id that routing resolves.
510    /// Every by-id lookup (relay reservation, drain-by-id, liveness, census,
511    /// live roots, module-control RPCs) reads this map and nothing else.
512    modules_by_id: HashMap<String, ModuleConnection>,
513    /// The CANDIDATE slot of a blue/green swap: a second process registered
514    /// under an id that already has an active endpoint. It has a full endpoint
515    /// identity (so its own connection can be looked up and torn down) but no
516    /// by-id lookup sees it, so nothing is routed to it until `cutover_candidate`
517    /// promotes it.
518    candidates_by_id: HashMap<String, ModuleConnection>,
519    /// Former active endpoints demoted by `cutover_candidate`, until their
520    /// connection is removed. Membership is what distinguishes an endpoint that
521    /// was SUPERSEDED by a promotion (still alive, still carrying bound routes,
522    /// about to be drained) from one that is merely STALE (replaced some other
523    /// way, which is treated as a fault on the acking connection).
524    superseded_endpoints: HashMap<ModuleEndpointId, ModuleConnection>,
525    endpoint_by_connection: HashMap<ConnectionId, ModuleEndpointId>,
526    module_id_by_endpoint: HashMap<ModuleEndpointId, String>,
527    /// Endpoints mid-drain, keyed to the reason the drain was begun with. The
528    /// value serves the census ("draining because restart"); membership alone
529    /// still answers every admission-gate check.
530    draining_endpoints: HashMap<ModuleEndpointId, RouteCloseReason>,
531    closing_connections: HashSet<ConnectionId>,
532    next_generation: u64,
533    reserved_client: HashMap<ClientRouteKey, ModuleRouteKey>,
534    reserved_module: HashMap<ModuleRouteKey, ClientRouteKey>,
535    next_client_channel: HashMap<ConnectionId, u16>,
536    next_module_channel: HashMap<ModuleEndpointId, u16>,
537    client_slot_epochs: HashMap<ClientRouteKey, u32>,
538    module_slot_epochs: HashMap<ModuleRouteKey, u32>,
539    last_published_epoch: HashMap<ClientRouteKey, u32>,
540    client_to_module: HashMap<ClientRouteKey, Arc<RouteBinding>>,
541    module_to_client: HashMap<ModuleRouteKey, Arc<RouteBinding>>,
542    status: HashMap<(ClientRouteKey, u32), String>,
543    pending_relays: HashMap<(ModuleEndpointId, u64), PendingRouteBindRelayEntry>,
544    /// Pending relays removed without a module reply before `route.open`
545    /// marked them enqueued. `route.open` consumes its key when it marks the
546    /// relay (and sends the GOODBYE the removal could not send safely) or when
547    /// it aborts, so an entry lives only while that open is mid-send.
548    relays_released_before_enqueue: HashSet<(ModuleEndpointId, u64)>,
549    /// Each scope's current `(scope_epoch, version)`, keyed `(owner, ref)`,
550    /// published by a scope sync in the same step that changes the record. A
551    /// bind commit reads it here so it never takes the scope table's lock
552    /// while holding this one. Absent means the scope is not live.
553    scope_tags: HashMap<(String, String), ScopeTag>,
554    next_control_corr: HashMap<ModuleEndpointId, u64>,
555    pending_control_rpcs: HashMap<(ModuleEndpointId, u64), PendingModuleControlRpcEntry>,
556    health_probe_tombstones: HashMap<(ModuleEndpointId, u64), HealthProbeTombstone>,
557}
558
559#[derive(Debug, Clone)]
560pub(crate) struct CloseReason {
561    code: &'static str,
562    message: String,
563}
564
565impl CloseReason {
566    pub(crate) fn new(code: &'static str, message: impl Into<String>) -> Self {
567        Self {
568            code,
569            message: message.into(),
570        }
571    }
572}
573
574impl fmt::Display for CloseReason {
575    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
576        write!(f, "{}: {}", self.code, self.message)
577    }
578}
579
580pub(crate) type ConnectionCloseReceiver = oneshot::Receiver<CloseReason>;
581
582/// Dynamic forwarding state shared by the control plane and data-plane router.
583#[derive(Debug, Default)]
584pub struct ForwardingTable {
585    inner: Arc<RwLock<ForwardingInner>>,
586    #[cfg(feature = "test-support")]
587    route_closing_hold: Mutex<Option<oneshot::Receiver<()>>>,
588    close_registry: Mutex<HashMap<ConnectionId, oneshot::Sender<CloseReason>>>,
589    counters: DaemonCounters,
590    /// Per-target-module bind-relay breaker state. It lives here, beside the
591    /// module connections it describes, because this is where a module
592    /// connection's identity is established and therefore where a stale
593    /// verdict has to be discarded.
594    route_bind_breakers: RouteBindBreakers,
595    /// Current route.bind relays keyed by target module. Admission is shared
596    /// across every client connection that points at the same endpoint.
597    route_bind_concurrency: RouteBindConcurrency,
598    /// Start and end of each module's `route.open` outage. Held here for the
599    /// same reason as the breakers: every handler built over this table must
600    /// see one outage per module, or each would log its own opening line.
601    route_outages: Arc<crate::route_outage::RouteOutageTracker>,
602}
603
604impl ForwardingTable {
605    /// Hold the next supervisor drain between enqueueing `module.draining` and
606    /// `route.closing`, so a test can forward a response in that window. Dropping
607    /// the returned sender also releases the hold, including on a test panic.
608    #[cfg(feature = "test-support")]
609    pub fn hold_route_closing_for_test(&self) -> oneshot::Sender<()> {
610        let (release, hold) = oneshot::channel();
611        *self.route_closing_hold.lock().unwrap() = Some(hold);
612        release
613    }
614
615    #[cfg(feature = "test-support")]
616    pub(crate) async fn wait_before_route_closing_for_test(&self) {
617        let hold = self.route_closing_hold.lock().unwrap().take();
618        if let Some(hold) = hold {
619            let _ = hold.await;
620        }
621    }
622
623    pub(crate) fn operator_confirms(&self) -> Arc<crate::operator_confirm::OperatorConfirms> {
624        Arc::clone(
625            &self
626                .inner
627                .read()
628                .unwrap_or_else(|p| p.into_inner())
629                .operator_confirms,
630        )
631    }
632
633    #[cfg(test)]
634    pub(crate) fn inject_operator_principal(&self, key: ModuleRouteKey, principal: Principal) {
635        let mut inner = self.write_inner().unwrap();
636        let old = inner.module_to_client.remove(&key).unwrap();
637        let client = ClientRouteKey {
638            connection_id: old.client_connection_id,
639            channel: old.client_channel,
640        };
641        inner.client_to_module.remove(&client);
642        let mut binding = Arc::try_unwrap(old).expect("test binding has no outstanding readers");
643        binding.principal = principal;
644        let binding = Arc::new(binding);
645        inner.client_to_module.insert(client, Arc::clone(&binding));
646        inner.module_to_client.insert(key, binding);
647    }
648
649    /// Keep the binding protected through confirmation admission. Route release
650    /// takes the write lock and then the same confirm-state lock, so it cannot
651    /// miss an admitted request or leave it waiting on a vanished route.
652    pub(crate) fn with_operator_route<T>(
653        &self,
654        connection_id: ConnectionId,
655        channel: u16,
656        epoch: u32,
657        admit: impl FnOnce(Option<&RouteBinding>) -> T,
658    ) -> Result<T, ForwardingError> {
659        let inner = self.read_inner()?;
660        let binding = inner
661            .endpoint_by_connection
662            .get(&connection_id)
663            .and_then(|endpoint| {
664                inner.module_to_client.get(&ModuleRouteKey {
665                    endpoint: *endpoint,
666                    channel,
667                })
668            })
669            .filter(|binding| binding.module_epoch == epoch);
670        Ok(admit(binding.map(Arc::as_ref)))
671    }
672
673    pub(crate) fn counters(&self) -> DaemonCounters {
674        self.counters.clone()
675    }
676
677    pub(crate) fn route_bind_breakers(&self) -> RouteBindBreakers {
678        self.route_bind_breakers.clone()
679    }
680
681    pub(crate) fn route_bind_concurrency(&self) -> RouteBindConcurrency {
682        self.route_bind_concurrency.clone()
683    }
684
685    pub(crate) fn route_outages(&self) -> Arc<crate::route_outage::RouteOutageTracker> {
686        Arc::clone(&self.route_outages)
687    }
688
689    pub(crate) fn register_connection_close(
690        &self,
691        connection_id: ConnectionId,
692    ) -> ConnectionCloseReceiver {
693        let (sender, receiver) = oneshot::channel();
694        let replaced = self
695            .lock_close_registry()
696            .insert(connection_id, sender)
697            .is_some();
698        if replaced {
699            warn!(
700                connection_id = connection_id.get(),
701                "replaced existing connection close registration"
702            );
703        }
704        receiver
705    }
706
707    pub(crate) fn unregister_connection_close(&self, connection_id: ConnectionId) {
708        self.lock_close_registry().remove(&connection_id);
709    }
710
711    /// Close every established connection, modules and clients alike, as the
712    /// last step of an announced daemon shutdown. By the time this runs each
713    /// registered module has already been sent a module GOODBYE (see
714    /// `Supervisor::end_children_for_daemon_shutdown`), so the EOF that follows
715    /// is a planned stop; EOF with no GOODBYE before it stays reserved for a
716    /// daemon that went away unannounced. Closing while the daemon is still
717    /// alive lets it wait for the modules' own teardowns before it ends
718    /// whatever is left.
719    #[cfg(unix)]
720    pub(crate) fn close_all_connections(&self, reason: &CloseReason) -> usize {
721        let senders: Vec<_> = self.lock_close_registry().drain().collect();
722        let count = senders.len();
723        for (_, sender) in senders {
724            let _ = sender.send(reason.clone());
725        }
726        count
727    }
728
729    /// Every registered module connection, whichever slot it occupies (active,
730    /// swap candidate, or superseded incumbent), once each. Daemon shutdown
731    /// uses this to tell each of them it is a planned stop before closing it.
732    #[cfg(unix)]
733    pub(crate) fn module_connections(
734        &self,
735    ) -> Result<Vec<ModuleConnectionTarget>, ForwardingError> {
736        let inner = self.read_inner()?;
737        let mut seen = HashSet::new();
738        Ok(inner
739            .modules_by_id
740            .values()
741            .chain(inner.candidates_by_id.values())
742            .chain(inner.superseded_endpoints.values())
743            .filter(|module| seen.insert(module.endpoint))
744            .map(|module| ModuleConnectionTarget {
745                module_id: inner
746                    .module_id_by_endpoint
747                    .get(&module.endpoint)
748                    .cloned()
749                    .unwrap_or_default(),
750                endpoint: module.endpoint,
751                sink: module.sink.clone(),
752                negotiated_ver: module.negotiated_ver,
753            })
754            .collect())
755    }
756
757    /// Ask a registered connection to close. Returns true only for the request
758    /// that actually reached it; later requests for the same connection, and
759    /// requests for a connection that is not registered, return false.
760    pub(crate) fn request_connection_close(
761        &self,
762        connection_id: ConnectionId,
763        reason: CloseReason,
764    ) -> bool {
765        let sender = self.lock_close_registry().remove(&connection_id);
766        if let Some(sender) = sender {
767            debug!(
768                connection_id = connection_id.get(),
769                close_reason = %reason,
770                "requesting connection close"
771            );
772            let _ = sender.send(reason);
773            true
774        } else {
775            debug!(
776                connection_id = connection_id.get(),
777                close_reason = %reason,
778                "connection close request ignored for inactive connection"
779            );
780            false
781        }
782    }
783
784    pub fn register_module_connection(
785        &self,
786        connection_id: ConnectionId,
787        module_id: String,
788        negotiated_ver: u8,
789        concurrency: Concurrency,
790        sink: FrameSink,
791    ) -> Result<ModuleEndpointId, ForwardingError> {
792        self.register_module_connection_inner(
793            connection_id,
794            module_id,
795            negotiated_ver,
796            concurrency,
797            sink,
798            None,
799        )
800    }
801
802    /// Register a module connection and queue its HELLO_ACK as the first frame
803    /// on its sink, in the same write-lock critical section that makes the
804    /// endpoint visible.
805    ///
806    /// A module reads HELLO_ACK as the first frame after its HELLO and exits on
807    /// anything else. Once the endpoint is in `modules_by_id`, a `route.open`
808    /// on any other connection can queue a `route.bind` request onto this sink,
809    /// so the ack must already be queued by then. Every lookup that can queue a
810    /// frame for the module takes this lock, so nothing can get in ahead of it.
811    ///
812    /// If the ack cannot be queued (the sink is closed or full), the
813    /// registration fails with nothing inserted.
814    pub(crate) fn register_module_connection_acked(
815        &self,
816        connection_id: ConnectionId,
817        module_id: String,
818        negotiated_ver: u8,
819        concurrency: Concurrency,
820        sink: FrameSink,
821        hello_ack: Frame,
822    ) -> Result<ModuleEndpointId, ForwardingError> {
823        self.register_module_connection_inner(
824            connection_id,
825            module_id,
826            negotiated_ver,
827            concurrency,
828            sink,
829            Some(hello_ack),
830        )
831    }
832
833    fn register_module_connection_inner(
834        &self,
835        connection_id: ConnectionId,
836        module_id: String,
837        negotiated_ver: u8,
838        concurrency: Concurrency,
839        sink: FrameSink,
840        hello_ack: Option<Frame>,
841    ) -> Result<ModuleEndpointId, ForwardingError> {
842        let mut inner = self.write_inner()?;
843        if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
844            return Err(ForwardingError::ConnectionClosing { connection_id });
845        }
846        check_module_connection_role_locked(&inner, connection_id)?;
847        // Every refusal check has passed and nothing has been mutated yet, so
848        // a failed enqueue leaves the table exactly as it was.
849        enqueue_hello_ack_locked(&sink, connection_id, hello_ack)?;
850
851        inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
852        let endpoint = ModuleEndpointId {
853            connection_id,
854            generation: inner.next_generation,
855        };
856        inner.endpoint_by_connection.insert(connection_id, endpoint);
857        inner
858            .module_id_by_endpoint
859            .insert(endpoint, module_id.clone());
860        inner.next_module_channel.insert(endpoint, 1);
861        inner.next_control_corr.insert(endpoint, 1);
862        inner.modules_by_id.insert(
863            module_id.clone(),
864            ModuleConnection {
865                endpoint,
866                sink,
867                negotiated_ver,
868                concurrency,
869            },
870        );
871        drop(inner);
872
873        // A new module connection has arrived under this id, so anything the
874        // bind-relay breaker learned was learned about a process that is no
875        // longer the one behind this name. See
876        // `RouteBindBreakers::reset_for_new_module_connection`.
877        //
878        // Keyed on ARRIVAL rather than on teardown deliberately: a connection
879        // going away is not evidence about anything, and a module whose
880        // connection drops without coming back should keep its verdict until
881        // something actually registers in its place. This also covers the
882        // unclean replacements -- a module killed mid-bind, or one whose
883        // connection was closing -- because registration is the single path by
884        // which any module connection becomes usable.
885        if let Some(discarded) = self
886            .route_bind_breakers
887            .reset_for_new_module_connection(&module_id)
888        {
889            info!(
890                module_id = %module_id,
891                discarded_consecutive_timeouts = discarded,
892                "route.bind breaker state discarded: a new module connection replaced the process it described"
893            );
894        }
895        Ok(endpoint)
896    }
897
898    /// Register a blue/green swap candidate for `module_id` into the candidate
899    /// slot, alongside whatever endpoint is active for the id.
900    ///
901    /// Unlike [`Self::register_module_connection`], this never touches the
902    /// active slot and never resets the bind-relay breaker: the incumbent is
903    /// still the process serving the id, so what the breaker learned about it
904    /// is still true. The breaker is reset at cutover instead, when the process
905    /// behind the name actually changes. A second candidate for the same id is
906    /// refused.
907    ///
908    /// Production registers candidates through
909    /// [`Self::register_candidate_module_connection_acked`]; this ack-less form
910    /// is for tests that build forwarding state directly.
911    #[cfg(test)]
912    pub(crate) fn register_candidate_module_connection(
913        &self,
914        connection_id: ConnectionId,
915        module_id: String,
916        negotiated_ver: u8,
917        concurrency: Concurrency,
918        sink: FrameSink,
919    ) -> Result<ModuleEndpointId, ForwardingError> {
920        self.register_candidate_module_connection_inner(
921            connection_id,
922            module_id,
923            negotiated_ver,
924            concurrency,
925            sink,
926            None,
927        )
928    }
929
930    /// The swap-candidate counterpart of
931    /// [`Self::register_module_connection_acked`]: the HELLO_ACK is queued on
932    /// the candidate's sink before its endpoint is inserted. No by-id lookup
933    /// sees a candidate, but its own connection's endpoint does become
934    /// resolvable here, and at cutover it becomes routable; the ack has to be
935    /// ahead of anything either can queue.
936    pub(crate) fn register_candidate_module_connection_acked(
937        &self,
938        connection_id: ConnectionId,
939        module_id: String,
940        negotiated_ver: u8,
941        concurrency: Concurrency,
942        sink: FrameSink,
943        hello_ack: Frame,
944    ) -> Result<ModuleEndpointId, ForwardingError> {
945        self.register_candidate_module_connection_inner(
946            connection_id,
947            module_id,
948            negotiated_ver,
949            concurrency,
950            sink,
951            Some(hello_ack),
952        )
953    }
954
955    fn register_candidate_module_connection_inner(
956        &self,
957        connection_id: ConnectionId,
958        module_id: String,
959        negotiated_ver: u8,
960        concurrency: Concurrency,
961        sink: FrameSink,
962        hello_ack: Option<Frame>,
963    ) -> Result<ModuleEndpointId, ForwardingError> {
964        let mut inner = self.write_inner()?;
965        if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
966            return Err(ForwardingError::ConnectionClosing { connection_id });
967        }
968        check_module_connection_role_locked(&inner, connection_id)?;
969        if inner.candidates_by_id.contains_key(&module_id) {
970            return Err(ForwardingError::CandidateSlotOccupied { module_id });
971        }
972        // As in `register_module_connection_inner`: after every refusal check,
973        // before any mutation.
974        enqueue_hello_ack_locked(&sink, connection_id, hello_ack)?;
975
976        inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
977        let endpoint = ModuleEndpointId {
978            connection_id,
979            generation: inner.next_generation,
980        };
981        inner.endpoint_by_connection.insert(connection_id, endpoint);
982        inner
983            .module_id_by_endpoint
984            .insert(endpoint, module_id.clone());
985        inner.next_module_channel.insert(endpoint, 1);
986        inner.next_control_corr.insert(endpoint, 1);
987        inner.candidates_by_id.insert(
988            module_id,
989            ModuleConnection {
990                endpoint,
991                sink,
992                negotiated_ver,
993                concurrency,
994            },
995        );
996        Ok(endpoint)
997    }
998
999    /// Promote `module_id`'s candidate to the active slot, in ONE forwarding
1000    /// write-lock critical section.
1001    ///
1002    /// That single section is the linearization point of a swap. Every relay
1003    /// reservation resolves the active slot under the same lock, so each one
1004    /// lands wholly before cutover (on the incumbent) or wholly after it (on the
1005    /// promoted candidate), never on a mix. A relay reserved on the incumbent
1006    /// before cutover can no longer commit: `commit_route_locked` requires the
1007    /// reservation's endpoint to be the active one, and an ack for it arriving
1008    /// later is answered by the superseded-endpoint arm of
1009    /// `complete_pending_relay` instead of ending the incumbent's connection.
1010    ///
1011    /// Both endpoints keep their identities: nothing is re-keyed, so the
1012    /// incumbent's bound routes and its pending correlation keys stay exactly
1013    /// where they are until it is drained with [`Self::begin_endpoint_drain`],
1014    /// using the incumbent endpoint this returns. Returns `Ok(None)` when there
1015    /// is no candidate for the id.
1016    pub(crate) fn cutover_candidate(
1017        &self,
1018        module_id: &str,
1019    ) -> Result<Option<ForwardingCutover>, ForwardingError> {
1020        let mut inner = self.write_inner()?;
1021        if inner.daemon_draining {
1022            return Err(ForwardingError::ModuleReloading {
1023                module_id: module_id.to_string(),
1024            });
1025        }
1026        let Some(candidate) = inner.candidates_by_id.remove(module_id) else {
1027            return Ok(None);
1028        };
1029        let promoted = candidate.endpoint;
1030        let incumbent = inner.modules_by_id.insert(module_id.to_string(), candidate);
1031        let incumbent = incumbent.map(|incumbent| {
1032            let endpoint = incumbent.endpoint;
1033            inner.superseded_endpoints.insert(endpoint, incumbent);
1034            endpoint
1035        });
1036        drop(inner);
1037
1038        // A different process now answers for this id; see the matching reset
1039        // in `register_module_connection` for why a verdict about the old one
1040        // must not carry over.
1041        if let Some(discarded) = self
1042            .route_bind_breakers
1043            .reset_for_new_module_connection(module_id)
1044        {
1045            info!(
1046                module_id = %module_id,
1047                discarded_consecutive_timeouts = discarded,
1048                "route.bind breaker state discarded: a swap candidate was promoted over the process it described"
1049            );
1050        }
1051        Ok(Some(ForwardingCutover {
1052            promoted,
1053            incumbent,
1054        }))
1055    }
1056
1057    #[allow(clippy::too_many_arguments)]
1058    pub(crate) async fn begin_route_bind_relay_for(
1059        &self,
1060        client_connection_id: ConnectionId,
1061        client_sink: FrameSink,
1062        client_negotiated_ver: u8,
1063        client_corr: u64,
1064        module_id: &str,
1065        principal: Principal,
1066        scope: Option<BoundScope>,
1067        project_root: Option<ProjectRootId>,
1068        deadline: Instant,
1069    ) -> Result<PendingRouteBindRelay, ForwardingError> {
1070        // Reserve egress capacity before taking the forwarding lock. The permit is
1071        // held until the bind reaches one terminal state, so an accepted bind can
1072        // publish its table entry and RouteOpen response in one critical section.
1073        let client_permit =
1074            client_sink
1075                .reserve_owned()
1076                .await
1077                .map_err(|_| ForwardingError::ClientEgressClosed {
1078                    connection_id: client_connection_id,
1079                })?;
1080        self.begin_route_bind_relay_inner(
1081            client_connection_id,
1082            client_sink,
1083            client_negotiated_ver,
1084            client_corr,
1085            module_id,
1086            principal,
1087            scope,
1088            project_root,
1089            deadline,
1090            client_permit,
1091        )
1092    }
1093
1094    #[cfg(test)]
1095    pub(crate) fn begin_route_bind_relay_for_test(
1096        &self,
1097        client_connection_id: ConnectionId,
1098        client_sink: FrameSink,
1099        client_corr: u64,
1100        module_id: &str,
1101    ) -> Result<PendingRouteBindRelay, ForwardingError> {
1102        let permit =
1103            client_sink
1104                .try_reserve_owned()
1105                .map_err(|_| ForwardingError::ClientEgressClosed {
1106                    connection_id: client_connection_id,
1107                })?;
1108        self.begin_route_bind_relay_inner(
1109            client_connection_id,
1110            client_sink,
1111            subc_protocol::PROTOCOL_VERSION,
1112            client_corr,
1113            module_id,
1114            Principal::Direct,
1115            None,
1116            None,
1117            Instant::now() + std::time::Duration::from_secs(60),
1118            permit,
1119        )
1120    }
1121
1122    pub(crate) fn begin_module_control_rpc_for(
1123        &self,
1124        module_id: &str,
1125        expected_op: &str,
1126        deadline: Instant,
1127    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1128        self.begin_module_control_rpc_inner(module_id, expected_op, deadline, None, false)
1129    }
1130
1131    pub(crate) fn begin_health_probe_rpc_for(
1132        &self,
1133        module_id: &str,
1134        expected_op: &str,
1135        probe_started_at: Instant,
1136        deadline: Instant,
1137    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1138        self.begin_module_control_rpc_inner(
1139            module_id,
1140            expected_op,
1141            deadline,
1142            Some(probe_started_at),
1143            false,
1144        )
1145    }
1146
1147    pub(crate) fn begin_drain_health_probe_rpc_for(
1148        &self,
1149        module_id: &str,
1150        expected_op: &str,
1151        probe_started_at: Instant,
1152        deadline: Instant,
1153    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1154        self.begin_module_control_rpc_inner(
1155            module_id,
1156            expected_op,
1157            deadline,
1158            Some(probe_started_at),
1159            true,
1160        )
1161    }
1162
1163    /// A health probe addressed to one endpoint, in whichever slot it is.
1164    ///
1165    /// Every other control RPC resolves the id's ACTIVE endpoint, which is how
1166    /// a swap candidate (not active until cutover) and a superseded incumbent
1167    /// (not active after it) are unreachable by them. A swap needs to probe
1168    /// exactly those two: the candidate before promoting it, and the incumbent
1169    /// for its busy gauges while it drains.
1170    pub(crate) fn begin_endpoint_health_probe_rpc_for(
1171        &self,
1172        endpoint: ModuleEndpointId,
1173        expected_op: &str,
1174        probe_started_at: Instant,
1175        deadline: Instant,
1176    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1177        let inner = self.write_inner()?;
1178        let module = module_connection_for_endpoint_locked(&inner, endpoint)
1179            .cloned()
1180            .ok_or(ForwardingError::NoModuleConnection)?;
1181        let module_id = inner
1182            .module_id_by_endpoint
1183            .get(&endpoint)
1184            .cloned()
1185            .unwrap_or_default();
1186        // A draining incumbent is exactly what this probes for busy gauges, so
1187        // draining is allowed, as it is for the by-id drain probe.
1188        self.begin_control_rpc_locked(
1189            inner,
1190            &module_id,
1191            module,
1192            expected_op,
1193            deadline,
1194            Some(probe_started_at),
1195            true,
1196        )
1197    }
1198
1199    fn begin_module_control_rpc_inner(
1200        &self,
1201        module_id: &str,
1202        expected_op: &str,
1203        deadline: Instant,
1204        health_probe_started_at: Option<Instant>,
1205        allow_draining: bool,
1206    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1207        let inner = self.write_inner()?;
1208        let module = inner
1209            .modules_by_id
1210            .get(module_id)
1211            .cloned()
1212            .ok_or(ForwardingError::NoModuleConnection)?;
1213        self.begin_control_rpc_locked(
1214            inner,
1215            module_id,
1216            module,
1217            expected_op,
1218            deadline,
1219            health_probe_started_at,
1220            allow_draining,
1221        )
1222    }
1223
1224    #[allow(clippy::too_many_arguments)]
1225    fn begin_control_rpc_locked(
1226        &self,
1227        mut inner: RwLockWriteGuard<'_, ForwardingInner>,
1228        module_id: &str,
1229        module: ModuleConnection,
1230        expected_op: &str,
1231        deadline: Instant,
1232        health_probe_started_at: Option<Instant>,
1233        allow_draining: bool,
1234    ) -> Result<PendingModuleControlRpc, ForwardingError> {
1235        if !allow_draining && inner.draining_endpoints.contains_key(&module.endpoint) {
1236            return Err(ForwardingError::ModuleReloading {
1237                module_id: module_id.to_string(),
1238            });
1239        }
1240        if inner
1241            .closing_connections
1242            .contains(&module.endpoint.connection_id)
1243        {
1244            return Err(ForwardingError::ConnectionClosing {
1245                connection_id: module.endpoint.connection_id,
1246            });
1247        }
1248        if health_probe_started_at.is_some() {
1249            // Recurring health probes are serial per endpoint. Once the next one
1250            // starts, an older answer can no longer improve the current snapshot,
1251            // so retaining more than the newest unanswered probe has no value.
1252            inner
1253                .health_probe_tombstones
1254                .retain(|(endpoint, _), _| *endpoint != module.endpoint);
1255        }
1256        let corr = match inner.allocate_control_corr(module.endpoint) {
1257            Ok(corr) => corr,
1258            Err(err) => {
1259                drop(inner);
1260                self.request_connection_close(
1261                    module.endpoint.connection_id,
1262                    CloseReason::new(
1263                        "control_correlation_exhausted",
1264                        "daemon-originated channel-0 correlation space exhausted",
1265                    ),
1266                );
1267                return Err(err);
1268            }
1269        };
1270        let (sender, receiver) = oneshot::channel();
1271        inner.pending_control_rpcs.insert(
1272            (module.endpoint, corr),
1273            PendingModuleControlRpcEntry {
1274                expected_op: expected_op.to_string(),
1275                deadline,
1276                health_probe_started_at,
1277                sender,
1278            },
1279        );
1280
1281        Ok(PendingModuleControlRpc {
1282            endpoint: module.endpoint,
1283            module_sink: module.sink,
1284            negotiated_ver: module.negotiated_ver,
1285            corr,
1286            receiver,
1287        })
1288    }
1289
1290    #[allow(clippy::too_many_arguments)]
1291    fn begin_route_bind_relay_inner(
1292        &self,
1293        client_connection_id: ConnectionId,
1294        client_sink: FrameSink,
1295        client_negotiated_ver: u8,
1296        client_corr: u64,
1297        expected_module_id: &str,
1298        principal: Principal,
1299        scope: Option<BoundScope>,
1300        project_root: Option<ProjectRootId>,
1301        deadline: Instant,
1302        client_permit: crate::router::EgressPermit,
1303    ) -> Result<PendingRouteBindRelay, ForwardingError> {
1304        let mut inner = self.write_inner()?;
1305        // HELLO and route.open can run on different tasks. Decide the role under
1306        // the same lock that installs endpoints and reserves client channels so
1307        // neither task can turn one connection into both kinds of owner.
1308        if inner
1309            .endpoint_by_connection
1310            .contains_key(&client_connection_id)
1311        {
1312            return Err(ForwardingError::ConnectionRoleConflict {
1313                connection_id: client_connection_id,
1314            });
1315        }
1316        if inner.closing_connections.contains(&client_connection_id) {
1317            return Err(ForwardingError::ConnectionClosing {
1318                connection_id: client_connection_id,
1319            });
1320        }
1321        let module = inner
1322            .modules_by_id
1323            .get(expected_module_id)
1324            .cloned()
1325            .ok_or(ForwardingError::NoModuleConnection)?;
1326        if inner.draining_endpoints.contains_key(&module.endpoint) {
1327            return Err(ForwardingError::ModuleReloading {
1328                module_id: expected_module_id.to_string(),
1329            });
1330        }
1331        if inner
1332            .closing_connections
1333            .contains(&module.endpoint.connection_id)
1334        {
1335            return Err(ForwardingError::ConnectionClosing {
1336                connection_id: module.endpoint.connection_id,
1337            });
1338        }
1339
1340        let corr = match inner.allocate_control_corr(module.endpoint) {
1341            Ok(corr) => corr,
1342            Err(err) => {
1343                drop(inner);
1344                self.request_connection_close(
1345                    module.endpoint.connection_id,
1346                    CloseReason::new(
1347                        "control_correlation_exhausted",
1348                        "daemon-originated channel-0 correlation space exhausted",
1349                    ),
1350                );
1351                return Err(err);
1352            }
1353        };
1354        let (client_channel, client_epoch, module_channel, module_epoch) =
1355            inner.allocate_route_slots(client_connection_id, module.endpoint)?;
1356        let client_key = ClientRouteKey {
1357            connection_id: client_connection_id,
1358            channel: client_channel,
1359        };
1360        let module_key = ModuleRouteKey {
1361            endpoint: module.endpoint,
1362            channel: module_channel,
1363        };
1364        let reservation = RouteReservation {
1365            client_key,
1366            module_key,
1367            client_epoch,
1368            module_epoch,
1369            project_root,
1370        };
1371        let response_body = serde_json::to_vec(&ClientControlResponse::RouteOpen {
1372            route_channel: client_channel,
1373            route_epoch: client_epoch,
1374        })
1375        .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
1376        let route_open_frame = Frame::build_with_version(
1377            client_negotiated_ver,
1378            FrameType::Response,
1379            Flags::new(false, Priority::Passive, false),
1380            0,
1381            0,
1382            client_corr,
1383            response_body,
1384        )
1385        .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
1386        let (sender, receiver) = oneshot::channel();
1387        inner.reserved_client.insert(client_key, module_key);
1388        inner.reserved_module.insert(module_key, client_key);
1389        inner.pending_relays.insert(
1390            (module.endpoint, corr),
1391            PendingRouteBindRelayEntry {
1392                reservation,
1393                client_sink,
1394                client_negotiated_ver,
1395                client_permit,
1396                route_open_frame,
1397                principal,
1398                scope,
1399                deadline,
1400                relay_enqueued: false,
1401                sender,
1402            },
1403        );
1404
1405        Ok(PendingRouteBindRelay {
1406            endpoint: module.endpoint,
1407            module_sink: module.sink,
1408            negotiated_ver: module.negotiated_ver,
1409            client_channel,
1410            client_epoch,
1411            module_channel,
1412            module_epoch,
1413            corr,
1414            receiver,
1415        })
1416    }
1417
1418    /// Called by `route.open` right after the route.bind entered the module's
1419    /// egress queue.
1420    ///
1421    /// The module reader runs concurrently with `route.open`, so a fast module
1422    /// can receive the bind, ack it, and have the reader commit the route
1423    /// before this runs. A missing entry therefore does not mean the bind was
1424    /// abandoned: only a removal recorded in `relays_released_before_enqueue`
1425    /// does. Telling the two apart is what stops `route.open` from sending a
1426    /// GOODBYE for a route that was just committed.
1427    pub(crate) fn mark_route_bind_relay_enqueued(
1428        &self,
1429        endpoint: ModuleEndpointId,
1430        corr: u64,
1431    ) -> Result<RouteBindRelayMark, ForwardingError> {
1432        let mut inner = self.write_inner()?;
1433        if let Some(pending) = inner.pending_relays.get_mut(&(endpoint, corr)) {
1434            pending.relay_enqueued = true;
1435            return Ok(RouteBindRelayMark::Marked);
1436        }
1437        if inner
1438            .relays_released_before_enqueue
1439            .remove(&(endpoint, corr))
1440        {
1441            return Ok(RouteBindRelayMark::ReleasedWithoutReply);
1442        }
1443        Ok(RouteBindRelayMark::SettledByReply)
1444    }
1445
1446    pub(crate) fn release_client_route(
1447        &self,
1448        client_connection_id: ConnectionId,
1449        client_channel: u16,
1450        expected_epoch: u32,
1451    ) -> Result<RouteRelease, ForwardingError> {
1452        let mut inner = self.write_inner()?;
1453        let release = release_client_route_locked(
1454            &mut inner,
1455            ClientRouteKey {
1456                connection_id: client_connection_id,
1457                channel: client_channel,
1458            },
1459            expected_epoch,
1460        );
1461        self.record_route_release(&release);
1462        Ok(release)
1463    }
1464
1465    pub(crate) fn release_module_route(
1466        &self,
1467        module_connection_id: ConnectionId,
1468        module_channel: u16,
1469        expected_epoch: u32,
1470    ) -> Result<RouteRelease, ForwardingError> {
1471        let mut inner = self.write_inner()?;
1472        let Some(endpoint) = inner
1473            .endpoint_by_connection
1474            .get(&module_connection_id)
1475            .copied()
1476        else {
1477            return Ok(RouteRelease::Absent);
1478        };
1479        let release = release_module_route_locked(
1480            &mut inner,
1481            ModuleRouteKey {
1482                endpoint,
1483                channel: module_channel,
1484            },
1485            expected_epoch,
1486        );
1487        self.record_route_release(&release);
1488        Ok(release)
1489    }
1490
1491    pub(crate) fn abort_pending_relay(
1492        &self,
1493        endpoint: ModuleEndpointId,
1494        corr: u64,
1495        outcome: RouteBindRelayOutcome,
1496    ) -> Result<Option<GoodbyeTarget>, ForwardingError> {
1497        let mut inner = self.write_inner()?;
1498        let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1499            // If another path dropped the relay while `route.open` was still
1500            // sending the bind, and that open is now giving up without marking
1501            // it (its send failed or it was cancelled), forget the record so it
1502            // cannot outlive the open.
1503            inner
1504                .relays_released_before_enqueue
1505                .remove(&(endpoint, corr));
1506            return Ok(None);
1507        };
1508        release_reserved_route_locked(
1509            &mut inner,
1510            pending.reservation.client_key,
1511            pending.reservation.module_key,
1512        );
1513        let target = pending
1514            .relay_enqueued
1515            .then(|| abandoned_route_target(&inner, &pending.reservation));
1516        let _ = pending.sender.send(outcome);
1517        Ok(target.flatten())
1518    }
1519
1520    pub(crate) fn cancel_module_control_rpc(
1521        &self,
1522        endpoint: ModuleEndpointId,
1523        corr: u64,
1524    ) -> Result<(), ForwardingError> {
1525        self.write_inner()?
1526            .pending_control_rpcs
1527            .remove(&(endpoint, corr));
1528        Ok(())
1529    }
1530
1531    pub(crate) fn tombstone_health_probe_rpc(
1532        &self,
1533        endpoint: ModuleEndpointId,
1534        corr: u64,
1535    ) -> Result<bool, ForwardingError> {
1536        let key = (endpoint, corr);
1537        let expires_at = Instant::now() + HEALTH_PROBE_TOMBSTONE_TTL;
1538        {
1539            let mut inner = self.write_inner()?;
1540            let Some(pending) = inner.pending_control_rpcs.remove(&key) else {
1541                return Ok(false);
1542            };
1543            let Some(probe_started_at) = pending.health_probe_started_at else {
1544                inner.pending_control_rpcs.insert(key, pending);
1545                return Ok(false);
1546            };
1547            let module_id = inner
1548                .module_id_by_endpoint
1549                .get(&endpoint)
1550                .cloned()
1551                .unwrap_or_else(|| "unknown".to_string());
1552            inner.health_probe_tombstones.insert(
1553                key,
1554                HealthProbeTombstone {
1555                    expected_op: pending.expected_op,
1556                    module_id,
1557                    probe_started_at,
1558                    expires_at,
1559                },
1560            );
1561        }
1562        self.schedule_health_probe_tombstone_expiration(key, expires_at);
1563        Ok(true)
1564    }
1565
1566    fn schedule_health_probe_tombstone_expiration(
1567        &self,
1568        key: (ModuleEndpointId, u64),
1569        expires_at: Instant,
1570    ) {
1571        let inner = Arc::downgrade(&self.inner);
1572        tokio::spawn(async move {
1573            tokio::time::sleep_until(expires_at).await;
1574            let Some(inner) = inner.upgrade() else {
1575                return;
1576            };
1577            let Ok(mut inner) = inner.write() else {
1578                return;
1579            };
1580            let expired = inner
1581                .health_probe_tombstones
1582                .get(&key)
1583                .is_some_and(|tombstone| tombstone.expires_at <= Instant::now());
1584            if expired {
1585                inner.health_probe_tombstones.remove(&key);
1586            }
1587        });
1588    }
1589
1590    pub(crate) fn complete_pending_relay(
1591        &self,
1592        connection_id: ConnectionId,
1593        corr: u64,
1594        outcome: RouteBindRelayOutcome,
1595    ) -> Result<PendingRelayCompletion, ForwardingError> {
1596        let mut inner = self.write_inner()?;
1597        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1598            return Ok(PendingRelayCompletion {
1599                settled: false,
1600                abandoned: None,
1601            });
1602        };
1603        let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1604            return Ok(PendingRelayCompletion {
1605                settled: false,
1606                abandoned: None,
1607            });
1608        };
1609
1610        if Instant::now() >= pending.deadline {
1611            release_reserved_route_locked(
1612                &mut inner,
1613                pending.reservation.client_key,
1614                pending.reservation.module_key,
1615            );
1616            let abandoned = matches!(outcome, RouteBindRelayOutcome::Accepted)
1617                .then(|| abandoned_route_target(&inner, &pending.reservation))
1618                .flatten();
1619            let _ = pending
1620                .sender
1621                .send(RouteBindRelayOutcome::Rejected(ErrorBody {
1622                    code: "module_timeout".to_string(),
1623                    message: "route.bind response arrived after its daemon deadline".to_string(),
1624                    detail: None,
1625                }));
1626            return Ok(PendingRelayCompletion {
1627                settled: true,
1628                abandoned,
1629            });
1630        }
1631
1632        match outcome {
1633            // Two shapes of "the client is not there to receive this route" that
1634            // must resolve identically: its egress is already closed, or it is
1635            // marked closing (its connection loop has been asked to end, but has
1636            // not drained yet, so the sink is still open).
1637            //
1638            // Only the first used to be caught here. The second fell through to
1639            // `commit_route_locked`, which refuses a closing client with
1640            // `ConnectionClosing` -- and this function is called from the MODULE
1641            // connection's frame handler, so that refusal ended the module's
1642            // connection instead of this one client's route. One dying client's
1643            // route.open then took down a connection carrying every other
1644            // client's routes to that module.
1645            //
1646            // The remedy for both is the same, which is why they share an arm:
1647            // give back the reserved handle pair, tell the waiting route.open the
1648            // route is gone, and report the module-side channel so the caller can
1649            // send a channel-scoped GOODBYE for the binding the module just
1650            // created. Nothing here touches the module connection.
1651            RouteBindRelayOutcome::Accepted
1652                if pending.client_sink.is_closed()
1653                    || inner
1654                        .closing_connections
1655                        .contains(&pending.reservation.client_key.connection_id) =>
1656            {
1657                let reason = if pending.client_sink.is_closed() {
1658                    "client egress closed before route publication"
1659                } else {
1660                    "client connection is closing before route publication"
1661                };
1662                release_reserved_route_locked(
1663                    &mut inner,
1664                    pending.reservation.client_key,
1665                    pending.reservation.module_key,
1666                );
1667                // Not gated on `relay_enqueued`: an ack proves the module
1668                // received the bind, and `route.open` sends nothing for a
1669                // relay a reply already settled, so this is the only GOODBYE.
1670                let abandoned = abandoned_route_target(&inner, &pending.reservation);
1671                let _ = pending
1672                    .sender
1673                    .send(RouteBindRelayOutcome::ModuleGone(reason.to_string()));
1674                return Ok(PendingRelayCompletion {
1675                    settled: true,
1676                    abandoned,
1677                });
1678            }
1679            // The endpoint that acked this bind is no longer accepting new
1680            // routes, for one of two reasons. Either it was the module's active
1681            // endpoint when this relay was reserved and a blue/green swap has
1682            // since promoted a candidate (the replacement process registered
1683            // beside it) into the active slot, superseding it; or a daemon
1684            // drain has gated it, and the ack arrived before the ordered module
1685            // drain settled pending relays. The module did nothing wrong: it
1686            // bound the route it was asked to bind, and the endpoint still
1687            // carries every other client's routes until it is drained. So the
1688            // ack is settled here, while the pending entry still holds the
1689            // client's sender and the reserved handle pair: release the pair,
1690            // refuse the waiting route.open with `module_reloading`, which
1691            // callers retry (after a swap the retry reserves on the promoted
1692            // endpoint), and hand back the module-side channel so the caller
1693            // sends one channel-scoped GOODBYE for the binding the module just
1694            // created. Returning an error instead would end the module
1695            // connection all those other routes share; nothing here touches it.
1696            //
1697            // An endpoint that stopped being active any other way is stale
1698            // (replaced, but not by a swap and not draining). It still falls
1699            // through to `commit_route_locked`, which refuses it with
1700            // `StaleModuleEndpoint` exactly as before swaps existed.
1701            RouteBindRelayOutcome::Accepted
1702                if inner.superseded_endpoints.contains_key(&endpoint)
1703                    || inner.draining_endpoints.contains_key(&endpoint) =>
1704            {
1705                release_reserved_route_locked(
1706                    &mut inner,
1707                    pending.reservation.client_key,
1708                    pending.reservation.module_key,
1709                );
1710                // Not gated on `relay_enqueued`: an ack proves the module
1711                // received the bind, whether or not the enqueue mark was set.
1712                let abandoned = abandoned_route_target(&inner, &pending.reservation);
1713                let module_id = inner
1714                    .module_id_by_endpoint
1715                    .get(&endpoint)
1716                    .cloned()
1717                    .unwrap_or_else(|| "unknown".to_string());
1718                let _ = pending
1719                    .sender
1720                    .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
1721                        "module_reloading",
1722                        format!("module_id '{module_id}' is reloading"),
1723                    )));
1724                return Ok(PendingRelayCompletion {
1725                    settled: true,
1726                    abandoned,
1727                });
1728            }
1729            // The scope the open was admitted under ended or changed after
1730            // admission. The module bound a route under a stamp that is no
1731            // longer true, so it must not become routable. Settled here like
1732            // the superseded arm above, never as an error from
1733            // `commit_route_locked`, which would close the module's whole
1734            // connection and every other client's routes to it: release the
1735            // pair, answer the waiting route.open by name, and hand back the
1736            // module-side channel for one channel-scoped GOODBYE.
1737            RouteBindRelayOutcome::Accepted
1738                if pending
1739                    .scope
1740                    .as_ref()
1741                    .is_some_and(|scope| scope_refusal_locked(&inner, scope).is_some()) =>
1742            {
1743                let refusal = pending
1744                    .scope
1745                    .as_ref()
1746                    .and_then(|scope| scope_refusal_locked(&inner, scope))
1747                    .expect("guard matched a refusal under the same lock");
1748                release_reserved_route_locked(
1749                    &mut inner,
1750                    pending.reservation.client_key,
1751                    pending.reservation.module_key,
1752                );
1753                let abandoned = abandoned_route_target(&inner, &pending.reservation);
1754                let _ = pending
1755                    .sender
1756                    .send(RouteBindRelayOutcome::Rejected(refusal));
1757                return Ok(PendingRelayCompletion {
1758                    settled: true,
1759                    abandoned,
1760                });
1761            }
1762            RouteBindRelayOutcome::Accepted => {
1763                let abandoned = commit_route_locked(&mut inner, pending)?;
1764                return Ok(PendingRelayCompletion {
1765                    settled: true,
1766                    abandoned,
1767                });
1768            }
1769            terminal => {
1770                release_reserved_route_locked(
1771                    &mut inner,
1772                    pending.reservation.client_key,
1773                    pending.reservation.module_key,
1774                );
1775                let _ = pending.sender.send(terminal);
1776            }
1777        }
1778        Ok(PendingRelayCompletion {
1779            settled: true,
1780            abandoned: None,
1781        })
1782    }
1783
1784    pub(crate) fn pending_module_control_op(
1785        &self,
1786        connection_id: ConnectionId,
1787        corr: u64,
1788    ) -> Result<Option<String>, ForwardingError> {
1789        let inner = self.read_inner()?;
1790        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1791            return Ok(None);
1792        };
1793        let key = (endpoint, corr);
1794        Ok(inner
1795            .pending_control_rpcs
1796            .get(&key)
1797            .map(|pending| pending.expected_op.clone())
1798            .or_else(|| {
1799                inner
1800                    .health_probe_tombstones
1801                    .get(&key)
1802                    .filter(|tombstone| tombstone.expires_at > Instant::now())
1803                    .map(|tombstone| tombstone.expected_op.clone())
1804            }))
1805    }
1806
1807    pub(crate) fn complete_module_control_rpc(
1808        &self,
1809        connection_id: ConnectionId,
1810        corr: u64,
1811        actual_op: Option<&str>,
1812        outcome: ModuleControlRpcOutcome,
1813    ) -> Result<ModuleControlRpcCompletion, ForwardingError> {
1814        let now = Instant::now();
1815        let mut inner = self.write_inner()?;
1816        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1817            return Ok(ModuleControlRpcCompletion::Unknown);
1818        };
1819        let key = (endpoint, corr);
1820        if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
1821            if now >= pending.deadline {
1822                let late_health_answer = pending.health_probe_started_at.map(|probe_started_at| {
1823                    ModuleControlRpcCompletion::LateHealthAnswer {
1824                        module_id: inner
1825                            .module_id_by_endpoint
1826                            .get(&endpoint)
1827                            .cloned()
1828                            .unwrap_or_else(|| "unknown".to_string()),
1829                        latency: now.saturating_duration_since(probe_started_at),
1830                    }
1831                });
1832                let _ = pending
1833                    .sender
1834                    .send(ModuleControlRpcOutcome::DeadlineElapsed);
1835                return Ok(late_health_answer.unwrap_or(ModuleControlRpcCompletion::Settled));
1836            }
1837            let outcome = match actual_op {
1838                Some(actual) if actual != pending.expected_op => {
1839                    ModuleControlRpcOutcome::UnexpectedOp {
1840                        expected: pending.expected_op,
1841                        actual: actual.to_string(),
1842                    }
1843                }
1844                _ => outcome,
1845            };
1846            let _ = pending.sender.send(outcome);
1847            return Ok(ModuleControlRpcCompletion::Settled);
1848        }
1849
1850        let Some(tombstone) = inner.health_probe_tombstones.remove(&key) else {
1851            return Ok(ModuleControlRpcCompletion::Unknown);
1852        };
1853        if tombstone.expires_at <= now {
1854            return Ok(ModuleControlRpcCompletion::Unknown);
1855        }
1856        Ok(ModuleControlRpcCompletion::LateHealthAnswer {
1857            module_id: tombstone.module_id,
1858            latency: now.saturating_duration_since(tombstone.probe_started_at),
1859        })
1860    }
1861
1862    #[cfg(test)]
1863    pub(crate) fn health_probe_tombstone_count(&self) -> Result<usize, ForwardingError> {
1864        Ok(self.read_inner()?.health_probe_tombstones.len())
1865    }
1866
1867    #[cfg(test)]
1868    pub(crate) fn closing_connection_count(&self) -> Result<usize, ForwardingError> {
1869        Ok(self.read_inner()?.closing_connections.len())
1870    }
1871
1872    /// Reserved-but-uncommitted route handles, counted on both index sides, so
1873    /// a test can assert a reservation pair was actually given back.
1874    #[cfg(test)]
1875    pub(crate) fn reserved_route_count(&self) -> Result<(usize, usize), ForwardingError> {
1876        let inner = self.read_inner()?;
1877        Ok((inner.reserved_client.len(), inner.reserved_module.len()))
1878    }
1879
1880    pub(crate) fn module_endpoint_for_connection(
1881        &self,
1882        connection_id: ConnectionId,
1883    ) -> Result<Option<ModuleEndpointId>, ForwardingError> {
1884        Ok(self
1885            .read_inner()?
1886            .endpoint_by_connection
1887            .get(&connection_id)
1888            .copied())
1889    }
1890
1891    /// Looks up the module registered on a data-plane connection so route-drop
1892    /// diagnostics name the emitter instead of only reporting a daemon total.
1893    pub(crate) fn module_id_for_connection(
1894        &self,
1895        connection_id: ConnectionId,
1896    ) -> Result<Option<String>, ForwardingError> {
1897        let inner = self.read_inner()?;
1898        Ok(inner
1899            .endpoint_by_connection
1900            .get(&connection_id)
1901            .and_then(|endpoint| inner.module_id_by_endpoint.get(endpoint))
1902            .cloned())
1903    }
1904
1905    /// Whether the daemon ever allocated `(channel, epoch)` on this module
1906    /// connection. Epochs on a module channel are handed out as 1, 2, 3, ...
1907    /// and the last one handed out is remembered until the connection ends, so
1908    /// every epoch from 1 up to that one was a real route at some point. Used
1909    /// only on the drop path, to tell a module still sending on a route the
1910    /// daemon released from one sending on a route that never existed.
1911    pub(crate) fn module_route_epoch_was_allocated(
1912        &self,
1913        connection_id: ConnectionId,
1914        channel: u16,
1915        epoch: u32,
1916    ) -> Result<bool, ForwardingError> {
1917        let inner = self.read_inner()?;
1918        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1919            return Ok(false);
1920        };
1921        Ok(inner
1922            .module_slot_epochs
1923            .get(&ModuleRouteKey { endpoint, channel })
1924            .is_some_and(|last| epoch != 0 && epoch <= *last))
1925    }
1926
1927    pub(crate) fn has_live_module_connection(
1928        &self,
1929        module_id: &str,
1930    ) -> Result<bool, ForwardingError> {
1931        Ok(self.read_inner()?.modules_by_id.contains_key(module_id))
1932    }
1933
1934    pub(crate) fn lookup_data_route(
1935        &self,
1936        connection_id: ConnectionId,
1937        channel: u16,
1938        epoch: u32,
1939    ) -> Result<DataRoute, ForwardingError> {
1940        let inner = self.read_inner()?;
1941        let state = if let Some(endpoint) =
1942            inner.endpoint_by_connection.get(&connection_id).copied()
1943        {
1944            let key = ModuleRouteKey { endpoint, channel };
1945            match inner.module_to_client.get(&key) {
1946                Some(route) if route.module_epoch == epoch => {
1947                    DataRouteState::Bound(Arc::clone(route))
1948                }
1949                Some(_) => DataRouteState::EpochMismatch,
1950                None if inner.reserved_module.contains_key(&key)
1951                    && inner.module_slot_epochs.get(&key).copied() == Some(epoch) =>
1952                {
1953                    DataRouteState::Reserved
1954                }
1955                None if inner.reserved_module.contains_key(&key) => DataRouteState::EpochMismatch,
1956                None => DataRouteState::Absent,
1957            }
1958        } else {
1959            let key = ClientRouteKey {
1960                connection_id,
1961                channel,
1962            };
1963            match inner.client_to_module.get(&key) {
1964                Some(route) if route.client_epoch == epoch => {
1965                    DataRouteState::Bound(Arc::clone(route))
1966                }
1967                Some(_) => DataRouteState::EpochMismatch,
1968                None if inner.reserved_client.contains_key(&key)
1969                    && inner.client_slot_epochs.get(&key).copied() == Some(epoch) =>
1970                {
1971                    DataRouteState::Reserved
1972                }
1973                None if inner.reserved_client.contains_key(&key) => DataRouteState::EpochMismatch,
1974                None => DataRouteState::Absent,
1975            }
1976        };
1977        Ok(
1978            if inner.endpoint_by_connection.contains_key(&connection_id) {
1979                DataRoute::Module(state)
1980            } else {
1981                DataRoute::Client(state)
1982            },
1983        )
1984    }
1985
1986    #[cfg(test)]
1987    pub(crate) fn inject_client_slot_epoch(
1988        &self,
1989        connection_id: ConnectionId,
1990        channel: u16,
1991        last_epoch: u32,
1992    ) {
1993        let mut inner = self.write_inner().expect("forwarding lock");
1994        inner.client_slot_epochs.insert(
1995            ClientRouteKey {
1996                connection_id,
1997                channel,
1998            },
1999            last_epoch,
2000        );
2001        inner.next_client_channel.insert(connection_id, channel);
2002    }
2003
2004    #[cfg(test)]
2005    pub(crate) fn inject_module_slot_epoch(
2006        &self,
2007        endpoint: ModuleEndpointId,
2008        channel: u16,
2009        last_epoch: u32,
2010    ) {
2011        let mut inner = self.write_inner().expect("forwarding lock");
2012        inner
2013            .module_slot_epochs
2014            .insert(ModuleRouteKey { endpoint, channel }, last_epoch);
2015        inner.next_module_channel.insert(endpoint, channel);
2016    }
2017
2018    #[cfg(test)]
2019    pub(crate) fn inject_control_corr(&self, endpoint: ModuleEndpointId, next_corr: u64) {
2020        self.write_inner()
2021            .expect("forwarding lock")
2022            .next_control_corr
2023            .insert(endpoint, next_corr);
2024    }
2025
2026    pub(crate) fn cache_status(
2027        &self,
2028        endpoint: ModuleEndpointId,
2029        module_channel: u16,
2030        module_epoch: u32,
2031        status: String,
2032    ) -> Result<bool, ForwardingError> {
2033        let mut inner = self.write_inner()?;
2034        if !inner.module_id_by_endpoint.contains_key(&endpoint) {
2035            return Err(ForwardingError::StaleModuleEndpoint);
2036        }
2037
2038        let module_key = ModuleRouteKey {
2039            endpoint,
2040            channel: module_channel,
2041        };
2042        let handle = if let Some(route) = inner.module_to_client.get(&module_key) {
2043            (route.module_epoch == module_epoch).then_some((
2044                ClientRouteKey {
2045                    connection_id: route.client_connection_id,
2046                    channel: route.client_channel,
2047                },
2048                route.client_epoch,
2049            ))
2050        } else if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2051            (inner.module_slot_epochs.get(&module_key).copied() == Some(module_epoch)).then_some((
2052                client_key,
2053                inner
2054                    .client_slot_epochs
2055                    .get(&client_key)
2056                    .copied()
2057                    .unwrap_or(0),
2058            ))
2059        } else {
2060            None
2061        };
2062
2063        if let Some(handle) = handle {
2064            inner.status.insert(handle, status);
2065            Ok(true)
2066        } else {
2067            debug!(
2068                module_channel,
2069                module_epoch,
2070                generation = endpoint.generation,
2071                connection_id = endpoint.connection_id.get(),
2072                "dropping stale status update for module route handle"
2073            );
2074            Ok(false)
2075        }
2076    }
2077
2078    pub(crate) fn route_poll_snapshot(
2079        &self,
2080        client_connection_id: ConnectionId,
2081        client_channel: u16,
2082        client_epoch: u32,
2083    ) -> Result<RoutePollSnapshot, ForwardingError> {
2084        let inner = self.read_inner()?;
2085        let client_key = ClientRouteKey {
2086            connection_id: client_connection_id,
2087            channel: client_channel,
2088        };
2089        let Some(route) = inner.client_to_module.get(&client_key) else {
2090            return Ok(RoutePollSnapshot::Absent);
2091        };
2092        if route.client_epoch != client_epoch
2093            || !inner
2094                .module_id_by_endpoint
2095                .contains_key(&route.module_endpoint)
2096        {
2097            return Ok(RoutePollSnapshot::Absent);
2098        }
2099        Ok(RoutePollSnapshot::Bound {
2100            module_id: route.module_id.clone(),
2101            status: inner.status.get(&(client_key, client_epoch)).cloned(),
2102        })
2103    }
2104
2105    pub fn active_binding_count(&self) -> Result<usize, ForwardingError> {
2106        Ok(self.read_inner()?.client_to_module.len())
2107    }
2108
2109    /// How many distinct client connections hold at least one committed route,
2110    /// alongside the largest number of routes any single connection holds.
2111    ///
2112    /// `connected_clients` alone cannot distinguish many clients with a route
2113    /// each from one client accumulating hundreds, and those have opposite
2114    /// causes. Reading it required an out-of-band `lsof` during a live
2115    /// investigation, and the count of connections was mistaken for a count of
2116    /// client processes — which sent two of us after cleanup paths that were
2117    /// working correctly.
2118    pub fn client_route_concentration(&self) -> Result<(usize, usize), ForwardingError> {
2119        let inner = self.read_inner()?;
2120        let mut per_connection: HashMap<ConnectionId, usize> = HashMap::new();
2121        for key in inner.client_to_module.keys() {
2122            *per_connection.entry(key.connection_id).or_insert(0) += 1;
2123        }
2124        let max = per_connection.values().copied().max().unwrap_or(0);
2125        Ok((per_connection.len(), max))
2126    }
2127
2128    pub fn has_route_channel(&self, route_channel: u16) -> Result<bool, ForwardingError> {
2129        let inner = self.read_inner()?;
2130        Ok(inner
2131            .client_to_module
2132            .keys()
2133            .any(|key| key.channel == route_channel))
2134    }
2135
2136    /// Whether daemon-wide shutdown has begun draining all providers.
2137    /// `None` means normal operation; during daemon shutdown, return the
2138    /// endpoint's measured result, or false if no completed measurement exists.
2139    pub(crate) fn daemon_drain_result_for_connection(
2140        &self,
2141        connection_id: ConnectionId,
2142    ) -> Result<Option<bool>, ForwardingError> {
2143        let inner = self.read_inner()?;
2144        Ok(inner.daemon_draining.then(|| {
2145            inner
2146                .endpoint_by_connection
2147                .get(&connection_id)
2148                .and_then(|endpoint| inner.daemon_drain_results.get(endpoint))
2149                .copied()
2150                .unwrap_or(false)
2151        }))
2152    }
2153
2154    #[cfg(unix)]
2155    pub(crate) fn record_daemon_drain_result(
2156        &self,
2157        endpoint: ModuleEndpointId,
2158        drained: bool,
2159    ) -> Result<(), ForwardingError> {
2160        let mut inner = self.write_inner()?;
2161        // Only update an endpoint captured at the start of daemon shutdown.
2162        // Teardown removes its entry, so a late wait cannot recreate the result
2163        // or attribute it to a replacement connection.
2164        if let Some(result) = inner.daemon_drain_results.get_mut(&endpoint) {
2165            *result = drained;
2166        }
2167        Ok(())
2168    }
2169
2170    /// Gate every provider atomically, including registrations racing shutdown.
2171    /// No supervisor lock is held while taking the forwarding lock.
2172    #[cfg(unix)]
2173    pub(crate) fn begin_daemon_drain(
2174        &self,
2175    ) -> Result<Vec<(String, ModuleEndpointId)>, ForwardingError> {
2176        let mut inner = self.write_inner()?;
2177        if inner.daemon_draining {
2178            return Ok(Vec::new());
2179        }
2180        inner.daemon_draining = true;
2181        let modules = inner
2182            .modules_by_id
2183            .iter()
2184            .map(|(id, module)| (id.clone(), module.endpoint))
2185            .collect::<Vec<_>>();
2186        for (_, endpoint) in &modules {
2187            inner.daemon_drain_results.insert(*endpoint, false);
2188            inner
2189                .draining_endpoints
2190                .insert(*endpoint, RouteCloseReason::Restart);
2191        }
2192        // Swap candidates and superseded incumbents are not routable, but they
2193        // are live endpoints that could still be sent module-control work, so
2194        // the daemon-wide gate covers them too. Only the active endpoints are
2195        // returned to the supervisor for measurement, one per module id.
2196        let off_slot_endpoints = inner
2197            .candidates_by_id
2198            .values()
2199            .map(|module| module.endpoint)
2200            .chain(inner.superseded_endpoints.keys().copied())
2201            .collect::<Vec<_>>();
2202        for endpoint in off_slot_endpoints {
2203            inner
2204                .draining_endpoints
2205                .insert(endpoint, RouteCloseReason::Restart);
2206        }
2207        Ok(modules)
2208    }
2209
2210    /// Begin draining whatever endpoint is ACTIVE for `module_id`.
2211    ///
2212    /// After a swap's cutover the active endpoint is the promoted candidate, so
2213    /// this must not be used to drain the incumbent; use
2214    /// [`Self::begin_endpoint_drain`] with the endpoint `cutover_candidate`
2215    /// returned.
2216    pub(crate) fn begin_module_drain(
2217        &self,
2218        module_id: &str,
2219        reason: RouteCloseReason,
2220    ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
2221        let mut inner = self.write_inner()?;
2222        let Some(module) = inner.modules_by_id.get(module_id).cloned() else {
2223            return Ok(None);
2224        };
2225        Ok(Some(begin_drain_locked(
2226            &mut inner, module_id, module, reason,
2227        )))
2228    }
2229
2230    /// Begin draining one specific endpoint, whichever slot it is in.
2231    ///
2232    /// This is how a swap drains its incumbent after cutover: the incumbent's
2233    /// endpoint is captured by `cutover_candidate`, and resolving it by module id
2234    /// instead would find the promoted candidate and leave neither process
2235    /// routable. Returns `Ok(None)` when the endpoint is no longer registered.
2236    pub(crate) fn begin_endpoint_drain(
2237        &self,
2238        endpoint: ModuleEndpointId,
2239        reason: RouteCloseReason,
2240    ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
2241        let mut inner = self.write_inner()?;
2242        let Some(module) = module_connection_for_endpoint_locked(&inner, endpoint).cloned() else {
2243            return Ok(None);
2244        };
2245        let module_id = inner
2246            .module_id_by_endpoint
2247            .get(&endpoint)
2248            .cloned()
2249            .expect("an endpoint resolved to a module connection has a module id");
2250        Ok(Some(begin_drain_locked(
2251            &mut inner, &module_id, module, reason,
2252        )))
2253    }
2254}
2255
2256/// Mark `module`'s endpoint draining and settle everything still pending on it.
2257/// Shared by the by-id and by-endpoint drain entry points, which differ only in
2258/// how they find the endpoint.
2259fn begin_drain_locked(
2260    inner: &mut ForwardingInner,
2261    module_id: &str,
2262    module: ModuleConnection,
2263    reason: RouteCloseReason,
2264) -> ModuleDrainTarget {
2265    {
2266        let endpoint = module.endpoint;
2267        inner.draining_endpoints.insert(endpoint, reason);
2268
2269        let flows = inner
2270            .client_to_module
2271            .values()
2272            .filter(|route| route.module_endpoint == endpoint)
2273            .map(|route| Arc::clone(&route.flow))
2274            .collect::<Vec<_>>();
2275        let excluded_subscriptions = flows
2276            .into_iter()
2277            .map(|flow| flow.begin_drain())
2278            .fold(0u32, u32::saturating_add);
2279
2280        let pending_keys = inner
2281            .pending_relays
2282            .keys()
2283            .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2284            .copied()
2285            .collect::<Vec<_>>();
2286        let mut abandoned_bindings = Vec::new();
2287        for key in pending_keys {
2288            let Some(pending) = inner.pending_relays.remove(&key) else {
2289                continue;
2290            };
2291            release_reserved_route_locked(
2292                inner,
2293                pending.reservation.client_key,
2294                pending.reservation.module_key,
2295            );
2296            if pending.relay_enqueued {
2297                if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2298                    abandoned_bindings.push(target);
2299                }
2300            } else {
2301                // The bind may still be on its way into the module's queue;
2302                // a GOODBYE sent now could overtake it. `route.open` sends
2303                // it once its send completes.
2304                inner.relays_released_before_enqueue.insert(key);
2305            }
2306            let _ = pending
2307                .sender
2308                .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
2309                    "module_reloading",
2310                    format!("module_id '{module_id}' is reloading"),
2311                )));
2312        }
2313
2314        let pending_control_keys = inner
2315            .pending_control_rpcs
2316            .keys()
2317            .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2318            .copied()
2319            .collect::<Vec<_>>();
2320        for key in pending_control_keys {
2321            if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
2322                let _ = pending
2323                    .sender
2324                    .send(ModuleControlRpcOutcome::ModuleGone(format!(
2325                        "module '{module_id}' began draining during module-control RPC"
2326                    )));
2327            }
2328        }
2329
2330        ModuleDrainTarget {
2331            endpoint,
2332            sink: module.sink,
2333            negotiated_ver: module.negotiated_ver,
2334            abandoned_bindings,
2335            excluded_subscriptions,
2336        }
2337    }
2338}
2339
2340/// What a drain that timed out was still waiting on, for the log line that
2341/// reports the timeout. Per-request ages are not tracked.
2342#[derive(Debug, Default, PartialEq, Eq)]
2343pub(crate) struct DrainHoldouts {
2344    /// Requests still counted against the drain (subscriptions flagged as
2345    /// such are excluded from the drain and not counted here).
2346    pub(crate) requests: usize,
2347    /// Routes holding at least one of those requests.
2348    pub(crate) routes: usize,
2349    /// Every route on the endpoint, for scale.
2350    pub(crate) total_routes: usize,
2351    /// The client connections holding the most requests, largest first, at
2352    /// most three: enough to name the consumer without listing every route.
2353    pub(crate) top_connections: Vec<(u64, usize)>,
2354    /// The held requests themselves, as the module saw them, so its own log
2355    /// can say what each one was: `(module channel, corr)`, ordered by channel
2356    /// then corr, at most [`DRAIN_HELD_REQUESTS_LISTED`]. The daemon never reads
2357    /// request bodies, so it cannot name a request's method; the module can,
2358    /// from the channel and corr. A request is released only when the module
2359    /// sends a terminal frame (Response, Error or StreamEnd) with its corr on
2360    /// its route, so every pair listed here is one the module ended, if at all,
2361    /// without sending that frame.
2362    pub(crate) held: Vec<(u16, u64)>,
2363}
2364
2365/// How many held requests the drain-timeout line lists by channel and corr.
2366pub(crate) const DRAIN_HELD_REQUESTS_LISTED: usize = 32;
2367
2368impl ForwardingTable {
2369    /// Summarise the requests one endpoint's drain is still waiting on.
2370    pub(crate) fn endpoint_drain_holdouts(
2371        &self,
2372        endpoint: ModuleEndpointId,
2373    ) -> Result<DrainHoldouts, ForwardingError> {
2374        let inner = self.read_inner()?;
2375        let mut holdouts = DrainHoldouts::default();
2376        let mut by_connection: HashMap<u64, usize> = HashMap::new();
2377        for (key, route) in &inner.client_to_module {
2378            if route.module_endpoint != endpoint {
2379                continue;
2380            }
2381            holdouts.total_routes += 1;
2382            let held = route.flow.drain_in_flight();
2383            if held == 0 {
2384                continue;
2385            }
2386            holdouts.requests += held;
2387            holdouts.routes += 1;
2388            *by_connection.entry(key.connection_id.get()).or_default() += held;
2389            holdouts.held.extend(
2390                route
2391                    .flow
2392                    .drain_held_corrs()
2393                    .into_iter()
2394                    .map(|corr| (route.module_channel, corr)),
2395            );
2396        }
2397        holdouts.held.sort_unstable();
2398        holdouts.held.truncate(DRAIN_HELD_REQUESTS_LISTED);
2399        let mut connections = by_connection.into_iter().collect::<Vec<_>>();
2400        connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
2401        connections.truncate(3);
2402        holdouts.top_connections = connections;
2403        Ok(holdouts)
2404    }
2405
2406    pub(crate) fn endpoint_in_flight_count(
2407        &self,
2408        endpoint: ModuleEndpointId,
2409    ) -> Result<usize, ForwardingError> {
2410        let inner = self.read_inner()?;
2411        Ok(inner
2412            .client_to_module
2413            .values()
2414            .filter(|route| route.module_endpoint == endpoint)
2415            .map(|route| route.flow.drain_in_flight())
2416            .sum())
2417    }
2418
2419    pub(crate) fn endpoint_is_draining(
2420        &self,
2421        endpoint: ModuleEndpointId,
2422    ) -> Result<bool, ForwardingError> {
2423        Ok(self
2424            .read_inner()?
2425            .draining_endpoints
2426            .contains_key(&endpoint))
2427    }
2428
2429    pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
2430        let inner = self.read_inner()?;
2431        Ok(inner
2432            .modules_by_id
2433            .get(module_id)
2434            .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
2435    }
2436
2437    pub(crate) fn release_module_endpoint_routes(
2438        &self,
2439        endpoint: ModuleEndpointId,
2440    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2441        let mut inner = self.write_inner()?;
2442        let routes = inner
2443            .module_to_client
2444            .iter()
2445            .filter(|(module_key, _)| module_key.endpoint == endpoint)
2446            .map(|(module_key, route)| (*module_key, route.module_epoch))
2447            .collect::<Vec<_>>();
2448        let mut released = Vec::with_capacity(routes.len());
2449        for (module_key, epoch) in routes {
2450            if let RouteRelease::Removed(target) =
2451                release_module_route_locked(&mut inner, module_key, epoch)
2452            {
2453                released.push(target);
2454            }
2455        }
2456        Ok(released)
2457    }
2458
2459    /// Enumerate one endpoint's current routes without contacting the module.
2460    ///
2461    /// The read lock makes this safe while the endpoint drains: the returned
2462    /// `draining` marker describes the same table state that owns the route,
2463    /// rather than inferring liveness from a module that may be stopping.
2464    pub(crate) fn endpoint_routes(
2465        &self,
2466        endpoint: ModuleEndpointId,
2467    ) -> Result<Vec<EndpointRoute>, ForwardingError> {
2468        let inner = self.read_inner()?;
2469        Ok(endpoint_routes_locked(&inner, endpoint))
2470    }
2471
2472    /// Snapshot all live endpoint route sets under one forwarding-table read lock.
2473    pub(crate) fn route_census(
2474        &self,
2475        module_id: Option<&str>,
2476    ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
2477        let inner = self.read_inner()?;
2478        let mut endpoints = inner
2479            .modules_by_id
2480            .iter()
2481            .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
2482            .map(|(id, module)| (id.clone(), module.endpoint))
2483            .collect::<Vec<_>>();
2484        endpoints.sort_by(|left, right| left.0.cmp(&right.0));
2485        Ok(endpoints
2486            .into_iter()
2487            .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
2488            .collect())
2489    }
2490
2491    /// Snapshot only the endpoint currently routable under this module id.
2492    pub(crate) fn live_roots(
2493        &self,
2494        module_id: &str,
2495    ) -> Result<ModuleControlResponseToModule, ForwardingError> {
2496        let inner = self.read_inner()?;
2497        let endpoint = inner
2498            .modules_by_id
2499            .get(module_id)
2500            .map(|module| module.endpoint);
2501        let mut roots = BTreeMap::new();
2502        let mut unknown_root_bindings = 0;
2503        let mut total_bindings = 0;
2504        if let Some(endpoint) = endpoint {
2505            for binding in inner
2506                .module_to_client
2507                .values()
2508                .filter(|binding| binding.module_endpoint == endpoint)
2509            {
2510                total_bindings += 1;
2511                if let Some(root) = &binding.project_root {
2512                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2513                    entry.0 += 1;
2514                } else {
2515                    unknown_root_bindings += 1;
2516                }
2517            }
2518            for pending in inner
2519                .pending_relays
2520                .values()
2521                .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
2522            {
2523                total_bindings += 1;
2524                if let Some(root) = &pending.reservation.project_root {
2525                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2526                    entry.1 += 1;
2527                } else {
2528                    unknown_root_bindings += 1;
2529                }
2530            }
2531        }
2532        Ok(ModuleControlResponseToModule::LiveRoots {
2533            roots: roots
2534                .into_iter()
2535                .map(|(project_root, (bound, pending))| LiveRoot {
2536                    project_root,
2537                    bound,
2538                    pending,
2539                })
2540                .collect(),
2541            unknown_root_bindings,
2542            total_bindings,
2543        })
2544    }
2545
2546    /// True if this connection already owns committed or reserved CLIENT routes.
2547    /// A module registers (HELLO) before serving and never opens client routes, so
2548    /// a connection that has client routes must not also become a module endpoint
2549    /// — otherwise one connection holds both client and module state and cleanup
2550    /// only releases one side.
2551    pub(crate) fn connection_has_client_routes(
2552        &self,
2553        connection_id: ConnectionId,
2554    ) -> Result<bool, ForwardingError> {
2555        let inner = self.read_inner()?;
2556        Ok(connection_has_client_routes_locked(&inner, connection_id))
2557    }
2558
2559    pub(crate) fn cleanup_connection(
2560        &self,
2561        connection_id: ConnectionId,
2562    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2563        self.cleanup_connection_counted(connection_id)
2564            .map(|cleanup| cleanup.released)
2565    }
2566
2567    /// [`Self::cleanup_connection`], also reporting how many pending
2568    /// route.bind relays to the closed module were aborted. That count is the
2569    /// `abandoned` figure of the `route.closed` push sent for a lost module
2570    /// connection; it is always zero for a client connection.
2571    pub(crate) fn cleanup_connection_counted(
2572        &self,
2573        connection_id: ConnectionId,
2574    ) -> Result<ConnectionCleanup, ForwardingError> {
2575        let mut inner = self.write_inner()?;
2576        inner.closing_connections.insert(connection_id);
2577        let cleanup = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
2578            remove_module_connection_locked(&mut inner, endpoint)
2579        } else {
2580            ConnectionCleanup {
2581                released: Self::cleanup_client_connection_locked(&mut inner, connection_id),
2582                abandoned_relays: 0,
2583            }
2584        };
2585        // The closing mark refuses new work for a connection whose teardown is
2586        // still pending. This is the latest point at which lifting it is safe:
2587        // teardown has just removed every per-connection entry above, under
2588        // this same write lock, so no lookup can still find live state for the
2589        // id; and ids come from a monotonic counter that never reissues one,
2590        // so the id can never name a future connection either. Keeping the
2591        // mark past this point only grew the set by one entry per connection
2592        // for the life of the daemon.
2593        inner.closing_connections.remove(&connection_id);
2594        Ok(cleanup)
2595    }
2596
2597    fn cleanup_client_connection_locked(
2598        inner: &mut ForwardingInner,
2599        connection_id: ConnectionId,
2600    ) -> Vec<GoodbyeTarget> {
2601        let routes = inner
2602            .client_to_module
2603            .iter()
2604            .filter(|(key, _)| key.connection_id == connection_id)
2605            .map(|(key, route)| (*key, route.client_epoch))
2606            .collect::<Vec<_>>();
2607        let mut released = Vec::with_capacity(routes.len());
2608        for (client_key, epoch) in routes {
2609            if let RouteRelease::Removed(target) =
2610                release_client_route_locked(inner, client_key, epoch)
2611            {
2612                released.push(target);
2613            }
2614        }
2615
2616        let pending_keys = inner
2617            .pending_relays
2618            .iter()
2619            .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2620            .map(|(key, _)| *key)
2621            .collect::<Vec<_>>();
2622        for key in pending_keys {
2623            let Some(pending) = inner.pending_relays.remove(&key) else {
2624                continue;
2625            };
2626            release_reserved_route_locked(
2627                inner,
2628                pending.reservation.client_key,
2629                pending.reservation.module_key,
2630            );
2631            if pending.relay_enqueued {
2632                if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2633                    released.push(target);
2634                }
2635            } else {
2636                // The bind may still be on its way into the module's queue;
2637                // a GOODBYE sent now could overtake it. `route.open` sends
2638                // it once its send completes.
2639                inner.relays_released_before_enqueue.insert(key);
2640            }
2641            let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2642                "client connection closed during route.bind relay".to_string(),
2643            ));
2644        }
2645
2646        let orphaned = inner
2647            .reserved_client
2648            .iter()
2649            .filter(|(key, _)| key.connection_id == connection_id)
2650            .map(|(client, module)| (*client, *module))
2651            .collect::<Vec<_>>();
2652        for (client_key, module_key) in orphaned {
2653            release_reserved_route_locked(inner, client_key, module_key);
2654        }
2655        inner.next_client_channel.remove(&connection_id);
2656        inner
2657            .client_slot_epochs
2658            .retain(|key, _| key.connection_id != connection_id);
2659        inner
2660            .last_published_epoch
2661            .retain(|key, _| key.connection_id != connection_id);
2662        inner
2663            .status
2664            .retain(|(key, _), _| key.connection_id != connection_id);
2665
2666        released
2667    }
2668
2669    /// Close a client connection whose egress queue refused a frame for the
2670    /// route `(connection_id, channel)` at `expected_epoch`. The whole
2671    /// connection closes because its queue is shared by every route on it.
2672    ///
2673    /// The first request that actually closes the connection logs one WARN with
2674    /// the diagnosis: which principals the connection's routes belonged to,
2675    /// which module's frame did not fit and on which client channel, and what
2676    /// the queue held at that moment.
2677    pub(crate) fn escalate_client_delivery_failure(
2678        &self,
2679        connection_id: ConnectionId,
2680        channel: u16,
2681        expected_epoch: u32,
2682        reason: CloseReason,
2683        undelivered: UndeliveredFrame<'_>,
2684    ) -> Result<bool, ForwardingError> {
2685        let principals = {
2686            let mut inner = self.write_inner()?;
2687            let key = ClientRouteKey {
2688                connection_id,
2689                channel,
2690            };
2691            if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2692                None
2693            } else {
2694                inner.closing_connections.insert(connection_id);
2695                Some(connection_principals_locked(&inner, connection_id))
2696            }
2697        };
2698        let Some(principals) = principals else {
2699            return Ok(false);
2700        };
2701        let backlog = undelivered.sink.backlog();
2702        let close_reason = reason.to_string();
2703        if self.request_connection_close(connection_id, reason) {
2704            warn!(
2705                connection_id = connection_id.get(),
2706                principals = %principals,
2707                module_id = undelivered.module_id.unwrap_or("unknown"),
2708                client_channel = channel,
2709                queued_bytes = backlog.queued_bytes,
2710                queued_frames = backlog.queued_frames,
2711                oldest_queued_ms = backlog
2712                    .oldest_age
2713                    .map(|age| age.as_millis() as u64)
2714                    .unwrap_or(0),
2715                close_reason = %close_reason,
2716                "closing client connection: its egress queue could not take a frame"
2717            );
2718        }
2719        Ok(true)
2720    }
2721
2722    /// Publish the tags a scope sync changed and close the live routes its
2723    /// drain table selects, in one critical section.
2724    ///
2725    /// The caller holds the scope table's write lock (scope table, then this
2726    /// table, always). Publishing and selecting under one lock is what closes
2727    /// the race with a bind commit: a commit before this sees the old tag and
2728    /// its route is selected here like any other; a commit after it sees the
2729    /// new tag and is refused.
2730    ///
2731    /// Every client route index entry is scanned, which includes routes on a
2732    /// swap's superseded endpoint (the old process of a module being replaced
2733    /// blue/green, which keeps its existing routes until they drain): those
2734    /// stay indexed until drained, so ending a scope reaches them too.
2735    pub(crate) fn publish_scope_changes(
2736        &self,
2737        changes: &[ScopeTagChange],
2738    ) -> Result<Vec<ScopeDrainedRoute>, ForwardingError> {
2739        let mut inner = self.write_inner()?;
2740        let mut by_scope: HashMap<(&str, &str), &ScopeTagChange> = HashMap::new();
2741        for change in changes {
2742            let key = (change.owner.clone(), change.scope_ref.clone());
2743            match change.after {
2744                Some(tag) => {
2745                    inner.scope_tags.insert(key, tag);
2746                }
2747                None => {
2748                    inner.scope_tags.remove(&key);
2749                }
2750            }
2751            by_scope.insert((change.owner.as_str(), change.scope_ref.as_str()), change);
2752        }
2753        let mut selected = Vec::new();
2754        for (client_key, route) in &inner.client_to_module {
2755            let Some(scope) = &route.scope else {
2756                continue;
2757            };
2758            let Some(change) = by_scope.get(&(scope.owner.as_str(), scope.scope_ref.as_str()))
2759            else {
2760                continue;
2761            };
2762            if change.before.map(|tag| tag.scope_epoch) != Some(scope.tag.scope_epoch) {
2763                continue;
2764            }
2765            let reason = match &change.drain {
2766                ScopeDrain::Nothing => continue,
2767                ScopeDrain::All(reason) => *reason,
2768                ScopeDrain::Carriers(narrowed) => {
2769                    let owner = Principal::Reserved {
2770                        module_id: scope.owner.clone(),
2771                    };
2772                    let hit = route.principal != owner
2773                        && narrowed.iter().any(|(principal, allowed)| {
2774                            *principal == route.principal
2775                                && allowed
2776                                    .as_ref()
2777                                    .is_none_or(|targets| !targets.contains(&route.module_id))
2778                        });
2779                    if !hit {
2780                        continue;
2781                    }
2782                    RouteCloseReason::ScopeCarrierRemoved
2783                }
2784            };
2785            selected.push((*client_key, route.client_epoch, reason));
2786        }
2787        let mut drained = Vec::new();
2788        for (client_key, client_epoch, reason) in selected {
2789            let Some(route) = inner.client_to_module.get(&client_key).cloned() else {
2790                continue;
2791            };
2792            let release = release_client_route_locked(&mut inner, client_key, client_epoch);
2793            self.record_route_release(&release);
2794            if let RouteRelease::Removed(module) = release {
2795                drained.push(ScopeDrainedRoute {
2796                    reason,
2797                    scope: route
2798                        .scope
2799                        .clone()
2800                        .expect("selected scope drain has a bound scope"),
2801                    module_id: route.module_id.clone(),
2802                    client: GoodbyeTarget {
2803                        connection_id: route.client_connection_id,
2804                        sink: route.client_sink.clone(),
2805                        negotiated_ver: route.client_negotiated_ver,
2806                        channel: route.client_channel,
2807                        epoch: route.client_epoch,
2808                        kind: GoodbyeTargetKind::Client,
2809                        module_id: Some(route.module_id.clone()),
2810                    },
2811                    module,
2812                });
2813            }
2814        }
2815        Ok(drained)
2816    }
2817
2818    /// The published tag of one scope, for tests of the commit re-check.
2819    #[cfg(test)]
2820    pub(crate) fn published_scope_tag(&self, owner: &str, scope_ref: &str) -> Option<ScopeTag> {
2821        self.read_inner()
2822            .ok()?
2823            .scope_tags
2824            .get(&(owner.to_string(), scope_ref.to_string()))
2825            .copied()
2826    }
2827
2828    fn record_route_release(&self, release: &RouteRelease) {
2829        match release {
2830            RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2831            RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2832            RouteRelease::Absent => {}
2833        }
2834    }
2835
2836    fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2837        self.inner.read().map_err(|_| ForwardingError::Poisoned)
2838    }
2839
2840    fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2841        self.inner.write().map_err(|_| ForwardingError::Poisoned)
2842    }
2843
2844    fn lock_close_registry(
2845        &self,
2846    ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2847        self.close_registry
2848            .lock()
2849            .unwrap_or_else(|poisoned| poisoned.into_inner())
2850    }
2851}
2852
2853impl ForwardingInner {
2854    fn allocate_route_slots(
2855        &mut self,
2856        connection_id: ConnectionId,
2857        endpoint: ModuleEndpointId,
2858    ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2859        let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2860        let mut client_channel = client_start;
2861        let client_channel = loop {
2862            let key = ClientRouteKey {
2863                connection_id,
2864                channel: client_channel,
2865            };
2866            let eligible = !self.client_to_module.contains_key(&key)
2867                && !self.reserved_client.contains_key(&key)
2868                && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2869            if eligible {
2870                break client_channel;
2871            }
2872            client_channel = next_channel(client_channel);
2873            if client_channel == client_start {
2874                return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2875            }
2876        };
2877
2878        let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2879        let mut module_channel = module_start;
2880        let module_channel = loop {
2881            let key = ModuleRouteKey {
2882                endpoint,
2883                channel: module_channel,
2884            };
2885            let eligible = !self.module_to_client.contains_key(&key)
2886                && !self.reserved_module.contains_key(&key)
2887                && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2888            if eligible {
2889                break module_channel;
2890            }
2891            module_channel = next_channel(module_channel);
2892            if module_channel == module_start {
2893                return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2894            }
2895        };
2896
2897        let client_key = ClientRouteKey {
2898            connection_id,
2899            channel: client_channel,
2900        };
2901        let module_key = ModuleRouteKey {
2902            endpoint,
2903            channel: module_channel,
2904        };
2905        let client_epoch = self
2906            .client_slot_epochs
2907            .get(&client_key)
2908            .copied()
2909            .unwrap_or(0)
2910            + 1;
2911        let module_epoch = self
2912            .module_slot_epochs
2913            .get(&module_key)
2914            .copied()
2915            .unwrap_or(0)
2916            + 1;
2917        self.client_slot_epochs.insert(client_key, client_epoch);
2918        self.module_slot_epochs.insert(module_key, module_epoch);
2919        self.next_client_channel
2920            .insert(connection_id, next_channel(client_channel));
2921        self.next_module_channel
2922            .insert(endpoint, next_channel(module_channel));
2923        Ok((client_channel, client_epoch, module_channel, module_epoch))
2924    }
2925
2926    fn allocate_control_corr(
2927        &mut self,
2928        endpoint: ModuleEndpointId,
2929    ) -> Result<u64, ForwardingError> {
2930        let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2931        if candidate == 0 {
2932            self.closing_connections.insert(endpoint.connection_id);
2933            return Err(ForwardingError::RelayCorrelationExhausted);
2934        }
2935        self.next_control_corr.insert(
2936            endpoint,
2937            if candidate == u64::MAX {
2938                0
2939            } else {
2940                candidate + 1
2941            },
2942        );
2943        Ok(candidate)
2944    }
2945}
2946
2947fn connection_has_client_routes_locked(
2948    inner: &ForwardingInner,
2949    connection_id: ConnectionId,
2950) -> bool {
2951    inner
2952        .client_to_module
2953        .keys()
2954        .any(|key| key.connection_id == connection_id)
2955        || inner
2956            .reserved_client
2957            .keys()
2958            .any(|key| key.connection_id == connection_id)
2959}
2960
2961fn check_module_connection_role_locked(
2962    inner: &ForwardingInner,
2963    connection_id: ConnectionId,
2964) -> Result<(), ForwardingError> {
2965    if inner.endpoint_by_connection.contains_key(&connection_id)
2966        || connection_has_client_routes_locked(inner, connection_id)
2967    {
2968        return Err(ForwardingError::ConnectionRoleConflict { connection_id });
2969    }
2970    Ok(())
2971}
2972
2973fn next_channel(channel: u16) -> u16 {
2974    let next = channel.wrapping_add(1);
2975    if next == 0 {
2976        1
2977    } else {
2978        next
2979    }
2980}
2981
2982fn endpoint_routes_locked(
2983    inner: &ForwardingInner,
2984    endpoint: ModuleEndpointId,
2985) -> Vec<EndpointRoute> {
2986    let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2987    let draining = drain_reason.is_some();
2988    let mut routes = inner
2989        .module_to_client
2990        .iter()
2991        .filter(|(module_key, _)| module_key.endpoint == endpoint)
2992        .map(|(_, route)| EndpointRoute {
2993            goodbye_target: GoodbyeTarget {
2994                connection_id: route.client_connection_id,
2995                sink: route.client_sink.clone(),
2996                negotiated_ver: route.client_negotiated_ver,
2997                channel: route.client_channel,
2998                epoch: route.client_epoch,
2999                kind: GoodbyeTargetKind::Client,
3000                module_id: Some(route.module_id.clone()),
3001            },
3002            principal: route.principal.clone(),
3003            bound_at: route.bound_at,
3004            draining,
3005            drain_reason,
3006        })
3007        .collect::<Vec<_>>();
3008    routes.sort_by_key(|route| {
3009        (
3010            route.goodbye_target.connection_id.get(),
3011            route.goodbye_target.channel,
3012            route.goodbye_target.epoch,
3013        )
3014    });
3015    routes
3016}
3017
3018/// Why a bind admitted under `scope` may no longer commit, read from the
3019/// published tags: `scope_ended` when the scope is gone or at another epoch,
3020/// `scope_changed` (retryable) when only its content moved.
3021fn scope_refusal_locked(inner: &ForwardingInner, scope: &BoundScope) -> Option<ErrorBody> {
3022    let key = (scope.owner.clone(), scope.scope_ref.clone());
3023    match inner.scope_tags.get(&key) {
3024        Some(current) if *current == scope.tag => None,
3025        Some(current) if current.scope_epoch == scope.tag.scope_epoch => Some(ErrorBody::new(
3026            subc_protocol::error_codes::SCOPE_CHANGED,
3027            format!(
3028                "scope '{}' of {} changed while the route was being bound; re-open it",
3029                scope.scope_ref, scope.owner
3030            ),
3031        )),
3032        _ => Some(ErrorBody::new(
3033            subc_protocol::error_codes::SCOPE_ENDED,
3034            format!(
3035                "scope '{}' of {} at scope_epoch {} ended while the route was being bound",
3036                scope.scope_ref, scope.owner, scope.tag.scope_epoch
3037            ),
3038        )),
3039    }
3040}
3041
3042fn release_reserved_route_locked(
3043    inner: &mut ForwardingInner,
3044    client_key: ClientRouteKey,
3045    module_key: ModuleRouteKey,
3046) {
3047    if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
3048        inner.reserved_client.remove(&client_key);
3049    }
3050    if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
3051        inner.reserved_module.remove(&module_key);
3052    }
3053    inner.status.retain(|(key, _), _| *key != client_key);
3054}
3055
3056fn release_client_route_locked(
3057    inner: &mut ForwardingInner,
3058    client_key: ClientRouteKey,
3059    expected_epoch: u32,
3060) -> RouteRelease {
3061    let Some(route) = inner.client_to_module.get(&client_key) else {
3062        return RouteRelease::Absent;
3063    };
3064    if route.client_epoch != expected_epoch {
3065        return RouteRelease::Stale;
3066    }
3067    let route = inner
3068        .client_to_module
3069        .remove(&client_key)
3070        .expect("route checked under the same forwarding lock");
3071    route.flow.close();
3072    if let Some(scope) = &route.scope {
3073        scope.release_bind();
3074    }
3075    inner.module_to_client.remove(&ModuleRouteKey {
3076        endpoint: route.module_endpoint,
3077        channel: route.module_channel,
3078    });
3079    inner.operator_confirms.route_closed(ModuleRouteKey {
3080        endpoint: route.module_endpoint,
3081        channel: route.module_channel,
3082    });
3083    inner.status.remove(&(client_key, expected_epoch));
3084    RouteRelease::Removed(GoodbyeTarget {
3085        connection_id: route.module_endpoint.connection_id,
3086        sink: route.module_sink.clone(),
3087        negotiated_ver: route.module_negotiated_ver,
3088        channel: route.module_channel,
3089        epoch: route.module_epoch,
3090        kind: GoodbyeTargetKind::Module,
3091        module_id: Some(route.module_id.clone()),
3092    })
3093}
3094
3095fn release_module_route_locked(
3096    inner: &mut ForwardingInner,
3097    module_key: ModuleRouteKey,
3098    expected_epoch: u32,
3099) -> RouteRelease {
3100    let Some(route) = inner.module_to_client.get(&module_key) else {
3101        return RouteRelease::Absent;
3102    };
3103    if route.module_epoch != expected_epoch {
3104        return RouteRelease::Stale;
3105    }
3106    let route = inner
3107        .module_to_client
3108        .remove(&module_key)
3109        .expect("route checked under the same forwarding lock");
3110    inner.operator_confirms.route_closed(module_key);
3111    route.flow.close();
3112    if let Some(scope) = &route.scope {
3113        scope.release_bind();
3114    }
3115    let client_key = ClientRouteKey {
3116        connection_id: route.client_connection_id,
3117        channel: route.client_channel,
3118    };
3119    inner.client_to_module.remove(&client_key);
3120    inner.status.remove(&(client_key, route.client_epoch));
3121    RouteRelease::Removed(GoodbyeTarget {
3122        connection_id: route.client_connection_id,
3123        sink: route.client_sink.clone(),
3124        negotiated_ver: route.client_negotiated_ver,
3125        channel: route.client_channel,
3126        epoch: route.client_epoch,
3127        kind: GoodbyeTargetKind::Client,
3128        module_id: Some(route.module_id.clone()),
3129    })
3130}
3131
3132fn commit_route_locked(
3133    inner: &mut ForwardingInner,
3134    pending: PendingRouteBindRelayEntry,
3135) -> Result<Option<GoodbyeTarget>, ForwardingError> {
3136    let reservation = pending.reservation;
3137    if inner
3138        .closing_connections
3139        .contains(&reservation.client_key.connection_id)
3140    {
3141        return Err(ForwardingError::ConnectionClosing {
3142            connection_id: reservation.client_key.connection_id,
3143        });
3144    }
3145    let module_id = inner
3146        .module_id_by_endpoint
3147        .get(&reservation.module_key.endpoint)
3148        .cloned()
3149        .ok_or(ForwardingError::StaleModuleEndpoint)?;
3150    if inner
3151        .draining_endpoints
3152        .contains_key(&reservation.module_key.endpoint)
3153    {
3154        return Err(ForwardingError::ModuleReloading { module_id });
3155    }
3156    if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
3157        || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
3158    {
3159        return Err(ForwardingError::UnknownReservation {
3160            client_channel: reservation.client_key.channel,
3161            module_channel: reservation.module_key.channel,
3162        });
3163    }
3164    let module = inner
3165        .modules_by_id
3166        .get(&module_id)
3167        .filter(|module| module.endpoint == reservation.module_key.endpoint)
3168        .cloned()
3169        .ok_or(ForwardingError::StaleModuleEndpoint)?;
3170    let binding = Arc::new(RouteBinding {
3171        client_connection_id: reservation.client_key.connection_id,
3172        client_sink: pending.client_sink,
3173        client_negotiated_ver: pending.client_negotiated_ver,
3174        client_channel: reservation.client_key.channel,
3175        client_epoch: reservation.client_epoch,
3176        module_id,
3177        module_endpoint: reservation.module_key.endpoint,
3178        module_sink: module.sink,
3179        module_negotiated_ver: module.negotiated_ver,
3180        module_channel: reservation.module_key.channel,
3181        module_epoch: reservation.module_epoch,
3182        principal: pending.principal,
3183        project_root: reservation.project_root.clone(),
3184        bound_at: Instant::now(),
3185        flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
3186        scope: pending.scope,
3187    });
3188    inner
3189        .client_to_module
3190        .insert(reservation.client_key, Arc::clone(&binding));
3191    inner
3192        .module_to_client
3193        .insert(reservation.module_key, Arc::clone(&binding));
3194    let previous_published = inner
3195        .last_published_epoch
3196        .insert(reservation.client_key, reservation.client_epoch);
3197
3198    // Sending through the reserved slot cannot fail, but it reports a receiver
3199    // that closed after reservation and before this locked publication point.
3200    // The frame is stamped and charged to the queued-byte count here, when it
3201    // actually enters the queue, not when the slot was reserved.
3202    let client_writer_closed = pending.client_permit.send(pending.route_open_frame);
3203    if client_writer_closed {
3204        // Not gated on `relay_enqueued`: only a module ack reaches this
3205        // commit, which proves the module holds the binding, and `route.open`
3206        // sends nothing for a relay a reply already settled.
3207        let abandoned = abandoned_route_target(inner, &reservation);
3208        if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
3209            route.flow.close();
3210        }
3211        inner.module_to_client.remove(&reservation.module_key);
3212        inner
3213            .status
3214            .remove(&(reservation.client_key, reservation.client_epoch));
3215        match previous_published {
3216            Some(epoch) => {
3217                inner
3218                    .last_published_epoch
3219                    .insert(reservation.client_key, epoch);
3220            }
3221            None => {
3222                inner.last_published_epoch.remove(&reservation.client_key);
3223            }
3224        }
3225        let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
3226            "client egress closed during route publication".to_string(),
3227        ));
3228        return Ok(abandoned);
3229    }
3230
3231    if let Some(scope) = &binding.scope {
3232        scope.commit_presence();
3233    }
3234    let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
3235    Ok(None)
3236}
3237
3238/// The live module connection registered under exactly `endpoint`, whichever
3239/// slot holds it: active, swap candidate, or superseded incumbent.
3240///
3241/// Resolving by endpoint rather than by module id is what keeps an incumbent
3242/// addressable after a swap promoted a candidate over its id. An endpoint that
3243/// is in none of the three (a stale one, replaced without a promotion) resolves
3244/// to nothing, as it always has.
3245fn module_connection_for_endpoint_locked(
3246    inner: &ForwardingInner,
3247    endpoint: ModuleEndpointId,
3248) -> Option<&ModuleConnection> {
3249    let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
3250    inner
3251        .modules_by_id
3252        .get(module_id)
3253        .filter(|module| module.endpoint == endpoint)
3254        .or_else(|| {
3255            inner
3256                .candidates_by_id
3257                .get(module_id)
3258                .filter(|module| module.endpoint == endpoint)
3259        })
3260        .or_else(|| inner.superseded_endpoints.get(&endpoint))
3261}
3262
3263fn abandoned_route_target(
3264    inner: &ForwardingInner,
3265    reservation: &RouteReservation,
3266) -> Option<GoodbyeTarget> {
3267    let module_id = inner
3268        .module_id_by_endpoint
3269        .get(&reservation.module_key.endpoint)?;
3270    let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
3271    (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
3272        connection_id: module.endpoint.connection_id,
3273        sink: module.sink.clone(),
3274        negotiated_ver: module.negotiated_ver,
3275        channel: reservation.module_key.channel,
3276        epoch: reservation.module_epoch,
3277        kind: GoodbyeTargetKind::Module,
3278        module_id: Some(module_id.clone()),
3279    })
3280}
3281
3282/// Queue a registering module's HELLO_ACK on its sink. Called with the
3283/// forwarding write lock held, before the endpoint is inserted, which is what
3284/// puts the ack ahead of any `route.bind` or control RPC routed to the module.
3285/// `try_send` never waits, so holding the lock across it is safe.
3286fn enqueue_hello_ack_locked(
3287    sink: &FrameSink,
3288    connection_id: ConnectionId,
3289    hello_ack: Option<Frame>,
3290) -> Result<(), ForwardingError> {
3291    let Some(hello_ack) = hello_ack else {
3292        return Ok(());
3293    };
3294    sink.try_send(hello_ack)
3295        .map_err(|_| ForwardingError::ModuleEgressUnavailable { connection_id })
3296}
3297
3298fn remove_module_connection_locked(
3299    inner: &mut ForwardingInner,
3300    endpoint: ModuleEndpointId,
3301) -> ConnectionCleanup {
3302    // Commit the disconnect before its routes are released, so their teardown
3303    // cannot replace module_closed with route_closed.
3304    inner.operator_confirms.module_closed(endpoint);
3305    inner.draining_endpoints.remove(&endpoint);
3306    inner.daemon_drain_results.remove(&endpoint);
3307    let module_id = inner.module_id_by_endpoint.remove(&endpoint);
3308    if let Some(module_id) = module_id.as_ref() {
3309        if inner
3310            .modules_by_id
3311            .get(module_id)
3312            .is_some_and(|module| module.endpoint == endpoint)
3313        {
3314            inner.modules_by_id.remove(module_id);
3315        }
3316        if inner
3317            .candidates_by_id
3318            .get(module_id)
3319            .is_some_and(|module| module.endpoint == endpoint)
3320        {
3321            inner.candidates_by_id.remove(module_id);
3322        }
3323    }
3324    inner.superseded_endpoints.remove(&endpoint);
3325    inner.endpoint_by_connection.remove(&endpoint.connection_id);
3326    inner.next_module_channel.remove(&endpoint);
3327    inner.next_control_corr.remove(&endpoint);
3328    inner
3329        .health_probe_tombstones
3330        .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
3331    inner
3332        .module_slot_epochs
3333        .retain(|key, _| key.endpoint != endpoint);
3334    let reserved_module_keys: Vec<ModuleRouteKey> = inner
3335        .reserved_module
3336        .keys()
3337        .filter(|module_key| module_key.endpoint == endpoint)
3338        .copied()
3339        .collect();
3340    for module_key in reserved_module_keys {
3341        if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
3342            release_reserved_route_locked(inner, client_key, module_key);
3343        }
3344    }
3345
3346    let pending_keys: Vec<_> = inner
3347        .pending_relays
3348        .keys()
3349        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
3350        .copied()
3351        .collect();
3352    let pending: Vec<_> = pending_keys
3353        .into_iter()
3354        .filter_map(|key| {
3355            inner
3356                .pending_relays
3357                .remove(&key)
3358                .map(|pending| (key, pending))
3359        })
3360        .collect();
3361    let abandoned_relays = u32::try_from(pending.len()).unwrap_or(u32::MAX);
3362    for (key, pending) in pending {
3363        // Recorded like every other removal without a reply, so a `route.open`
3364        // still sending the bind learns the relay was released rather than
3365        // settled by the module.
3366        if !pending.relay_enqueued {
3367            inner.relays_released_before_enqueue.insert(key);
3368        }
3369        let module_label = module_id.as_deref().unwrap_or("unknown");
3370        let _ = pending
3371            .sender
3372            .send(RouteBindRelayOutcome::ModuleGone(format!(
3373                "module '{module_label}' connection closed during route.bind relay"
3374            )));
3375    }
3376
3377    let pending_control_keys: Vec<_> = inner
3378        .pending_control_rpcs
3379        .keys()
3380        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
3381        .copied()
3382        .collect();
3383    let pending_control: Vec<_> = pending_control_keys
3384        .into_iter()
3385        .filter_map(|key| inner.pending_control_rpcs.remove(&key))
3386        .collect();
3387    for pending in pending_control {
3388        let module_label = module_id.as_deref().unwrap_or("unknown");
3389        let _ = pending
3390            .sender
3391            .send(ModuleControlRpcOutcome::ModuleGone(format!(
3392                "module '{module_label}' connection closed during module-control RPC"
3393            )));
3394    }
3395
3396    let module_routes = inner
3397        .module_to_client
3398        .iter()
3399        .filter(|(module_key, _)| module_key.endpoint == endpoint)
3400        .map(|(module_key, route)| (*module_key, route.module_epoch))
3401        .collect::<Vec<_>>();
3402    let mut released = Vec::with_capacity(module_routes.len());
3403    for (module_key, epoch) in module_routes {
3404        if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
3405        {
3406            released.push(target);
3407        }
3408    }
3409    ConnectionCleanup {
3410        released,
3411        abandoned_relays,
3412    }
3413}
3414
3415#[derive(Debug, Clone, Copy)]
3416struct RequestCredit {
3417    subscription: bool,
3418    excluded_from_drain: bool,
3419}
3420
3421#[derive(Debug, Default)]
3422struct CreditLedger {
3423    by_corr: HashMap<u64, Vec<RequestCredit>>,
3424}
3425
3426impl CreditLedger {
3427    fn acquire(&mut self, corr: u64, subscription: bool) {
3428        self.by_corr.entry(corr).or_default().push(RequestCredit {
3429            subscription,
3430            excluded_from_drain: false,
3431        });
3432    }
3433
3434    fn release(&mut self, corr: u64) -> bool {
3435        let Some(credits) = self.by_corr.get_mut(&corr) else {
3436            return false;
3437        };
3438        let released = credits.pop().is_some();
3439        if credits.is_empty() {
3440            self.by_corr.remove(&corr);
3441        }
3442        released
3443    }
3444
3445    fn capture_subscription_exclusions(&mut self) -> u32 {
3446        let mut excluded = 0u32;
3447        for credit in self.by_corr.values_mut().flatten() {
3448            if credit.subscription && !credit.excluded_from_drain {
3449                credit.excluded_from_drain = true;
3450                excluded = excluded.saturating_add(1);
3451            }
3452        }
3453        excluded
3454    }
3455
3456    #[cfg(test)]
3457    fn in_flight(&self) -> usize {
3458        self.by_corr.values().map(Vec::len).sum()
3459    }
3460
3461    fn drain_in_flight(&self) -> usize {
3462        self.by_corr
3463            .values()
3464            .flatten()
3465            .filter(|credit| !credit.excluded_from_drain)
3466            .count()
3467    }
3468
3469    /// The correlation ids of the requests a drain still waits on, one entry
3470    /// per held credit, in ascending order.
3471    fn drain_held_corrs(&self) -> Vec<u64> {
3472        let mut corrs = self
3473            .by_corr
3474            .iter()
3475            .flat_map(|(corr, credits)| {
3476                credits
3477                    .iter()
3478                    .filter(|credit| !credit.excluded_from_drain)
3479                    .map(move |_| *corr)
3480            })
3481            .collect::<Vec<_>>();
3482        corrs.sort_unstable();
3483        corrs
3484    }
3485}
3486
3487#[derive(Debug, Default)]
3488struct ChannelFlowState {
3489    closed: bool,
3490    credits: CreditLedger,
3491}
3492
3493/// Per-channel request-credit accounting shared by the client and module route halves.
3494#[derive(Debug)]
3495pub(crate) struct ChannelFlow {
3496    sem: Semaphore,
3497    window: usize,
3498    state: Mutex<ChannelFlowState>,
3499}
3500
3501impl ChannelFlow {
3502    pub(crate) fn new(window: usize) -> Self {
3503        debug_assert!(window > 0, "flow-control window must be non-zero");
3504        Self {
3505            sem: Semaphore::new(window),
3506            window,
3507            state: Mutex::new(ChannelFlowState::default()),
3508        }
3509    }
3510
3511    #[cfg(test)]
3512    pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
3513        self.acquire_tagged(0, false).await
3514    }
3515
3516    pub(crate) async fn acquire_tagged(
3517        &self,
3518        corr: u64,
3519        subscription: bool,
3520    ) -> Result<(), ChannelFlowClosed> {
3521        let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
3522        let mut state = self
3523            .state
3524            .lock()
3525            .unwrap_or_else(|poisoned| poisoned.into_inner());
3526        if state.closed {
3527            return Err(ChannelFlowClosed);
3528        }
3529        state.credits.acquire(corr, subscription);
3530        permit.forget();
3531        Ok(())
3532    }
3533
3534    #[cfg(test)]
3535    pub(crate) fn release(&self) {
3536        self.release_corr(0);
3537    }
3538
3539    pub(crate) fn release_corr(&self, corr: u64) {
3540        let released = self
3541            .state
3542            .lock()
3543            .unwrap_or_else(|poisoned| poisoned.into_inner())
3544            .credits
3545            .release(corr);
3546        if !released {
3547            // Protocol-conforming modules emit exactly one terminal per request.
3548            // This guard is a best-effort safety net against window growth, not a
3549            // security boundary against malicious peers.
3550            warn!(
3551                window = self.window,
3552                available = self.sem.available_permits(),
3553                "flow-control over-release ignored"
3554            );
3555            return;
3556        }
3557        if !self.sem.is_closed() {
3558            self.sem.add_permits(1);
3559        }
3560    }
3561
3562    #[cfg(test)]
3563    pub(crate) fn in_flight(&self) -> usize {
3564        self.state
3565            .lock()
3566            .unwrap_or_else(|poisoned| poisoned.into_inner())
3567            .credits
3568            .in_flight()
3569    }
3570
3571    pub(crate) fn drain_in_flight(&self) -> usize {
3572        self.state
3573            .lock()
3574            .unwrap_or_else(|poisoned| poisoned.into_inner())
3575            .credits
3576            .drain_in_flight()
3577    }
3578
3579    pub(crate) fn drain_held_corrs(&self) -> Vec<u64> {
3580        self.state
3581            .lock()
3582            .unwrap_or_else(|poisoned| poisoned.into_inner())
3583            .credits
3584            .drain_held_corrs()
3585    }
3586
3587    #[cfg(test)]
3588    pub(crate) fn available_permits(&self) -> usize {
3589        self.sem.available_permits()
3590    }
3591
3592    pub(crate) fn begin_drain(&self) -> u32 {
3593        let mut state = self
3594            .state
3595            .lock()
3596            .unwrap_or_else(|poisoned| poisoned.into_inner());
3597        state.closed = true;
3598        self.sem.close();
3599        state.credits.capture_subscription_exclusions()
3600    }
3601
3602    pub(crate) fn close(&self) {
3603        self.state
3604            .lock()
3605            .unwrap_or_else(|poisoned| poisoned.into_inner())
3606            .closed = true;
3607        self.sem.close();
3608    }
3609}
3610
3611#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3612pub(crate) struct ChannelFlowClosed;
3613
3614impl fmt::Display for ChannelFlowClosed {
3615    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3616        write!(f, "flow-control window closed")
3617    }
3618}
3619
3620impl Error for ChannelFlowClosed {}
3621
3622fn window_for(concurrency: &Concurrency) -> usize {
3623    match concurrency {
3624        Concurrency::Serial => 1,
3625        Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
3626        Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
3627    }
3628}
3629
3630#[derive(Debug, Clone, PartialEq, Eq)]
3631pub enum ForwardingError {
3632    /// A connection cannot hold both client routes and a module endpoint, or
3633    /// replace a module identity while that endpoint is still registered.
3634    ConnectionRoleConflict {
3635        connection_id: ConnectionId,
3636    },
3637    NoModuleConnection,
3638    ModuleReloading {
3639        module_id: String,
3640    },
3641    StaleModuleEndpoint,
3642    UnknownReservation {
3643        client_channel: u16,
3644        module_channel: u16,
3645    },
3646    ClientRouteChannelExhausted {
3647        connection_id: ConnectionId,
3648    },
3649    ModuleRouteChannelExhausted {
3650        endpoint: ModuleEndpointId,
3651    },
3652    RelayCorrelationExhausted,
3653    ConnectionClosing {
3654        connection_id: ConnectionId,
3655    },
3656    ClientEgressClosed {
3657        connection_id: ConnectionId,
3658    },
3659    RouteOpenBuild(String),
3660    /// A swap candidate is already registered for this module id.
3661    CandidateSlotOccupied {
3662        module_id: String,
3663    },
3664    /// A registering module's HELLO_ACK could not be queued because its
3665    /// outbound queue is closed or full.
3666    ModuleEgressUnavailable {
3667        connection_id: ConnectionId,
3668    },
3669    Poisoned,
3670}
3671
3672impl fmt::Display for ForwardingError {
3673    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3674        match self {
3675            Self::ConnectionRoleConflict { connection_id } => write!(f, "connection {} already holds an incompatible role", connection_id.get()),
3676            Self::NoModuleConnection => write!(f, "no module connection is registered"),
3677            Self::ModuleReloading { module_id } => {
3678                write!(f, "module_id '{module_id}' is reloading")
3679            }
3680            Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
3681            Self::UnknownReservation {
3682                client_channel,
3683                module_channel,
3684            } => write!(
3685                f,
3686                "route reservation client channel {client_channel} / module channel {module_channel} was not found"
3687            ),
3688            Self::ClientRouteChannelExhausted { connection_id } => write!(
3689                f,
3690                "no client route channels are available for connection {}",
3691                connection_id.get()
3692            ),
3693            Self::ModuleRouteChannelExhausted { endpoint } => write!(
3694                f,
3695                "no module route channels are available for endpoint generation {} on connection {}",
3696                endpoint.generation,
3697                endpoint.connection_id.get()
3698            ),
3699            Self::RelayCorrelationExhausted => {
3700                write!(f, "module control correlation ids are exhausted")
3701            }
3702            Self::ConnectionClosing { connection_id } => write!(
3703                f,
3704                "connection {} is closing and cannot accept route allocation",
3705                connection_id.get()
3706            ),
3707            Self::ClientEgressClosed { connection_id } => write!(
3708                f,
3709                "client connection {} egress is closed",
3710                connection_id.get()
3711            ),
3712            Self::RouteOpenBuild(message) => {
3713                write!(f, "failed to prebuild route.open response: {message}")
3714            }
3715            Self::CandidateSlotOccupied { module_id } => write!(
3716                f,
3717                "module_id '{module_id}' already has a swap candidate registered"
3718            ),
3719            Self::ModuleEgressUnavailable { connection_id } => write!(
3720                f,
3721                "module connection {} egress is unavailable; HELLO_ACK could not be queued",
3722                connection_id.get()
3723            ),
3724            Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
3725        }
3726    }
3727}
3728
3729impl Error for ForwardingError {}
3730
3731#[cfg(test)]
3732mod tests {
3733    use std::time::Duration;
3734
3735    use super::*;
3736    use tokio::sync::mpsc;
3737
3738    #[test]
3739    fn ordinary_long_running_request_is_not_excluded_from_drain() {
3740        let mut ledger = CreditLedger::default();
3741        ledger.acquire(1, false);
3742
3743        assert_eq!(ledger.capture_subscription_exclusions(), 0);
3744        assert_eq!(ledger.drain_in_flight(), 1);
3745    }
3746
3747    #[test]
3748    fn bit_set_subscription_is_excluded_and_counted() {
3749        let mut ledger = CreditLedger::default();
3750        ledger.acquire(1, true);
3751
3752        assert_eq!(ledger.capture_subscription_exclusions(), 1);
3753        assert_eq!(ledger.drain_in_flight(), 0);
3754    }
3755
3756    #[test]
3757    fn subscription_opened_after_drain_snapshot_is_not_excluded() {
3758        let mut ledger = CreditLedger::default();
3759        ledger.acquire(1, true);
3760        assert_eq!(ledger.capture_subscription_exclusions(), 1);
3761
3762        ledger.acquire(2, true);
3763
3764        assert_eq!(ledger.drain_in_flight(), 1);
3765    }
3766
3767    #[test]
3768    fn drain_with_no_subscriptions_reports_zero_excluded() {
3769        let mut ledger = CreditLedger::default();
3770        assert_eq!(ledger.capture_subscription_exclusions(), 0);
3771    }
3772
3773    fn test_hello_ack(corr: u64) -> Frame {
3774        Frame::build(
3775            FrameType::HelloAck,
3776            Flags::new(false, Priority::Passive, false),
3777            0,
3778            0,
3779            corr,
3780            Vec::new(),
3781        )
3782        .unwrap()
3783    }
3784
3785    /// An acked registration whose HELLO_ACK cannot be queued must leave no
3786    /// endpoint behind: a routable module that never got its ack would read a
3787    /// route.bind first and exit. A closed queue and a full one both refuse.
3788    #[test]
3789    fn acked_registration_that_cannot_queue_its_hello_ack_inserts_nothing() {
3790        let forwarding = ForwardingTable::default();
3791
3792        let (closed_tx, closed_rx) = mpsc::channel(8);
3793        drop(closed_rx);
3794        let closed = ConnectionId::new(1);
3795        assert_eq!(
3796            forwarding.register_module_connection_acked(
3797                closed,
3798                "closed".to_string(),
3799                2,
3800                Concurrency::ModuleManaged,
3801                FrameSink::new(closed_tx),
3802                test_hello_ack(1),
3803            ),
3804            Err(ForwardingError::ModuleEgressUnavailable {
3805                connection_id: closed
3806            })
3807        );
3808
3809        let (full_tx, _full_rx) = mpsc::channel(1);
3810        let full_sink = FrameSink::new(full_tx);
3811        full_sink.try_send(test_hello_ack(99)).unwrap();
3812        let full = ConnectionId::new(2);
3813        assert_eq!(
3814            forwarding.register_module_connection_acked(
3815                full,
3816                "full".to_string(),
3817                2,
3818                Concurrency::ModuleManaged,
3819                full_sink.clone(),
3820                test_hello_ack(2),
3821            ),
3822            Err(ForwardingError::ModuleEgressUnavailable {
3823                connection_id: full
3824            })
3825        );
3826        assert_eq!(
3827            forwarding.register_candidate_module_connection_acked(
3828                full,
3829                "full".to_string(),
3830                2,
3831                Concurrency::ModuleManaged,
3832                full_sink,
3833                test_hello_ack(3),
3834            ),
3835            Err(ForwardingError::ModuleEgressUnavailable {
3836                connection_id: full
3837            })
3838        );
3839
3840        for (connection, module_id) in [(closed, "closed"), (full, "full")] {
3841            assert_eq!(
3842                forwarding
3843                    .module_endpoint_for_connection(connection)
3844                    .unwrap(),
3845                None
3846            );
3847            let (client_tx, _client_rx) = mpsc::channel(8);
3848            assert_eq!(
3849                forwarding
3850                    .begin_route_bind_relay_for_test(
3851                        ConnectionId::new(50),
3852                        FrameSink::new(client_tx),
3853                        1,
3854                        module_id,
3855                    )
3856                    .err(),
3857                Some(ForwardingError::NoModuleConnection)
3858            );
3859        }
3860        assert!(forwarding.read_inner().unwrap().candidates_by_id.is_empty());
3861    }
3862
3863    /// Both acked registration forms put the HELLO_ACK on the module's queue
3864    /// by the time the endpoint can be resolved.
3865    #[test]
3866    fn acked_registration_queues_the_hello_ack_first() {
3867        let forwarding = ForwardingTable::default();
3868        let (active_tx, mut active_rx) = mpsc::channel(8);
3869        forwarding
3870            .register_module_connection_acked(
3871                ConnectionId::new(1),
3872                "acked".to_string(),
3873                2,
3874                Concurrency::ModuleManaged,
3875                FrameSink::new(active_tx),
3876                test_hello_ack(11),
3877            )
3878            .unwrap();
3879        let (candidate_tx, mut candidate_rx) = mpsc::channel(8);
3880        forwarding
3881            .register_candidate_module_connection_acked(
3882                ConnectionId::new(2),
3883                "acked".to_string(),
3884                2,
3885                Concurrency::ModuleManaged,
3886                FrameSink::new(candidate_tx),
3887                test_hello_ack(12),
3888            )
3889            .unwrap();
3890
3891        let active_first = active_rx.try_recv().unwrap().frame;
3892        assert_eq!(active_first.header.ty, FrameType::HelloAck);
3893        assert_eq!(active_first.header.corr, 11);
3894        let candidate_first = candidate_rx.try_recv().unwrap().frame;
3895        assert_eq!(candidate_first.header.ty, FrameType::HelloAck);
3896        assert_eq!(candidate_first.header.corr, 12);
3897    }
3898
3899    #[test]
3900    fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
3901        let forwarding = ForwardingTable::default();
3902        let module_connection = ConnectionId::new(10);
3903        let exhausted_client = ConnectionId::new(20);
3904        let second_client = ConnectionId::new(30);
3905        let (module_tx, _module_rx) = mpsc::channel(1);
3906        let endpoint = forwarding
3907            .register_module_connection(
3908                module_connection,
3909                "route-limit-provider".to_string(),
3910                1,
3911                Concurrency::ModuleManaged,
3912                FrameSink::new(module_tx),
3913            )
3914            .unwrap();
3915
3916        {
3917            let mut inner = forwarding.inner.write().unwrap();
3918            for channel in 1..=u16::MAX {
3919                inner.reserved_client.insert(
3920                    ClientRouteKey {
3921                        connection_id: exhausted_client,
3922                        channel,
3923                    },
3924                    ModuleRouteKey {
3925                        endpoint,
3926                        channel: 1,
3927                    },
3928                );
3929            }
3930        }
3931
3932        let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
3933        let err = forwarding
3934            .begin_route_bind_relay_for_test(
3935                exhausted_client,
3936                FrameSink::new(exhausted_tx),
3937                1,
3938                "route-limit-provider",
3939            )
3940            .unwrap_err();
3941        assert!(matches!(
3942            err,
3943            ForwardingError::ClientRouteChannelExhausted { connection_id }
3944                if connection_id == exhausted_client
3945        ));
3946
3947        let (second_tx, _second_rx) = mpsc::channel(1);
3948        let pending = forwarding
3949            .begin_route_bind_relay_for_test(
3950                second_client,
3951                FrameSink::new(second_tx),
3952                2,
3953                "route-limit-provider",
3954            )
3955            .unwrap();
3956        assert_eq!(pending.client_channel, 1);
3957    }
3958
3959    #[test]
3960    fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
3961        let forwarding = ForwardingTable::default();
3962        let module_connection = ConnectionId::new(40);
3963        let client = ConnectionId::new(50);
3964        let (module_tx, _module_rx) = mpsc::channel(1);
3965        forwarding
3966            .register_module_connection(
3967                module_connection,
3968                "slot-reuse-provider".to_string(),
3969                1,
3970                Concurrency::ModuleManaged,
3971                FrameSink::new(module_tx),
3972            )
3973            .unwrap();
3974
3975        let (client_tx, _client_rx) = mpsc::channel(1);
3976        let client_sink = FrameSink::new(client_tx);
3977        let mut wrapped_channel = None;
3978        for index in 0..=usize::from(u16::MAX) {
3979            let pending = forwarding
3980                .begin_route_bind_relay_for_test(
3981                    client,
3982                    client_sink.clone(),
3983                    index as u64 + 1,
3984                    "slot-reuse-provider",
3985                )
3986                .unwrap();
3987            if index == usize::from(u16::MAX) {
3988                wrapped_channel = Some(pending.module_channel);
3989            }
3990            forwarding
3991                .abort_pending_relay(
3992                    pending.endpoint,
3993                    pending.corr,
3994                    RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
3995                )
3996                .unwrap();
3997        }
3998
3999        assert_eq!(wrapped_channel, Some(1));
4000    }
4001
4002    #[test]
4003    fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
4004        let forwarding = ForwardingTable::default();
4005        let client = ConnectionId::new(60);
4006        forwarding
4007            .inner
4008            .write()
4009            .unwrap()
4010            .next_client_channel
4011            .insert(client, 41);
4012
4013        let released = forwarding.cleanup_connection(client).unwrap();
4014
4015        assert!(released.is_empty());
4016        assert!(!forwarding
4017            .inner
4018            .read()
4019            .unwrap()
4020            .next_client_channel
4021            .contains_key(&client));
4022    }
4023
4024    #[test]
4025    fn stale_module_cleanup_preserves_fast_reconnect_successor() {
4026        let forwarding = ForwardingTable::default();
4027        let module_id = "fast-reconnect-provider";
4028        let first_connection = ConnectionId::new(70);
4029        let second_connection = ConnectionId::new(80);
4030        let (first_tx, _first_rx) = mpsc::channel(1);
4031        let first_endpoint = forwarding
4032            .register_module_connection(
4033                first_connection,
4034                module_id.to_string(),
4035                1,
4036                Concurrency::ModuleManaged,
4037                FrameSink::new(first_tx),
4038            )
4039            .unwrap();
4040        let (second_tx, _second_rx) = mpsc::channel(1);
4041        let second_endpoint = forwarding
4042            .register_module_connection(
4043                second_connection,
4044                module_id.to_string(),
4045                1,
4046                Concurrency::ModuleManaged,
4047                FrameSink::new(second_tx),
4048            )
4049            .unwrap();
4050        assert_ne!(first_endpoint, second_endpoint);
4051
4052        let released = forwarding.cleanup_connection(first_connection).unwrap();
4053
4054        assert!(released.is_empty());
4055        assert_eq!(
4056            forwarding
4057                .inner
4058                .read()
4059                .unwrap()
4060                .modules_by_id
4061                .get(module_id)
4062                .map(|module| module.endpoint),
4063            Some(second_endpoint)
4064        );
4065        assert!(forwarding.has_live_module_connection(module_id).unwrap());
4066        let control_rpc = forwarding
4067            .begin_module_control_rpc_for(
4068                module_id,
4069                "health.check",
4070                Instant::now() + Duration::from_secs(1),
4071            )
4072            .unwrap();
4073        assert_eq!(control_rpc.endpoint, second_endpoint);
4074    }
4075
4076    fn route_fixture(
4077        module_id: &str,
4078    ) -> (
4079        ForwardingTable,
4080        ConnectionId,
4081        ModuleEndpointId,
4082        ConnectionId,
4083        FrameSink,
4084        mpsc::Receiver<crate::router::OutboundFrame>,
4085    ) {
4086        let forwarding = ForwardingTable::default();
4087        let module_connection = ConnectionId::new(100);
4088        let client_connection = ConnectionId::new(200);
4089        let (module_tx, _module_rx) = mpsc::channel(8);
4090        let endpoint = forwarding
4091            .register_module_connection(
4092                module_connection,
4093                module_id.to_string(),
4094                2,
4095                Concurrency::ModuleManaged,
4096                FrameSink::new(module_tx),
4097            )
4098            .unwrap();
4099        let (client_tx, client_rx) = mpsc::channel(8);
4100        (
4101            forwarding,
4102            module_connection,
4103            endpoint,
4104            client_connection,
4105            FrameSink::new(client_tx),
4106            client_rx,
4107        )
4108    }
4109
4110    #[test]
4111    #[cfg(unix)]
4112    fn daemon_drain_gates_current_and_racing_provider_registrations() {
4113        let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
4114        assert_eq!(
4115            forwarding.begin_daemon_drain().unwrap(),
4116            [("provider".to_string(), endpoint)]
4117        );
4118        assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
4119        assert!(matches!(
4120            forwarding.register_module_connection(
4121                ConnectionId::new(300),
4122                "late-provider".into(),
4123                2,
4124                Concurrency::ModuleManaged,
4125                sink,
4126            ),
4127            Err(ForwardingError::ConnectionClosing { .. })
4128        ));
4129    }
4130
4131    #[test]
4132    #[cfg(unix)]
4133    fn daemon_drain_results_are_endpoint_fenced_and_removed_on_cleanup() {
4134        let (forwarding, connection, endpoint, _, _, _) = route_fixture("provider");
4135        assert_eq!(
4136            forwarding
4137                .daemon_drain_result_for_connection(connection)
4138                .unwrap(),
4139            None
4140        );
4141        forwarding
4142            .record_daemon_drain_result(endpoint, true)
4143            .unwrap();
4144        forwarding.begin_daemon_drain().unwrap();
4145        assert_eq!(
4146            forwarding
4147                .daemon_drain_result_for_connection(connection)
4148                .unwrap(),
4149            Some(false)
4150        );
4151        let replacement = ModuleEndpointId {
4152            generation: endpoint.generation + 1,
4153            ..endpoint
4154        };
4155        forwarding
4156            .record_daemon_drain_result(replacement, true)
4157            .unwrap();
4158        assert_eq!(
4159            forwarding
4160                .daemon_drain_result_for_connection(connection)
4161                .unwrap(),
4162            Some(false)
4163        );
4164        forwarding
4165            .record_daemon_drain_result(endpoint, true)
4166            .unwrap();
4167        assert_eq!(
4168            forwarding
4169                .daemon_drain_result_for_connection(connection)
4170                .unwrap(),
4171            Some(true)
4172        );
4173        assert_eq!(
4174            forwarding
4175                .daemon_drain_result_for_connection(ConnectionId::new(300))
4176                .unwrap(),
4177            Some(false)
4178        );
4179        forwarding.cleanup_connection(connection).unwrap();
4180        forwarding
4181            .record_daemon_drain_result(endpoint, true)
4182            .unwrap();
4183        assert_eq!(
4184            forwarding
4185                .daemon_drain_result_for_connection(connection)
4186                .unwrap(),
4187            Some(false)
4188        );
4189        assert!(forwarding
4190            .read_inner()
4191            .unwrap()
4192            .daemon_drain_results
4193            .is_empty());
4194    }
4195
4196    #[test]
4197    #[cfg(unix)]
4198    fn late_bind_ack_during_daemon_drain_settles_without_leaking_or_closing_module() {
4199        let (forwarding, module_connection, endpoint, client_connection, sink, mut rx) =
4200            route_fixture("provider");
4201        let mut pending =
4202            begin_test_route(&forwarding, client_connection, sink.clone(), 1, "provider");
4203        assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4204        forwarding.begin_daemon_drain().unwrap();
4205        let completion = forwarding
4206            .complete_pending_relay(
4207                module_connection,
4208                pending.corr,
4209                RouteBindRelayOutcome::Accepted,
4210            )
4211            .expect("draining admission is not a fatal module error");
4212        assert!(completion.settled);
4213        let abandoned = completion
4214            .abandoned
4215            .expect("late accepted bind needs a route GOODBYE");
4216        assert_eq!(abandoned.channel, pending.module_channel);
4217        assert!(matches!(pending.receiver.try_recv().unwrap(),
4218            RouteBindRelayOutcome::Rejected(error) if error.code == "module_reloading"));
4219        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4220        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4221        assert!(
4222            rx.try_recv().is_err(),
4223            "no route.open may be published during drain"
4224        );
4225        assert_eq!(
4226            forwarding
4227                .module_endpoint_for_connection(module_connection)
4228                .unwrap(),
4229            Some(endpoint)
4230        );
4231        assert!(
4232            !forwarding
4233                .complete_pending_relay(
4234                    module_connection,
4235                    pending.corr,
4236                    RouteBindRelayOutcome::Accepted,
4237                )
4238                .unwrap()
4239                .settled
4240        );
4241    }
4242
4243    fn test_ping(corr: u64) -> Frame {
4244        Frame::build(
4245            FrameType::Ping,
4246            Flags::new(false, Priority::Passive, false),
4247            0,
4248            0,
4249            corr,
4250            Vec::new(),
4251        )
4252        .unwrap()
4253    }
4254
4255    fn begin_test_route(
4256        forwarding: &ForwardingTable,
4257        client_connection: ConnectionId,
4258        client_sink: FrameSink,
4259        corr: u64,
4260        module_id: &str,
4261    ) -> PendingRouteBindRelay {
4262        forwarding
4263            .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
4264            .unwrap()
4265    }
4266
4267    #[test]
4268    fn module_registration_and_client_reservation_are_mutually_exclusive_in_both_orders() {
4269        for candidate in [false, true] {
4270            for register_first in [false, true] {
4271                let (forwarding, module_connection, _, client_connection, sink, _rx) =
4272                    route_fixture("target");
4273                let register = || {
4274                    if candidate {
4275                        forwarding.register_candidate_module_connection(
4276                            client_connection,
4277                            "source".into(),
4278                            2,
4279                            Concurrency::ModuleManaged,
4280                            sink.clone(),
4281                        )
4282                    } else {
4283                        forwarding.register_module_connection(
4284                            client_connection,
4285                            "source".into(),
4286                            2,
4287                            Concurrency::ModuleManaged,
4288                            sink.clone(),
4289                        )
4290                    }
4291                };
4292                if register_first {
4293                    register().unwrap();
4294                    assert!(
4295                        forwarding
4296                            .begin_route_bind_relay_for_test(
4297                                client_connection,
4298                                sink.clone(),
4299                                1,
4300                                "target",
4301                            )
4302                            .is_err(),
4303                        "a deferred route.open cannot reserve after HELLO"
4304                    );
4305                    assert!(forwarding
4306                        .cleanup_connection(client_connection)
4307                        .unwrap()
4308                        .is_empty());
4309                } else {
4310                    let pending =
4311                        begin_test_route(&forwarding, client_connection, sink.clone(), 1, "target");
4312                    assert!(
4313                        register().is_err(),
4314                        "HELLO cannot register after a route reservation"
4315                    );
4316                    forwarding
4317                        .complete_pending_relay(
4318                            module_connection,
4319                            pending.corr,
4320                            RouteBindRelayOutcome::Accepted,
4321                        )
4322                        .unwrap();
4323                    assert_eq!(
4324                        forwarding
4325                            .cleanup_connection(client_connection)
4326                            .unwrap()
4327                            .len(),
4328                        1
4329                    );
4330                }
4331                assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4332                assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4333            }
4334        }
4335    }
4336
4337    /// A pending route.open reserves its queue slot with `reserve_owned` and
4338    /// sends its prebuilt response through that slot when the module accepts.
4339    /// A queue already holding many data frames (far more than the old 64-frame
4340    /// queue) must not stop the open from reserving or completing, and the
4341    /// response must be charged to the queue's byte count when it is sent.
4342    #[tokio::test]
4343    async fn pending_route_open_completes_behind_queued_data_frames() {
4344        assert_eq!(
4345            crate::server::MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
4346            8,
4347            "the per-connection pending route.open limit is its own constant"
4348        );
4349        let (forwarding, module_connection, _endpoint, client, _unused_sink, _unused_rx) =
4350            route_fixture("open-behind-data");
4351        let (sink, mut client_rx) = crate::server::connection_egress();
4352        const DATA_FRAMES: usize = 1_000;
4353        let data = |corr: u64| {
4354            Frame::build(
4355                FrameType::StreamData,
4356                Flags::new(false, Priority::Interactive, false),
4357                9,
4358                1,
4359                corr,
4360                vec![b'x'; 200],
4361            )
4362            .unwrap()
4363        };
4364        for corr in 0..DATA_FRAMES as u64 {
4365            sink.try_send(data(corr)).unwrap();
4366        }
4367        let data_bytes = DATA_FRAMES * (subc_protocol::HEADER_LEN + 200);
4368        assert_eq!(sink.backlog().queued_bytes, data_bytes);
4369
4370        let pending = tokio::time::timeout(
4371            Duration::from_secs(5),
4372            forwarding.begin_route_bind_relay_for(
4373                client,
4374                sink.clone(),
4375                subc_protocol::PROTOCOL_VERSION,
4376                4_242,
4377                "open-behind-data",
4378                Principal::Direct,
4379                None,
4380                None,
4381                Instant::now() + Duration::from_secs(60),
4382            ),
4383        )
4384        .await
4385        .expect("reserving the route.open slot must not wait behind data frames")
4386        .unwrap();
4387        forwarding
4388            .complete_pending_relay(
4389                module_connection,
4390                pending.corr,
4391                RouteBindRelayOutcome::Accepted,
4392            )
4393            .unwrap();
4394
4395        let backlog = sink.backlog();
4396        assert_eq!(backlog.queued_frames, DATA_FRAMES + 1);
4397        assert!(
4398            backlog.queued_bytes > data_bytes,
4399            "the route.open response must be counted in queued bytes: {backlog:?}"
4400        );
4401        for corr in 0..DATA_FRAMES as u64 {
4402            assert_eq!(client_rx.recv().await.unwrap().header.corr, corr);
4403        }
4404        let open = client_rx.recv().await.unwrap();
4405        assert_eq!(open.header.corr, 4_242);
4406        assert_eq!(open.header.ty, FrameType::Response);
4407        drop(open);
4408        assert_eq!(sink.backlog().queued_bytes, 0);
4409        assert_eq!(sink.backlog().queued_frames, 0);
4410    }
4411
4412    /// The drain-timeout line reports these numbers, so they must count the
4413    /// requests the drain is waiting on and no others: a route with nothing in
4414    /// flight is not a holdout, and a flagged subscription is excluded from the
4415    /// drain and so from the count.
4416    #[tokio::test]
4417    async fn drain_holdouts_count_held_requests_and_name_the_connection() {
4418        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4419            route_fixture("holdouts");
4420        let mut bound = |corr| {
4421            let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
4422            forwarding
4423                .complete_pending_relay(
4424                    module_connection,
4425                    route.corr,
4426                    RouteBindRelayOutcome::Accepted,
4427                )
4428                .unwrap();
4429            client_rx.try_recv().unwrap();
4430            match forwarding
4431                .lookup_data_route(client, route.client_channel, route.client_epoch)
4432                .unwrap()
4433            {
4434                DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4435                other => panic!("expected live route, got {other:?}"),
4436            }
4437        };
4438        let holding = bound(61);
4439        let _idle = bound(62);
4440        holding.flow.acquire_tagged(7, false).await.unwrap();
4441        holding.flow.acquire_tagged(2, false).await.unwrap();
4442        holding.flow.acquire_tagged(3, true).await.unwrap();
4443        forwarding
4444            .begin_module_drain("holdouts", RouteCloseReason::Restart)
4445            .unwrap();
4446
4447        let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
4448        assert_eq!(
4449            holdouts,
4450            DrainHoldouts {
4451                requests: 2,
4452                routes: 1,
4453                total_routes: 2,
4454                top_connections: vec![(client.get(), 2)],
4455                // The held requests by the module's channel and corr, ascending,
4456                // without the excluded subscription (corr 3).
4457                held: vec![(holding.module_channel, 2), (holding.module_channel, 7)],
4458            }
4459        );
4460    }
4461
4462    #[test]
4463    fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
4464        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
4465            route_fixture("census");
4466        let pending = begin_test_route(&forwarding, client, sink, 1, "census");
4467        forwarding
4468            .complete_pending_relay(
4469                module_connection,
4470                pending.corr,
4471                RouteBindRelayOutcome::Accepted,
4472            )
4473            .unwrap();
4474
4475        let routes = forwarding.endpoint_routes(endpoint).unwrap();
4476        assert_eq!(routes.len(), 1);
4477        assert!(matches!(routes[0].principal, Principal::Direct));
4478        assert_eq!(routes[0].goodbye_target.connection_id, client);
4479        assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
4480        assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
4481        assert!(!routes[0].draining);
4482
4483        forwarding
4484            .begin_module_drain("census", RouteCloseReason::Restart)
4485            .unwrap();
4486        let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
4487        assert_eq!(draining_routes.len(), 1);
4488        assert!(draining_routes[0].draining);
4489    }
4490
4491    #[test]
4492    fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
4493        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
4494        let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
4495        assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
4496        forwarding
4497            .abort_pending_relay(
4498                first.endpoint,
4499                first.corr,
4500                RouteBindRelayOutcome::ModuleGone("abort".into()),
4501            )
4502            .unwrap();
4503        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
4504        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
4505
4506        let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
4507        assert_eq!(second.client_channel, first.client_channel);
4508        assert_eq!(second.module_channel, first.module_channel);
4509        assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
4510    }
4511
4512    #[test]
4513    fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
4514        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4515            route_fixture("epoch-release");
4516        let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
4517        forwarding
4518            .complete_pending_relay(
4519                module_connection,
4520                first.corr,
4521                RouteBindRelayOutcome::Accepted,
4522            )
4523            .unwrap();
4524        assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
4525        assert!(matches!(
4526            forwarding
4527                .release_client_route(client, first.client_channel, first.client_epoch)
4528                .unwrap(),
4529            RouteRelease::Removed(_)
4530        ));
4531        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
4532        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
4533
4534        let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
4535        forwarding
4536            .complete_pending_relay(
4537                module_connection,
4538                second.corr,
4539                RouteBindRelayOutcome::Accepted,
4540            )
4541            .unwrap();
4542        assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
4543        assert!(matches!(
4544            forwarding
4545                .release_client_route(client, second.client_channel, first.client_epoch)
4546                .unwrap(),
4547            RouteRelease::Stale
4548        ));
4549        assert!(!forwarding
4550            .cache_status(
4551                endpoint,
4552                second.module_channel,
4553                first.module_epoch,
4554                "stale".into(),
4555            )
4556            .unwrap());
4557        assert!(forwarding
4558            .cache_status(
4559                endpoint,
4560                second.module_channel,
4561                second.module_epoch,
4562                "current".into(),
4563            )
4564            .unwrap());
4565        match forwarding
4566            .route_poll_snapshot(client, second.client_channel, second.client_epoch)
4567            .unwrap()
4568        {
4569            RoutePollSnapshot::Bound { status, .. } => {
4570                assert_eq!(status.as_deref(), Some("current"));
4571            }
4572            RoutePollSnapshot::Absent => panic!("successor binding was removed"),
4573        }
4574        let counters = forwarding.counters().snapshot();
4575        assert_eq!(counters["route_released_epoch_fenced"], 1);
4576        assert_eq!(counters["route_release_stale_skipped"], 1);
4577    }
4578
4579    #[test]
4580    fn max_epoch_reservation_retires_only_that_slot() {
4581        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
4582        forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
4583        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
4584        let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
4585        assert_eq!(
4586            (final_use.client_channel, final_use.client_epoch),
4587            (7, u32::MAX)
4588        );
4589        assert_eq!(
4590            (final_use.module_channel, final_use.module_epoch),
4591            (9, u32::MAX)
4592        );
4593        forwarding
4594            .abort_pending_relay(
4595                endpoint,
4596                final_use.corr,
4597                RouteBindRelayOutcome::ModuleGone("abort".into()),
4598            )
4599            .unwrap();
4600        forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
4601        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
4602        let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
4603        assert_ne!(next.client_channel, 7);
4604        assert_ne!(next.module_channel, 9);
4605        assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
4606    }
4607
4608    #[test]
4609    fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
4610        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
4611            route_fixture("corr-shared");
4612        let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
4613        assert_eq!(bind.corr, 1);
4614        forwarding
4615            .abort_pending_relay(
4616                endpoint,
4617                bind.corr,
4618                RouteBindRelayOutcome::ModuleGone("abort".into()),
4619            )
4620            .unwrap();
4621        let rpc = forwarding
4622            .begin_module_control_rpc_for(
4623                "corr-shared",
4624                "health.check",
4625                Instant::now() - Duration::from_millis(1),
4626            )
4627            .unwrap();
4628        assert_eq!(rpc.corr, 2);
4629        assert_eq!(
4630            forwarding
4631                .complete_module_control_rpc(
4632                    module_connection,
4633                    rpc.corr,
4634                    Some("health.check"),
4635                    ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
4636                        status: subc_protocol::session::HealthStatus::Ok,
4637                        detail: None,
4638                        metrics: None,
4639                    }),
4640                )
4641                .unwrap(),
4642            ModuleControlRpcCompletion::Settled
4643        );
4644        assert!(matches!(
4645            rpc.receiver.blocking_recv().unwrap(),
4646            ModuleControlRpcOutcome::DeadlineElapsed
4647        ));
4648    }
4649
4650    #[tokio::test(start_paused = true)]
4651    async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
4652        let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
4653        let probe_started_at = Instant::now();
4654        let rpc = forwarding
4655            .begin_health_probe_rpc_for(
4656                "tombstone-ttl",
4657                "health.check",
4658                probe_started_at,
4659                probe_started_at + Duration::from_secs(5),
4660            )
4661            .unwrap();
4662        assert!(forwarding
4663            .tombstone_health_probe_rpc(endpoint, rpc.corr)
4664            .unwrap());
4665        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
4666
4667        tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
4668        tokio::task::yield_now().await;
4669
4670        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
4671    }
4672
4673    #[test]
4674    fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
4675        let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
4676        let mut close = forwarding.register_connection_close(endpoint.connection_id);
4677        forwarding.inject_control_corr(endpoint, u64::MAX);
4678        let final_rpc = forwarding
4679            .begin_module_control_rpc_for(
4680                "corr-max",
4681                "health.check",
4682                Instant::now() + Duration::from_secs(1),
4683            )
4684            .unwrap();
4685        assert_eq!(final_rpc.corr, u64::MAX);
4686        forwarding
4687            .cancel_module_control_rpc(endpoint, final_rpc.corr)
4688            .unwrap();
4689        assert!(matches!(
4690            forwarding.begin_module_control_rpc_for(
4691                "corr-max",
4692                "health.check",
4693                Instant::now() + Duration::from_secs(1),
4694            ),
4695            Err(ForwardingError::RelayCorrelationExhausted)
4696        ));
4697        assert!(close.try_recv().is_ok());
4698    }
4699
4700    #[test]
4701    fn publication_epoch_controls_delivery_failure_escalation() {
4702        fn setup_successor(
4703            commit_successor: Option<bool>,
4704        ) -> (ForwardingTable, ConnectionId, u16, u32) {
4705            let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4706                route_fixture("escalation");
4707            let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
4708            forwarding
4709                .complete_pending_relay(
4710                    module_connection,
4711                    first.corr,
4712                    RouteBindRelayOutcome::Accepted,
4713                )
4714                .unwrap();
4715            client_rx.try_recv().unwrap();
4716            assert!(matches!(
4717                forwarding
4718                    .release_client_route(client, first.client_channel, first.client_epoch)
4719                    .unwrap(),
4720                RouteRelease::Removed(_)
4721            ));
4722            if let Some(commit_successor) = commit_successor {
4723                forwarding.inject_client_slot_epoch(
4724                    client,
4725                    first.client_channel,
4726                    first.client_epoch,
4727                );
4728                forwarding.inject_module_slot_epoch(
4729                    endpoint,
4730                    first.module_channel,
4731                    first.module_epoch,
4732                );
4733                let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
4734                if commit_successor {
4735                    forwarding
4736                        .complete_pending_relay(
4737                            module_connection,
4738                            successor.corr,
4739                            RouteBindRelayOutcome::Accepted,
4740                        )
4741                        .unwrap();
4742                    client_rx.try_recv().unwrap();
4743                } else {
4744                    forwarding
4745                        .abort_pending_relay(
4746                            endpoint,
4747                            successor.corr,
4748                            RouteBindRelayOutcome::ModuleGone("abort".into()),
4749                        )
4750                        .unwrap();
4751                }
4752            }
4753            (forwarding, client, first.client_channel, first.client_epoch)
4754        }
4755
4756        let probe_sink = FrameSink::new(mpsc::channel(1).0);
4757        let (no_successor, client, channel, epoch) = setup_successor(None);
4758        let mut close = no_successor.register_connection_close(client);
4759        assert!(no_successor
4760            .escalate_client_delivery_failure(
4761                client,
4762                channel,
4763                epoch,
4764                CloseReason::new("delivery", "failed"),
4765                UndeliveredFrame {
4766                    module_id: None,
4767                    sink: &probe_sink,
4768                },
4769            )
4770            .unwrap());
4771        assert!(close.try_recv().is_ok());
4772
4773        let (aborted, client, channel, epoch) = setup_successor(Some(false));
4774        let mut close = aborted.register_connection_close(client);
4775        assert!(aborted
4776            .escalate_client_delivery_failure(
4777                client,
4778                channel,
4779                epoch,
4780                CloseReason::new("delivery", "failed"),
4781                UndeliveredFrame {
4782                    module_id: None,
4783                    sink: &probe_sink,
4784                },
4785            )
4786            .unwrap());
4787        assert!(close.try_recv().is_ok());
4788
4789        let (published, client, channel, epoch) = setup_successor(Some(true));
4790        let mut close = published.register_connection_close(client);
4791        assert!(!published
4792            .escalate_client_delivery_failure(
4793                client,
4794                channel,
4795                epoch,
4796                CloseReason::new("delivery", "stale failure"),
4797                UndeliveredFrame {
4798                    module_id: None,
4799                    sink: &probe_sink,
4800                },
4801            )
4802            .unwrap());
4803        assert!(close.try_recv().is_err());
4804    }
4805
4806    #[test]
4807    fn route_concentration_separates_client_count_from_routes_per_client() {
4808        // The distinction this asserts is the one a bare connection count cannot
4809        // make: two connections holding one route each and one connection
4810        // holding two are the same total, and have opposite causes.
4811        let (forwarding, module_connection, _, client, sink, _client_rx) =
4812            route_fixture("concentration");
4813        assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
4814
4815        for corr in [70_u64, 71] {
4816            let pending =
4817                begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
4818            forwarding
4819                .complete_pending_relay(
4820                    module_connection,
4821                    pending.corr,
4822                    RouteBindRelayOutcome::Accepted,
4823                )
4824                .unwrap();
4825        }
4826
4827        // One connection, two routes — not two connections with a route each.
4828        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4829        assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
4830    }
4831
4832    #[test]
4833    fn cleanup_and_accepted_resolution_have_one_lock_winner() {
4834        let (forwarding, module_connection, _, client, sink, mut client_rx) =
4835            route_fixture("cleanup-race");
4836        let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
4837        forwarding
4838            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4839            .unwrap();
4840        let released = forwarding.cleanup_connection(client).unwrap();
4841        assert_eq!(released.len(), 1);
4842        let completion = forwarding
4843            .complete_pending_relay(
4844                module_connection,
4845                pending.corr,
4846                RouteBindRelayOutcome::Accepted,
4847            )
4848            .unwrap();
4849        assert!(!completion.settled);
4850        assert!(client_rx.try_recv().is_err());
4851        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4852
4853        let (forwarding, module_connection, _, client, sink, mut client_rx) =
4854            route_fixture("accepted-race");
4855        let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
4856        forwarding
4857            .complete_pending_relay(
4858                module_connection,
4859                pending.corr,
4860                RouteBindRelayOutcome::Accepted,
4861            )
4862            .unwrap();
4863        assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
4864        let released = forwarding.cleanup_connection(client).unwrap();
4865        assert_eq!(released.len(), 1);
4866        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4867    }
4868
4869    #[test]
4870    fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
4871        let (forwarding, module_connection, _, client, sink, mut client_rx) =
4872            route_fixture("drain-gap");
4873        let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
4874        forwarding
4875            .complete_pending_relay(
4876                module_connection,
4877                live.corr,
4878                RouteBindRelayOutcome::Accepted,
4879            )
4880            .unwrap();
4881        client_rx.try_recv().unwrap();
4882        let binding = match forwarding
4883            .lookup_data_route(client, live.client_channel, live.client_epoch)
4884            .unwrap()
4885        {
4886            DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4887            other => panic!("expected live route, got {other:?}"),
4888        };
4889
4890        let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
4891        forwarding
4892            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4893            .unwrap();
4894        let control_rpc = forwarding
4895            .begin_module_control_rpc_for(
4896                "drain-gap",
4897                "health.check",
4898                Instant::now() + Duration::from_secs(1),
4899            )
4900            .unwrap();
4901        let target = forwarding
4902            .begin_module_drain("drain-gap", RouteCloseReason::Reload)
4903            .unwrap()
4904            .unwrap();
4905        assert!(matches!(
4906            control_rpc.receiver.blocking_recv().unwrap(),
4907            ModuleControlRpcOutcome::ModuleGone(_)
4908        ));
4909        assert_eq!(target.abandoned_bindings.len(), 1);
4910        assert!(binding.flow.sem.is_closed());
4911        assert!(
4912            !forwarding
4913                .complete_pending_relay(
4914                    module_connection,
4915                    pending.corr,
4916                    RouteBindRelayOutcome::Accepted,
4917                )
4918                .unwrap()
4919                .settled
4920        );
4921        assert!(matches!(
4922            forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
4923            Err(ForwardingError::ModuleReloading { .. })
4924        ));
4925        let released = forwarding
4926            .release_module_endpoint_routes(target.endpoint)
4927            .unwrap();
4928        assert_eq!(released.len(), 1);
4929        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4930    }
4931
4932    /// A client can be marked closing while its egress is still open: the daemon
4933    /// asks a connection to close (here through the production path, a module
4934    /// frame that its egress refused) and the connection loop tears down a moment
4935    /// later. A route.bind ack that lands inside that window is answered on the
4936    /// MODULE connection's frame handler, so resolving it must not produce an
4937    /// error -- an error there ends the module connection, and that connection
4938    /// carries every other client's routes to the module.
4939    #[test]
4940    fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
4941        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4942            route_fixture("closing-client");
4943
4944        // A published route on this client: escalate_client_delivery_failure only
4945        // marks a connection closing for a route it has already published.
4946        let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
4947        forwarding
4948            .complete_pending_relay(
4949                module_connection,
4950                live.corr,
4951                RouteBindRelayOutcome::Accepted,
4952            )
4953            .unwrap();
4954        client_rx.try_recv().unwrap();
4955
4956        // A second route.open from the same client, relayed and awaiting its ack.
4957        let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
4958        forwarding
4959            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4960            .unwrap();
4961
4962        // The window: closing, but the sink is still open.
4963        assert!(forwarding
4964            .escalate_client_delivery_failure(
4965                client,
4966                live.client_channel,
4967                live.client_epoch,
4968                CloseReason::new(
4969                    "module_to_client_delivery_failed",
4970                    "client egress refused a module frame",
4971                ),
4972                UndeliveredFrame {
4973                    module_id: None,
4974                    sink: &sink,
4975                },
4976            )
4977            .unwrap());
4978        assert!(!sink.is_closed());
4979
4980        let completion = forwarding
4981            .complete_pending_relay(
4982                module_connection,
4983                pending.corr,
4984                RouteBindRelayOutcome::Accepted,
4985            )
4986            .expect("a closing client must not turn a module's ack into an error");
4987
4988        assert!(completion.settled);
4989        let abandoned = completion
4990            .abandoned
4991            .expect("the module must be told to drop the binding it just created");
4992        assert_eq!(abandoned.connection_id, module_connection);
4993        assert_eq!(abandoned.channel, pending.module_channel);
4994        assert_eq!(abandoned.epoch, pending.module_epoch);
4995        assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
4996        assert!(matches!(
4997            pending.receiver.blocking_recv().unwrap(),
4998            RouteBindRelayOutcome::ModuleGone(_)
4999        ));
5000        // No route was published to a client that is on its way out, and the
5001        // reserved handle pair went back.
5002        assert!(client_rx.try_recv().is_err());
5003        assert_eq!(forwarding.active_binding_count().unwrap(), 1);
5004
5005        // The module endpoint is untouched: still live, and still able to take a
5006        // route from another client.
5007        assert!(forwarding
5008            .has_live_module_connection("closing-client")
5009            .unwrap());
5010        let cotenant = ConnectionId::new(201);
5011        let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
5012        let cotenant_route = begin_test_route(
5013            &forwarding,
5014            cotenant,
5015            FrameSink::new(cotenant_tx),
5016            62,
5017            "closing-client",
5018        );
5019        assert_eq!(cotenant_route.endpoint, endpoint);
5020        forwarding
5021            .complete_pending_relay(
5022                module_connection,
5023                cotenant_route.corr,
5024                RouteBindRelayOutcome::Accepted,
5025            )
5026            .unwrap();
5027        assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
5028        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
5029    }
5030
5031    #[test]
5032    fn pending_route_permit_is_released_on_rejection_and_abort() {
5033        let forwarding = ForwardingTable::default();
5034        let module_connection = ConnectionId::new(300);
5035        let client = ConnectionId::new(301);
5036        let (module_tx, _module_rx) = mpsc::channel(1);
5037        let endpoint = forwarding
5038            .register_module_connection(
5039                module_connection,
5040                "permit".into(),
5041                2,
5042                Concurrency::ModuleManaged,
5043                FrameSink::new(module_tx),
5044            )
5045            .unwrap();
5046        let (client_tx, mut client_rx) = mpsc::channel(1);
5047        let sink = FrameSink::new(client_tx);
5048        let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
5049        assert!(sink.try_send(test_ping(999)).is_err());
5050        forwarding
5051            .complete_pending_relay(
5052                module_connection,
5053                rejected.corr,
5054                RouteBindRelayOutcome::Rejected(ErrorBody {
5055                    code: "no".into(),
5056                    message: "rejected".into(),
5057                    detail: None,
5058                }),
5059            )
5060            .unwrap();
5061        sink.try_send(test_ping(1000)).unwrap();
5062        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
5063
5064        let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
5065        assert!(sink.try_send(test_ping(1001)).is_err());
5066        forwarding
5067            .abort_pending_relay(
5068                endpoint,
5069                aborted.corr,
5070                RouteBindRelayOutcome::ModuleGone("abort".into()),
5071            )
5072            .unwrap();
5073        sink.try_send(test_ping(1002)).unwrap();
5074        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
5075
5076        let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
5077        forwarding
5078            .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
5079            .unwrap();
5080        drop(client_rx);
5081        let completion = forwarding
5082            .complete_pending_relay(
5083                module_connection,
5084                receiver_closed.corr,
5085                RouteBindRelayOutcome::Accepted,
5086            )
5087            .unwrap();
5088        assert!(completion.abandoned.is_some());
5089        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
5090    }
5091
5092    /// Every connection teardown passes through `cleanup_connection`, and
5093    /// connection ids come from a monotonic counter that never hands an id out
5094    /// twice. If the closing mark survives teardown, the set grows by one entry
5095    /// per connection for the life of the daemon -- the self-watchdog alone
5096    /// reconnects once a minute.
5097    #[test]
5098    fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
5099        let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
5100            route_fixture("closing-set-leak");
5101
5102        const CONNECTIONS: u64 = 32;
5103        for index in 0..CONNECTIONS {
5104            let client = ConnectionId::new(1000 + index);
5105            let (client_tx, _client_rx) = mpsc::channel(8);
5106            let route = begin_test_route(
5107                &forwarding,
5108                client,
5109                FrameSink::new(client_tx),
5110                index + 1,
5111                "closing-set-leak",
5112            );
5113            forwarding
5114                .complete_pending_relay(
5115                    module_connection,
5116                    route.corr,
5117                    RouteBindRelayOutcome::Accepted,
5118                )
5119                .unwrap();
5120            forwarding.cleanup_connection(client).unwrap();
5121        }
5122        forwarding.cleanup_connection(module_connection).unwrap();
5123
5124        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
5125    }
5126
5127    /// The closing mark exists to refuse new work for a connection that is on
5128    /// its way out but whose teardown has not run yet: the daemon asks the
5129    /// connection loop to end, and only when the loop reacts does
5130    /// `cleanup_connection` strip the connection's state. Inside that window an
5131    /// operation for the dying connection must still be refused; only after
5132    /// cleanup completes may the mark go.
5133    #[test]
5134    fn closing_connection_is_refused_new_work_until_cleanup_completes() {
5135        let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
5136            route_fixture("closing-gate");
5137
5138        // A published route: escalate_client_delivery_failure only marks a
5139        // connection closing for a route it has already published.
5140        let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
5141        forwarding
5142            .complete_pending_relay(
5143                module_connection,
5144                live.corr,
5145                RouteBindRelayOutcome::Accepted,
5146            )
5147            .unwrap();
5148        client_rx.try_recv().unwrap();
5149
5150        // Mark the connection closing through the production path without
5151        // running teardown, pinning the window open.
5152        assert!(forwarding
5153            .escalate_client_delivery_failure(
5154                client,
5155                live.client_channel,
5156                live.client_epoch,
5157                CloseReason::new(
5158                    "module_to_client_delivery_failed",
5159                    "client egress refused a module frame",
5160                ),
5161                UndeliveredFrame {
5162                    module_id: None,
5163                    sink: &sink,
5164                },
5165            )
5166            .unwrap());
5167        assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
5168
5169        // A late route.open for the closing client is refused ...
5170        assert!(matches!(
5171            forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
5172            Err(ForwardingError::ConnectionClosing { connection_id })
5173                if connection_id == client
5174        ));
5175        // ... and so is a late attempt to register the connection as a module.
5176        let (late_tx, _late_rx) = mpsc::channel(1);
5177        assert!(matches!(
5178            forwarding.register_module_connection(
5179                client,
5180                "late-module".into(),
5181                2,
5182                Concurrency::ModuleManaged,
5183                FrameSink::new(late_tx),
5184            ),
5185            Err(ForwardingError::ConnectionClosing { connection_id })
5186                if connection_id == client
5187        ));
5188
5189        // Teardown is the point that lifts the mark: it has just removed every
5190        // per-connection entry under the same lock, so the gate has nothing
5191        // left to protect for this id.
5192        forwarding.cleanup_connection(client).unwrap();
5193        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
5194    }
5195}
5196
5197/// Blue/green swap slots: a candidate registered beside the active endpoint,
5198/// promoted by `cutover_candidate`, with the old incumbent drained by endpoint.
5199#[cfg(test)]
5200mod swap_slot_tests {
5201    use std::time::Duration;
5202
5203    use super::*;
5204    use tokio::sync::mpsc;
5205
5206    const MODULE_ID: &str = "swapped";
5207
5208    struct SwapFixture {
5209        forwarding: ForwardingTable,
5210        incumbent_connection: ConnectionId,
5211        incumbent: ModuleEndpointId,
5212        candidate_connection: ConnectionId,
5213        candidate: ModuleEndpointId,
5214        _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
5215    }
5216
5217    fn swap_fixture() -> SwapFixture {
5218        let forwarding = ForwardingTable::default();
5219        let incumbent_connection = ConnectionId::new(100);
5220        let candidate_connection = ConnectionId::new(110);
5221        let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
5222        let incumbent = forwarding
5223            .register_module_connection(
5224                incumbent_connection,
5225                MODULE_ID.to_string(),
5226                2,
5227                Concurrency::ModuleManaged,
5228                FrameSink::new(incumbent_tx),
5229            )
5230            .unwrap();
5231        let (candidate_tx, candidate_rx) = mpsc::channel(8);
5232        let candidate = forwarding
5233            .register_candidate_module_connection(
5234                candidate_connection,
5235                MODULE_ID.to_string(),
5236                2,
5237                Concurrency::ModuleManaged,
5238                FrameSink::new(candidate_tx),
5239            )
5240            .unwrap();
5241        SwapFixture {
5242            forwarding,
5243            incumbent_connection,
5244            incumbent,
5245            candidate_connection,
5246            candidate,
5247            _module_rxs: vec![incumbent_rx, candidate_rx],
5248        }
5249    }
5250
5251    fn client(
5252        raw: u64,
5253    ) -> (
5254        ConnectionId,
5255        FrameSink,
5256        mpsc::Receiver<crate::router::OutboundFrame>,
5257    ) {
5258        let (tx, rx) = mpsc::channel(8);
5259        (ConnectionId::new(raw), FrameSink::new(tx), rx)
5260    }
5261
5262    fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
5263        forwarding
5264            .read_inner()
5265            .unwrap()
5266            .client_to_module
5267            .values()
5268            .map(|route| route.module_endpoint)
5269            .collect()
5270    }
5271
5272    #[test]
5273    fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
5274        let fixture = swap_fixture();
5275        let forwarding = &fixture.forwarding;
5276        assert_ne!(fixture.incumbent, fixture.candidate);
5277
5278        // Every by-id consumer still resolves the incumbent.
5279        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
5280        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
5281        let (client_connection, client_sink, _client_rx) = client(200);
5282        let pending = forwarding
5283            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5284            .unwrap();
5285        assert_eq!(pending.endpoint, fixture.incumbent);
5286        let rpc = forwarding
5287            .begin_module_control_rpc_for(
5288                MODULE_ID,
5289                "health.check",
5290                Instant::now() + Duration::from_secs(1),
5291            )
5292            .unwrap();
5293        assert_eq!(rpc.endpoint, fixture.incumbent);
5294        let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
5295        assert_eq!(census.len(), 1, "the census lists one endpoint per id");
5296
5297        // Connection-keyed lookups see the candidate, so its own frames resolve.
5298        assert_eq!(
5299            forwarding
5300                .module_endpoint_for_connection(fixture.candidate_connection)
5301                .unwrap(),
5302            Some(fixture.candidate)
5303        );
5304        assert_eq!(
5305            forwarding
5306                .module_id_for_connection(fixture.candidate_connection)
5307                .unwrap()
5308                .as_deref(),
5309            Some(MODULE_ID)
5310        );
5311
5312        // One candidate per id.
5313        let (other_tx, _other_rx) = mpsc::channel(1);
5314        assert_eq!(
5315            forwarding.register_candidate_module_connection(
5316                ConnectionId::new(120),
5317                MODULE_ID.to_string(),
5318                2,
5319                Concurrency::ModuleManaged,
5320                FrameSink::new(other_tx),
5321            ),
5322            Err(ForwardingError::CandidateSlotOccupied {
5323                module_id: MODULE_ID.to_string()
5324            })
5325        );
5326    }
5327
5328    /// The cutover linearization point. A relay reserved on the incumbent before
5329    /// cutover must never become a route on the incumbent, and every relay
5330    /// reserved after it must land on the promoted candidate.
5331    #[test]
5332    fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
5333        let fixture = swap_fixture();
5334        let forwarding = &fixture.forwarding;
5335        let (early_client, early_sink, _early_rx) = client(200);
5336        let mut early = forwarding
5337            .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
5338            .unwrap();
5339        assert_eq!(early.endpoint, fixture.incumbent);
5340        assert_eq!(
5341            forwarding
5342                .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
5343                .unwrap(),
5344            RouteBindRelayMark::Marked
5345        );
5346
5347        let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
5348        assert_eq!(
5349            cutover,
5350            ForwardingCutover {
5351                promoted: fixture.candidate,
5352                incumbent: Some(fixture.incumbent),
5353            }
5354        );
5355
5356        // A route.open reserved after cutover goes to the promoted candidate.
5357        let (late_client, late_sink, _late_rx) = client(201);
5358        let late = forwarding
5359            .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
5360            .unwrap();
5361        assert_eq!(
5362            late.endpoint, fixture.candidate,
5363            "a route.open after cutover was reserved on the incumbent"
5364        );
5365
5366        // The incumbent acks the early relay after cutover.
5367        let completion = forwarding
5368            .complete_pending_relay(
5369                fixture.incumbent_connection,
5370                early.corr,
5371                RouteBindRelayOutcome::Accepted,
5372            )
5373            .expect("a superseded endpoint's ack is not an error on its connection");
5374        assert!(completion.settled);
5375        assert!(
5376            !committed_endpoints(forwarding).contains(&fixture.incumbent),
5377            "a relay reserved before cutover committed a route on the incumbent"
5378        );
5379        let goodbye = completion
5380            .abandoned
5381            .expect("the incumbent is told to drop the binding it just created");
5382        assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
5383        assert_eq!(goodbye.channel, early.module_channel);
5384        assert_eq!(goodbye.epoch, early.module_epoch);
5385        assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
5386        match early.receiver.try_recv() {
5387            Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
5388            other => panic!("expected a retryable module_reloading answer, got {other:?}"),
5389        }
5390        assert!(matches!(
5391            forwarding
5392                .lookup_data_route(early_client, early.client_channel, early.client_epoch)
5393                .unwrap(),
5394            DataRoute::Client(DataRouteState::Absent)
5395        ));
5396
5397        // The late relay commits on the candidate; only its pair stays reserved
5398        // until then.
5399        assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
5400        forwarding
5401            .complete_pending_relay(
5402                fixture.candidate_connection,
5403                late.corr,
5404                RouteBindRelayOutcome::Accepted,
5405            )
5406            .unwrap();
5407        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
5408        assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
5409    }
5410
5411    #[test]
5412    fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
5413        let fixture = swap_fixture();
5414        let forwarding = &fixture.forwarding;
5415        // One bound route and one in-flight relay on the incumbent.
5416        let (bound_client, bound_sink, _bound_rx) = client(200);
5417        let bound = forwarding
5418            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
5419            .unwrap();
5420        forwarding
5421            .complete_pending_relay(
5422                fixture.incumbent_connection,
5423                bound.corr,
5424                RouteBindRelayOutcome::Accepted,
5425            )
5426            .unwrap();
5427        let (pending_client, pending_sink, _pending_rx) = client(201);
5428        let mut in_flight = forwarding
5429            .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
5430            .unwrap();
5431        forwarding
5432            .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
5433            .unwrap();
5434
5435        let incumbent = forwarding
5436            .cutover_candidate(MODULE_ID)
5437            .unwrap()
5438            .unwrap()
5439            .incumbent
5440            .unwrap();
5441        let target = forwarding
5442            .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
5443            .unwrap()
5444            .expect("the superseded incumbent is still registered");
5445
5446        assert_eq!(target.endpoint, fixture.incumbent);
5447        assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
5448        assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
5449        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
5450        assert_eq!(target.abandoned_bindings.len(), 1);
5451        assert_eq!(
5452            target.abandoned_bindings[0].channel,
5453            in_flight.module_channel
5454        );
5455        assert!(matches!(
5456            in_flight.receiver.try_recv(),
5457            Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
5458        ));
5459        assert_eq!(
5460            forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
5461            1,
5462            "the incumbent's bound route stays until its drain finishes"
5463        );
5464
5465        let (next_client, next_sink, _next_rx) = client(202);
5466        let next = forwarding
5467            .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
5468            .expect("the promoted candidate keeps accepting routes");
5469        assert_eq!(next.endpoint, fixture.candidate);
5470    }
5471
5472    /// An endpoint replaced WITHOUT a promotion (a successor registered over it
5473    /// as an ordinary active HELLO) is stale, not superseded, and its ack still
5474    /// fails the acking connection exactly as it did before swap slots existed:
5475    /// `StaleModuleEndpoint`, the reservation indexes already stripped, and the
5476    /// waiting client's sender dropped unanswered.
5477    #[test]
5478    fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
5479        let forwarding = ForwardingTable::default();
5480        let first_connection = ConnectionId::new(70);
5481        let (first_tx, _first_rx) = mpsc::channel(8);
5482        forwarding
5483            .register_module_connection(
5484                first_connection,
5485                MODULE_ID.to_string(),
5486                2,
5487                Concurrency::ModuleManaged,
5488                FrameSink::new(first_tx),
5489            )
5490            .unwrap();
5491        let (client_connection, client_sink, _client_rx) = client(200);
5492        let mut pending = forwarding
5493            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5494            .unwrap();
5495        let (second_tx, _second_rx) = mpsc::channel(8);
5496        forwarding
5497            .register_module_connection(
5498                ConnectionId::new(80),
5499                MODULE_ID.to_string(),
5500                2,
5501                Concurrency::ModuleManaged,
5502                FrameSink::new(second_tx),
5503            )
5504            .unwrap();
5505
5506        assert_eq!(
5507            forwarding
5508                .complete_pending_relay(
5509                    first_connection,
5510                    pending.corr,
5511                    RouteBindRelayOutcome::Accepted
5512                )
5513                .unwrap_err(),
5514            ForwardingError::StaleModuleEndpoint
5515        );
5516        assert!(committed_endpoints(&forwarding).is_empty());
5517        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
5518        assert!(matches!(
5519            pending.receiver.try_recv(),
5520            Err(oneshot::error::TryRecvError::Closed)
5521        ));
5522    }
5523
5524    #[test]
5525    fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
5526        // A candidate whose connection drops leaves the incumbent routable.
5527        let fixture = swap_fixture();
5528        let forwarding = &fixture.forwarding;
5529        assert!(forwarding
5530            .cleanup_connection(fixture.candidate_connection)
5531            .unwrap()
5532            .is_empty());
5533        assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
5534        let (client_connection, client_sink, _client_rx) = client(200);
5535        assert_eq!(
5536            forwarding
5537                .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5538                .unwrap()
5539                .endpoint,
5540            fixture.incumbent
5541        );
5542
5543        // After a cutover, the incumbent's teardown releases its own routes and
5544        // leaves the promoted candidate in place.
5545        let fixture = swap_fixture();
5546        let forwarding = &fixture.forwarding;
5547        let (bound_client, bound_sink, _bound_rx) = client(200);
5548        let bound = forwarding
5549            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
5550            .unwrap();
5551        forwarding
5552            .complete_pending_relay(
5553                fixture.incumbent_connection,
5554                bound.corr,
5555                RouteBindRelayOutcome::Accepted,
5556            )
5557            .unwrap();
5558        forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
5559        let released = forwarding
5560            .cleanup_connection(fixture.incumbent_connection)
5561            .unwrap();
5562        assert_eq!(released.len(), 1);
5563        assert_eq!(released[0].connection_id, bound_client);
5564        assert!(forwarding
5565            .read_inner()
5566            .unwrap()
5567            .superseded_endpoints
5568            .is_empty());
5569        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
5570        let (next_client, next_sink, _next_rx) = client(201);
5571        assert_eq!(
5572            forwarding
5573                .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
5574                .unwrap()
5575                .endpoint,
5576            fixture.candidate
5577        );
5578    }
5579}