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::{mpsc, 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    Frame, ProjectRootId,
25};
26
27/// Default per-channel request-credit window for modules that schedule internally.
28const DEFAULT_MODULE_MANAGED_WINDOW: usize = 32;
29
30/// High per-channel cap for stateless modules; this is an OOM guard, not scheduling policy.
31const STATELESS_PARALLEL_WINDOW: usize = 1024;
32
33/// A stopped probe cycle cannot retain its last unanswered correlation forever.
34/// Active endpoints replace the tombstone on their next serial health probe;
35/// this backstop covers endpoints that stop probing altogether.
36const HEALTH_PROBE_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
37
38/// Module connection identity used in forwarding keys.
39///
40/// The generation is bumped every time a module connection is registered so a future restart cannot
41/// accidentally reuse a stale `(connection_id, route_channel)` binding.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub struct ModuleEndpointId {
44    pub connection_id: ConnectionId,
45    pub generation: u64,
46}
47
48/// Client-local route key. A route channel is unique only within one client connection.
49#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
50pub(crate) struct ClientRouteKey {
51    pub connection_id: ConnectionId,
52    pub channel: u16,
53}
54
55/// Module-local route key. A route channel is unique only within one live module endpoint.
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
57pub(crate) struct ModuleRouteKey {
58    pub endpoint: ModuleEndpointId,
59    pub channel: u16,
60}
61
62#[derive(Debug)]
63pub(crate) struct RouteBinding {
64    pub client_connection_id: ConnectionId,
65    pub client_sink: FrameSink,
66    pub client_negotiated_ver: u8,
67    pub client_channel: u16,
68    pub client_epoch: u32,
69    pub module_id: String,
70    pub module_endpoint: ModuleEndpointId,
71    pub module_sink: FrameSink,
72    pub module_negotiated_ver: u8,
73    pub module_channel: u16,
74    pub module_epoch: u32,
75    pub principal: Principal,
76    pub project_root: Option<ProjectRootId>,
77    pub bound_at: Instant,
78    pub flow: Arc<ChannelFlow>,
79}
80
81#[derive(Debug, Clone)]
82pub(crate) enum DataRoute {
83    Client(DataRouteState),
84    Module(DataRouteState),
85}
86
87#[derive(Debug, Clone)]
88pub(crate) enum DataRouteState {
89    Bound(Arc<RouteBinding>),
90    Reserved,
91    EpochMismatch,
92    Absent,
93}
94
95/// Which kind of peer a route GOODBYE is being delivered to. This decides what
96/// happens when the GOODBYE cannot be enqueued (egress full/closed):
97/// - `Client`: escalate to closing that client connection (a socket close is a
98///   stronger teardown signal, and a full client egress means it is the slow
99///   client we would drop anyway).
100/// - `Module`: best-effort DROP, never close. A client-disconnect notifies the
101///   SHARED module that one client's route is gone; closing the module on its
102///   egress backpressure would tear down every co-tenant client (the exact
103///   cross-tenant blast radius this never-close rule exists to prevent — observed when a
104///   flooding dead client filled BOTH its own and the module's egress, so its
105///   route-gone GOODBYE to the module failed and closed the shared connection).
106///   subc has already removed the route from its forwarding state and drops
107///   stale module frames for the released channel (see router.rs), so subc's
108///   own routing is correct. The residual: under SUSTAINED module-egress
109///   backpressure a module-targeted route-gone notification can be lost, which a
110///   module using it for client-refcounting (e.g. AFT's session accounting)
111///   would miss. This is INTENTIONALLY ACCEPTED, not a gap. A consuming module
112///   must bound stale bindings with its own idle-activity reaper (last-touched
113///   TTL) independent of route-gone signals — AFT does exactly this, so a lost
114///   GOODBYE degrades to "the
115///   binding stays warm until its idle TTL" (bounded wasted resources), never an
116///   unbounded leak; disk-durable replay is unaffected. A dedicated reliable
117///   module control lane was evaluated and deliberately NOT
118///   built: it would add starvation-avoidance machinery to the thin core for a
119///   bounded warm-resource window that is not a correctness issue. never-close
120///   is the invariant that matters here.
121#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122pub(crate) enum GoodbyeTargetKind {
123    Client,
124    Module,
125}
126
127#[derive(Debug, Clone)]
128pub(crate) struct GoodbyeTarget {
129    pub connection_id: ConnectionId,
130    pub sink: FrameSink,
131    pub negotiated_ver: u8,
132    pub channel: u16,
133    pub epoch: u32,
134    pub kind: GoodbyeTargetKind,
135    /// The receiving module when this is a module-targeted relay. Client targets
136    /// have no module owner, while this identifier attributes a dropped relay.
137    pub module_id: Option<String>,
138}
139
140impl GoodbyeTarget {
141    /// True only when an undeliverable GOODBYE should escalate to closing the
142    /// target connection. Never escalate for module recipients.
143    pub(crate) fn close_on_delivery_failure(&self) -> bool {
144        matches!(self.kind, GoodbyeTargetKind::Client)
145    }
146}
147
148/// One route currently served by a module endpoint.
149///
150/// `goodbye_target` is deliberately retained alongside the census projection so
151/// a draining caller can address exactly the same route set that this read
152/// reports, without a second forwarding-table pass.
153#[derive(Debug, Clone)]
154pub(crate) struct EndpointRoute {
155    pub goodbye_target: GoodbyeTarget,
156    pub principal: Principal,
157    pub bound_at: Instant,
158    pub draining: bool,
159    /// WHY the endpoint is draining, when it is. Carried per-route so the
160    /// census can answer "closing because of what" without a second lookup;
161    /// `None` exactly when `draining` is false (one source: the drain map).
162    pub drain_reason: Option<RouteCloseReason>,
163}
164
165#[derive(Debug)]
166pub(crate) struct PendingRouteBindRelay {
167    pub endpoint: ModuleEndpointId,
168    pub module_sink: FrameSink,
169    pub negotiated_ver: u8,
170    pub client_channel: u16,
171    pub client_epoch: u32,
172    pub module_channel: u16,
173    pub module_epoch: u32,
174    pub corr: u64,
175    pub receiver: oneshot::Receiver<RouteBindRelayOutcome>,
176}
177
178#[derive(Debug, Clone)]
179pub(crate) struct ModuleDrainTarget {
180    pub endpoint: ModuleEndpointId,
181    pub sink: FrameSink,
182    pub negotiated_ver: u8,
183    pub abandoned_bindings: Vec<GoodbyeTarget>,
184    pub excluded_subscriptions: u32,
185}
186
187#[derive(Debug, Clone)]
188pub(crate) enum RouteBindRelayOutcome {
189    Accepted,
190    Rejected(ErrorBody),
191    ModuleGone(String),
192}
193
194/// What [`ForwardingTable::cutover_candidate`] did.
195#[derive(Debug, Clone, Copy, PartialEq, Eq)]
196pub(crate) struct ForwardingCutover {
197    /// The endpoint now in the active slot (the former candidate).
198    pub promoted: ModuleEndpointId,
199    /// The endpoint demoted out of the active slot, to be drained by endpoint.
200    /// `None` when the incumbent's connection was already gone.
201    pub incumbent: Option<ModuleEndpointId>,
202}
203
204#[derive(Debug, Clone)]
205pub(crate) struct PendingRelayCompletion {
206    pub settled: bool,
207    pub abandoned: Option<GoodbyeTarget>,
208}
209
210#[derive(Debug)]
211pub(crate) struct PendingModuleControlRpc {
212    pub endpoint: ModuleEndpointId,
213    pub module_sink: FrameSink,
214    pub negotiated_ver: u8,
215    pub corr: u64,
216    pub receiver: oneshot::Receiver<ModuleControlRpcOutcome>,
217}
218
219#[derive(Debug, Clone)]
220pub(crate) enum ModuleControlRpcOutcome {
221    Response(ModuleControlResponse),
222    Rejected(ErrorBody),
223    ModuleGone(String),
224    MalformedResponse(String),
225    UnexpectedOp { expected: String, actual: String },
226    DeadlineElapsed,
227}
228
229#[derive(Debug, Clone, PartialEq, Eq)]
230pub(crate) enum ModuleControlRpcCompletion {
231    Unknown,
232    Settled,
233    LateHealthAnswer {
234        module_id: String,
235        latency: Duration,
236    },
237}
238
239#[derive(Debug)]
240struct PendingModuleControlRpcEntry {
241    expected_op: String,
242    deadline: Instant,
243    health_probe_started_at: Option<Instant>,
244    sender: oneshot::Sender<ModuleControlRpcOutcome>,
245}
246
247#[derive(Debug)]
248struct HealthProbeTombstone {
249    expected_op: String,
250    module_id: String,
251    probe_started_at: Instant,
252    expires_at: Instant,
253}
254
255#[derive(Debug, Clone)]
256struct RouteReservation {
257    client_key: ClientRouteKey,
258    module_key: ModuleRouteKey,
259    client_epoch: u32,
260    module_epoch: u32,
261    project_root: Option<ProjectRootId>,
262}
263
264#[derive(Debug)]
265struct PendingRouteBindRelayEntry {
266    reservation: RouteReservation,
267    client_sink: FrameSink,
268    client_negotiated_ver: u8,
269    client_permit: mpsc::OwnedPermit<crate::router::OutboundFrame>,
270    route_open_frame: Frame,
271    principal: Principal,
272    deadline: Instant,
273    relay_enqueued: bool,
274    sender: oneshot::Sender<RouteBindRelayOutcome>,
275}
276
277#[derive(Debug, Clone)]
278pub(crate) enum RouteRelease {
279    Removed(GoodbyeTarget),
280    Stale,
281    Absent,
282}
283
284#[derive(Debug, Clone)]
285pub(crate) enum RoutePollSnapshot {
286    Bound {
287        module_id: String,
288        status: Option<String>,
289    },
290    Absent,
291}
292
293#[derive(Debug, Clone)]
294struct ModuleConnection {
295    endpoint: ModuleEndpointId,
296    sink: FrameSink,
297    negotiated_ver: u8,
298    concurrency: Concurrency,
299}
300
301#[derive(Debug, Default)]
302struct ForwardingInner {
303    daemon_draining: bool,
304    /// The ACTIVE slot: the one endpoint per module id that routing resolves.
305    /// Every by-id lookup (relay reservation, drain-by-id, liveness, census,
306    /// live roots, module-control RPCs) reads this map and nothing else.
307    modules_by_id: HashMap<String, ModuleConnection>,
308    /// The CANDIDATE slot of a blue/green swap: a second process registered
309    /// under an id that already has an active endpoint. It has a full endpoint
310    /// identity (so its own connection can be looked up and torn down) but no
311    /// by-id lookup sees it, so nothing is routed to it until `cutover_candidate`
312    /// promotes it.
313    candidates_by_id: HashMap<String, ModuleConnection>,
314    /// Former active endpoints demoted by `cutover_candidate`, until their
315    /// connection is removed. Membership is what distinguishes an endpoint that
316    /// was SUPERSEDED by a promotion (still alive, still carrying bound routes,
317    /// about to be drained) from one that is merely STALE (replaced some other
318    /// way, which is treated as a fault on the acking connection).
319    superseded_endpoints: HashMap<ModuleEndpointId, ModuleConnection>,
320    endpoint_by_connection: HashMap<ConnectionId, ModuleEndpointId>,
321    module_id_by_endpoint: HashMap<ModuleEndpointId, String>,
322    /// Endpoints mid-drain, keyed to the reason the drain was begun with. The
323    /// value serves the census ("draining because restart"); membership alone
324    /// still answers every admission-gate check.
325    draining_endpoints: HashMap<ModuleEndpointId, RouteCloseReason>,
326    closing_connections: HashSet<ConnectionId>,
327    next_generation: u64,
328    reserved_client: HashMap<ClientRouteKey, ModuleRouteKey>,
329    reserved_module: HashMap<ModuleRouteKey, ClientRouteKey>,
330    next_client_channel: HashMap<ConnectionId, u16>,
331    next_module_channel: HashMap<ModuleEndpointId, u16>,
332    client_slot_epochs: HashMap<ClientRouteKey, u32>,
333    module_slot_epochs: HashMap<ModuleRouteKey, u32>,
334    last_published_epoch: HashMap<ClientRouteKey, u32>,
335    client_to_module: HashMap<ClientRouteKey, Arc<RouteBinding>>,
336    module_to_client: HashMap<ModuleRouteKey, Arc<RouteBinding>>,
337    status: HashMap<(ClientRouteKey, u32), String>,
338    pending_relays: HashMap<(ModuleEndpointId, u64), PendingRouteBindRelayEntry>,
339    next_control_corr: HashMap<ModuleEndpointId, u64>,
340    pending_control_rpcs: HashMap<(ModuleEndpointId, u64), PendingModuleControlRpcEntry>,
341    health_probe_tombstones: HashMap<(ModuleEndpointId, u64), HealthProbeTombstone>,
342}
343
344#[derive(Debug, Clone)]
345pub(crate) struct CloseReason {
346    code: &'static str,
347    message: String,
348}
349
350impl CloseReason {
351    pub(crate) fn new(code: &'static str, message: impl Into<String>) -> Self {
352        Self {
353            code,
354            message: message.into(),
355        }
356    }
357}
358
359impl fmt::Display for CloseReason {
360    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
361        write!(f, "{}: {}", self.code, self.message)
362    }
363}
364
365pub(crate) type ConnectionCloseReceiver = oneshot::Receiver<CloseReason>;
366
367/// Dynamic forwarding state shared by the control plane and data-plane router.
368#[derive(Debug, Default)]
369pub struct ForwardingTable {
370    inner: Arc<RwLock<ForwardingInner>>,
371    close_registry: Mutex<HashMap<ConnectionId, oneshot::Sender<CloseReason>>>,
372    counters: DaemonCounters,
373    /// Per-target-module bind-relay breaker state. It lives here, beside the
374    /// module connections it describes, because this is where a module
375    /// connection's identity is established and therefore where a stale
376    /// verdict has to be discarded.
377    route_bind_breakers: RouteBindBreakers,
378    /// Current route.bind relays keyed by target module. Admission is shared
379    /// across every client connection that points at the same endpoint.
380    route_bind_concurrency: RouteBindConcurrency,
381}
382
383impl ForwardingTable {
384    pub(crate) fn counters(&self) -> DaemonCounters {
385        self.counters.clone()
386    }
387
388    pub(crate) fn route_bind_breakers(&self) -> RouteBindBreakers {
389        self.route_bind_breakers.clone()
390    }
391
392    pub(crate) fn route_bind_concurrency(&self) -> RouteBindConcurrency {
393        self.route_bind_concurrency.clone()
394    }
395
396    pub(crate) fn register_connection_close(
397        &self,
398        connection_id: ConnectionId,
399    ) -> ConnectionCloseReceiver {
400        let (sender, receiver) = oneshot::channel();
401        let replaced = self
402            .lock_close_registry()
403            .insert(connection_id, sender)
404            .is_some();
405        if replaced {
406            warn!(
407                connection_id = connection_id.get(),
408                "replaced existing connection close registration"
409            );
410        }
411        receiver
412    }
413
414    pub(crate) fn unregister_connection_close(&self, connection_id: ConnectionId) {
415        self.lock_close_registry().remove(&connection_id);
416    }
417
418    pub(crate) fn request_connection_close(
419        &self,
420        connection_id: ConnectionId,
421        reason: CloseReason,
422    ) {
423        let sender = self.lock_close_registry().remove(&connection_id);
424        if let Some(sender) = sender {
425            debug!(
426                connection_id = connection_id.get(),
427                close_reason = %reason,
428                "requesting connection close"
429            );
430            let _ = sender.send(reason);
431        } else {
432            debug!(
433                connection_id = connection_id.get(),
434                close_reason = %reason,
435                "connection close request ignored for inactive connection"
436            );
437        }
438    }
439
440    pub fn register_module_connection(
441        &self,
442        connection_id: ConnectionId,
443        module_id: String,
444        negotiated_ver: u8,
445        concurrency: Concurrency,
446        sink: FrameSink,
447    ) -> Result<ModuleEndpointId, ForwardingError> {
448        let mut inner = self.write_inner()?;
449        if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
450            return Err(ForwardingError::ConnectionClosing { connection_id });
451        }
452        if let Some(old_endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
453            let _ = remove_module_connection_locked(&mut inner, old_endpoint);
454        }
455
456        inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
457        let endpoint = ModuleEndpointId {
458            connection_id,
459            generation: inner.next_generation,
460        };
461        inner.endpoint_by_connection.insert(connection_id, endpoint);
462        inner
463            .module_id_by_endpoint
464            .insert(endpoint, module_id.clone());
465        inner.next_module_channel.insert(endpoint, 1);
466        inner.next_control_corr.insert(endpoint, 1);
467        inner.modules_by_id.insert(
468            module_id.clone(),
469            ModuleConnection {
470                endpoint,
471                sink,
472                negotiated_ver,
473                concurrency,
474            },
475        );
476        drop(inner);
477
478        // A new module connection has arrived under this id, so anything the
479        // bind-relay breaker learned was learned about a process that is no
480        // longer the one behind this name. See
481        // `RouteBindBreakers::reset_for_new_module_connection`.
482        //
483        // Keyed on ARRIVAL rather than on teardown deliberately: a connection
484        // going away is not evidence about anything, and a module whose
485        // connection drops without coming back should keep its verdict until
486        // something actually registers in its place. This also covers the
487        // unclean replacements -- a module killed mid-bind, or one whose
488        // connection was closing -- because registration is the single path by
489        // which any module connection becomes usable.
490        if let Some(discarded) = self
491            .route_bind_breakers
492            .reset_for_new_module_connection(&module_id)
493        {
494            info!(
495                module_id = %module_id,
496                discarded_consecutive_timeouts = discarded,
497                "route.bind breaker state discarded: a new module connection replaced the process it described"
498            );
499        }
500        Ok(endpoint)
501    }
502
503    /// Register a blue/green swap candidate for `module_id` into the candidate
504    /// slot, alongside whatever endpoint is active for the id.
505    ///
506    /// Unlike [`Self::register_module_connection`], this never touches the
507    /// active slot and never resets the bind-relay breaker: the incumbent is
508    /// still the process serving the id, so what the breaker learned about it
509    /// is still true. The breaker is reset at cutover instead, when the process
510    /// behind the name actually changes. A second candidate for the same id is
511    /// refused.
512    pub(crate) fn register_candidate_module_connection(
513        &self,
514        connection_id: ConnectionId,
515        module_id: String,
516        negotiated_ver: u8,
517        concurrency: Concurrency,
518        sink: FrameSink,
519    ) -> Result<ModuleEndpointId, ForwardingError> {
520        let mut inner = self.write_inner()?;
521        if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
522            return Err(ForwardingError::ConnectionClosing { connection_id });
523        }
524        if inner.candidates_by_id.contains_key(&module_id) {
525            return Err(ForwardingError::CandidateSlotOccupied { module_id });
526        }
527        if let Some(old_endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
528            let _ = remove_module_connection_locked(&mut inner, old_endpoint);
529        }
530
531        inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
532        let endpoint = ModuleEndpointId {
533            connection_id,
534            generation: inner.next_generation,
535        };
536        inner.endpoint_by_connection.insert(connection_id, endpoint);
537        inner
538            .module_id_by_endpoint
539            .insert(endpoint, module_id.clone());
540        inner.next_module_channel.insert(endpoint, 1);
541        inner.next_control_corr.insert(endpoint, 1);
542        inner.candidates_by_id.insert(
543            module_id,
544            ModuleConnection {
545                endpoint,
546                sink,
547                negotiated_ver,
548                concurrency,
549            },
550        );
551        Ok(endpoint)
552    }
553
554    /// Promote `module_id`'s candidate to the active slot, in ONE forwarding
555    /// write-lock critical section.
556    ///
557    /// That single section is the linearization point of a swap. Every relay
558    /// reservation resolves the active slot under the same lock, so each one
559    /// lands wholly before cutover (on the incumbent) or wholly after it (on the
560    /// promoted candidate), never on a mix. A relay reserved on the incumbent
561    /// before cutover can no longer commit: `commit_route_locked` requires the
562    /// reservation's endpoint to be the active one, and an ack for it arriving
563    /// later is answered by the superseded-endpoint arm of
564    /// `complete_pending_relay` instead of ending the incumbent's connection.
565    ///
566    /// Both endpoints keep their identities: nothing is re-keyed, so the
567    /// incumbent's bound routes and its pending correlation keys stay exactly
568    /// where they are until it is drained with [`Self::begin_endpoint_drain`],
569    /// using the incumbent endpoint this returns. Returns `Ok(None)` when there
570    /// is no candidate for the id.
571    pub(crate) fn cutover_candidate(
572        &self,
573        module_id: &str,
574    ) -> Result<Option<ForwardingCutover>, ForwardingError> {
575        let mut inner = self.write_inner()?;
576        if inner.daemon_draining {
577            return Err(ForwardingError::ModuleReloading {
578                module_id: module_id.to_string(),
579            });
580        }
581        let Some(candidate) = inner.candidates_by_id.remove(module_id) else {
582            return Ok(None);
583        };
584        let promoted = candidate.endpoint;
585        let incumbent = inner.modules_by_id.insert(module_id.to_string(), candidate);
586        let incumbent = incumbent.map(|incumbent| {
587            let endpoint = incumbent.endpoint;
588            inner.superseded_endpoints.insert(endpoint, incumbent);
589            endpoint
590        });
591        drop(inner);
592
593        // A different process now answers for this id; see the matching reset
594        // in `register_module_connection` for why a verdict about the old one
595        // must not carry over.
596        if let Some(discarded) = self
597            .route_bind_breakers
598            .reset_for_new_module_connection(module_id)
599        {
600            info!(
601                module_id = %module_id,
602                discarded_consecutive_timeouts = discarded,
603                "route.bind breaker state discarded: a swap candidate was promoted over the process it described"
604            );
605        }
606        Ok(Some(ForwardingCutover {
607            promoted,
608            incumbent,
609        }))
610    }
611
612    #[allow(clippy::too_many_arguments)]
613    pub(crate) async fn begin_route_bind_relay_for(
614        &self,
615        client_connection_id: ConnectionId,
616        client_sink: FrameSink,
617        client_negotiated_ver: u8,
618        client_corr: u64,
619        module_id: &str,
620        principal: Principal,
621        project_root: Option<ProjectRootId>,
622        deadline: Instant,
623    ) -> Result<PendingRouteBindRelay, ForwardingError> {
624        // Reserve egress capacity before taking the forwarding lock. The permit is
625        // held until the bind reaches one terminal state, so an accepted bind can
626        // publish its table entry and RouteOpen response in one critical section.
627        let client_permit =
628            client_sink
629                .reserve_owned()
630                .await
631                .map_err(|_| ForwardingError::ClientEgressClosed {
632                    connection_id: client_connection_id,
633                })?;
634        self.begin_route_bind_relay_inner(
635            client_connection_id,
636            client_sink,
637            client_negotiated_ver,
638            client_corr,
639            module_id,
640            principal,
641            project_root,
642            deadline,
643            client_permit,
644        )
645    }
646
647    #[cfg(test)]
648    pub(crate) fn begin_route_bind_relay_for_test(
649        &self,
650        client_connection_id: ConnectionId,
651        client_sink: FrameSink,
652        client_corr: u64,
653        module_id: &str,
654    ) -> Result<PendingRouteBindRelay, ForwardingError> {
655        let permit =
656            client_sink
657                .try_reserve_owned()
658                .map_err(|_| ForwardingError::ClientEgressClosed {
659                    connection_id: client_connection_id,
660                })?;
661        self.begin_route_bind_relay_inner(
662            client_connection_id,
663            client_sink,
664            subc_protocol::PROTOCOL_VERSION,
665            client_corr,
666            module_id,
667            Principal::Direct,
668            None,
669            Instant::now() + std::time::Duration::from_secs(60),
670            permit,
671        )
672    }
673
674    pub(crate) fn begin_module_control_rpc_for(
675        &self,
676        module_id: &str,
677        expected_op: &str,
678        deadline: Instant,
679    ) -> Result<PendingModuleControlRpc, ForwardingError> {
680        self.begin_module_control_rpc_inner(module_id, expected_op, deadline, None, false)
681    }
682
683    pub(crate) fn begin_health_probe_rpc_for(
684        &self,
685        module_id: &str,
686        expected_op: &str,
687        probe_started_at: Instant,
688        deadline: Instant,
689    ) -> Result<PendingModuleControlRpc, ForwardingError> {
690        self.begin_module_control_rpc_inner(
691            module_id,
692            expected_op,
693            deadline,
694            Some(probe_started_at),
695            false,
696        )
697    }
698
699    pub(crate) fn begin_drain_health_probe_rpc_for(
700        &self,
701        module_id: &str,
702        expected_op: &str,
703        probe_started_at: Instant,
704        deadline: Instant,
705    ) -> Result<PendingModuleControlRpc, ForwardingError> {
706        self.begin_module_control_rpc_inner(
707            module_id,
708            expected_op,
709            deadline,
710            Some(probe_started_at),
711            true,
712        )
713    }
714
715    /// A health probe addressed to one endpoint, in whichever slot it is.
716    ///
717    /// Every other control RPC resolves the id's ACTIVE endpoint, which is how
718    /// a swap candidate (not active until cutover) and a superseded incumbent
719    /// (not active after it) are unreachable by them. A swap needs to probe
720    /// exactly those two: the candidate before promoting it, and the incumbent
721    /// for its busy gauges while it drains.
722    pub(crate) fn begin_endpoint_health_probe_rpc_for(
723        &self,
724        endpoint: ModuleEndpointId,
725        expected_op: &str,
726        probe_started_at: Instant,
727        deadline: Instant,
728    ) -> Result<PendingModuleControlRpc, ForwardingError> {
729        let inner = self.write_inner()?;
730        let module = module_connection_for_endpoint_locked(&inner, endpoint)
731            .cloned()
732            .ok_or(ForwardingError::NoModuleConnection)?;
733        let module_id = inner
734            .module_id_by_endpoint
735            .get(&endpoint)
736            .cloned()
737            .unwrap_or_default();
738        // A draining incumbent is exactly what this probes for busy gauges, so
739        // draining is allowed, as it is for the by-id drain probe.
740        self.begin_control_rpc_locked(
741            inner,
742            &module_id,
743            module,
744            expected_op,
745            deadline,
746            Some(probe_started_at),
747            true,
748        )
749    }
750
751    fn begin_module_control_rpc_inner(
752        &self,
753        module_id: &str,
754        expected_op: &str,
755        deadline: Instant,
756        health_probe_started_at: Option<Instant>,
757        allow_draining: bool,
758    ) -> Result<PendingModuleControlRpc, ForwardingError> {
759        let inner = self.write_inner()?;
760        let module = inner
761            .modules_by_id
762            .get(module_id)
763            .cloned()
764            .ok_or(ForwardingError::NoModuleConnection)?;
765        self.begin_control_rpc_locked(
766            inner,
767            module_id,
768            module,
769            expected_op,
770            deadline,
771            health_probe_started_at,
772            allow_draining,
773        )
774    }
775
776    #[allow(clippy::too_many_arguments)]
777    fn begin_control_rpc_locked(
778        &self,
779        mut inner: RwLockWriteGuard<'_, ForwardingInner>,
780        module_id: &str,
781        module: ModuleConnection,
782        expected_op: &str,
783        deadline: Instant,
784        health_probe_started_at: Option<Instant>,
785        allow_draining: bool,
786    ) -> Result<PendingModuleControlRpc, ForwardingError> {
787        if !allow_draining && inner.draining_endpoints.contains_key(&module.endpoint) {
788            return Err(ForwardingError::ModuleReloading {
789                module_id: module_id.to_string(),
790            });
791        }
792        if inner
793            .closing_connections
794            .contains(&module.endpoint.connection_id)
795        {
796            return Err(ForwardingError::ConnectionClosing {
797                connection_id: module.endpoint.connection_id,
798            });
799        }
800        if health_probe_started_at.is_some() {
801            // Recurring health probes are serial per endpoint. Once the next one
802            // starts, an older answer can no longer improve the current snapshot,
803            // so retaining more than the newest unanswered probe has no value.
804            inner
805                .health_probe_tombstones
806                .retain(|(endpoint, _), _| *endpoint != module.endpoint);
807        }
808        let corr = match inner.allocate_control_corr(module.endpoint) {
809            Ok(corr) => corr,
810            Err(err) => {
811                drop(inner);
812                self.request_connection_close(
813                    module.endpoint.connection_id,
814                    CloseReason::new(
815                        "control_correlation_exhausted",
816                        "daemon-originated channel-0 correlation space exhausted",
817                    ),
818                );
819                return Err(err);
820            }
821        };
822        let (sender, receiver) = oneshot::channel();
823        inner.pending_control_rpcs.insert(
824            (module.endpoint, corr),
825            PendingModuleControlRpcEntry {
826                expected_op: expected_op.to_string(),
827                deadline,
828                health_probe_started_at,
829                sender,
830            },
831        );
832
833        Ok(PendingModuleControlRpc {
834            endpoint: module.endpoint,
835            module_sink: module.sink,
836            negotiated_ver: module.negotiated_ver,
837            corr,
838            receiver,
839        })
840    }
841
842    #[allow(clippy::too_many_arguments)]
843    fn begin_route_bind_relay_inner(
844        &self,
845        client_connection_id: ConnectionId,
846        client_sink: FrameSink,
847        client_negotiated_ver: u8,
848        client_corr: u64,
849        expected_module_id: &str,
850        principal: Principal,
851        project_root: Option<ProjectRootId>,
852        deadline: Instant,
853        client_permit: mpsc::OwnedPermit<crate::router::OutboundFrame>,
854    ) -> Result<PendingRouteBindRelay, ForwardingError> {
855        let mut inner = self.write_inner()?;
856        if inner.closing_connections.contains(&client_connection_id) {
857            return Err(ForwardingError::ConnectionClosing {
858                connection_id: client_connection_id,
859            });
860        }
861        let module = inner
862            .modules_by_id
863            .get(expected_module_id)
864            .cloned()
865            .ok_or(ForwardingError::NoModuleConnection)?;
866        if inner.draining_endpoints.contains_key(&module.endpoint) {
867            return Err(ForwardingError::ModuleReloading {
868                module_id: expected_module_id.to_string(),
869            });
870        }
871        if inner
872            .closing_connections
873            .contains(&module.endpoint.connection_id)
874        {
875            return Err(ForwardingError::ConnectionClosing {
876                connection_id: module.endpoint.connection_id,
877            });
878        }
879
880        let corr = match inner.allocate_control_corr(module.endpoint) {
881            Ok(corr) => corr,
882            Err(err) => {
883                drop(inner);
884                self.request_connection_close(
885                    module.endpoint.connection_id,
886                    CloseReason::new(
887                        "control_correlation_exhausted",
888                        "daemon-originated channel-0 correlation space exhausted",
889                    ),
890                );
891                return Err(err);
892            }
893        };
894        let (client_channel, client_epoch, module_channel, module_epoch) =
895            inner.allocate_route_slots(client_connection_id, module.endpoint)?;
896        let client_key = ClientRouteKey {
897            connection_id: client_connection_id,
898            channel: client_channel,
899        };
900        let module_key = ModuleRouteKey {
901            endpoint: module.endpoint,
902            channel: module_channel,
903        };
904        let reservation = RouteReservation {
905            client_key,
906            module_key,
907            client_epoch,
908            module_epoch,
909            project_root,
910        };
911        let response_body = serde_json::to_vec(&ClientControlResponse::RouteOpen {
912            route_channel: client_channel,
913            route_epoch: client_epoch,
914        })
915        .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
916        let route_open_frame = Frame::build_with_version(
917            client_negotiated_ver,
918            FrameType::Response,
919            Flags::new(false, Priority::Passive, false),
920            0,
921            0,
922            client_corr,
923            response_body,
924        )
925        .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
926        let (sender, receiver) = oneshot::channel();
927        inner.reserved_client.insert(client_key, module_key);
928        inner.reserved_module.insert(module_key, client_key);
929        inner.pending_relays.insert(
930            (module.endpoint, corr),
931            PendingRouteBindRelayEntry {
932                reservation,
933                client_sink,
934                client_negotiated_ver,
935                client_permit,
936                route_open_frame,
937                principal,
938                deadline,
939                relay_enqueued: false,
940                sender,
941            },
942        );
943
944        Ok(PendingRouteBindRelay {
945            endpoint: module.endpoint,
946            module_sink: module.sink,
947            negotiated_ver: module.negotiated_ver,
948            client_channel,
949            client_epoch,
950            module_channel,
951            module_epoch,
952            corr,
953            receiver,
954        })
955    }
956
957    pub(crate) fn mark_route_bind_relay_enqueued(
958        &self,
959        endpoint: ModuleEndpointId,
960        corr: u64,
961    ) -> Result<bool, ForwardingError> {
962        let mut inner = self.write_inner()?;
963        let Some(pending) = inner.pending_relays.get_mut(&(endpoint, corr)) else {
964            return Ok(false);
965        };
966        pending.relay_enqueued = true;
967        Ok(true)
968    }
969
970    pub(crate) fn release_client_route(
971        &self,
972        client_connection_id: ConnectionId,
973        client_channel: u16,
974        expected_epoch: u32,
975    ) -> Result<RouteRelease, ForwardingError> {
976        let mut inner = self.write_inner()?;
977        let release = release_client_route_locked(
978            &mut inner,
979            ClientRouteKey {
980                connection_id: client_connection_id,
981                channel: client_channel,
982            },
983            expected_epoch,
984        );
985        self.record_route_release(&release);
986        Ok(release)
987    }
988
989    pub(crate) fn release_module_route(
990        &self,
991        module_connection_id: ConnectionId,
992        module_channel: u16,
993        expected_epoch: u32,
994    ) -> Result<RouteRelease, ForwardingError> {
995        let mut inner = self.write_inner()?;
996        let Some(endpoint) = inner
997            .endpoint_by_connection
998            .get(&module_connection_id)
999            .copied()
1000        else {
1001            return Ok(RouteRelease::Absent);
1002        };
1003        let release = release_module_route_locked(
1004            &mut inner,
1005            ModuleRouteKey {
1006                endpoint,
1007                channel: module_channel,
1008            },
1009            expected_epoch,
1010        );
1011        self.record_route_release(&release);
1012        Ok(release)
1013    }
1014
1015    pub(crate) fn abort_pending_relay(
1016        &self,
1017        endpoint: ModuleEndpointId,
1018        corr: u64,
1019        outcome: RouteBindRelayOutcome,
1020    ) -> Result<Option<GoodbyeTarget>, ForwardingError> {
1021        let mut inner = self.write_inner()?;
1022        let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1023            return Ok(None);
1024        };
1025        release_reserved_route_locked(
1026            &mut inner,
1027            pending.reservation.client_key,
1028            pending.reservation.module_key,
1029        );
1030        let target = pending
1031            .relay_enqueued
1032            .then(|| abandoned_route_target(&inner, &pending.reservation));
1033        let _ = pending.sender.send(outcome);
1034        Ok(target.flatten())
1035    }
1036
1037    pub(crate) fn cancel_module_control_rpc(
1038        &self,
1039        endpoint: ModuleEndpointId,
1040        corr: u64,
1041    ) -> Result<(), ForwardingError> {
1042        self.write_inner()?
1043            .pending_control_rpcs
1044            .remove(&(endpoint, corr));
1045        Ok(())
1046    }
1047
1048    pub(crate) fn tombstone_health_probe_rpc(
1049        &self,
1050        endpoint: ModuleEndpointId,
1051        corr: u64,
1052    ) -> Result<bool, ForwardingError> {
1053        let key = (endpoint, corr);
1054        let expires_at = Instant::now() + HEALTH_PROBE_TOMBSTONE_TTL;
1055        {
1056            let mut inner = self.write_inner()?;
1057            let Some(pending) = inner.pending_control_rpcs.remove(&key) else {
1058                return Ok(false);
1059            };
1060            let Some(probe_started_at) = pending.health_probe_started_at else {
1061                inner.pending_control_rpcs.insert(key, pending);
1062                return Ok(false);
1063            };
1064            let module_id = inner
1065                .module_id_by_endpoint
1066                .get(&endpoint)
1067                .cloned()
1068                .unwrap_or_else(|| "unknown".to_string());
1069            inner.health_probe_tombstones.insert(
1070                key,
1071                HealthProbeTombstone {
1072                    expected_op: pending.expected_op,
1073                    module_id,
1074                    probe_started_at,
1075                    expires_at,
1076                },
1077            );
1078        }
1079        self.schedule_health_probe_tombstone_expiration(key, expires_at);
1080        Ok(true)
1081    }
1082
1083    fn schedule_health_probe_tombstone_expiration(
1084        &self,
1085        key: (ModuleEndpointId, u64),
1086        expires_at: Instant,
1087    ) {
1088        let inner = Arc::downgrade(&self.inner);
1089        tokio::spawn(async move {
1090            tokio::time::sleep_until(expires_at).await;
1091            let Some(inner) = inner.upgrade() else {
1092                return;
1093            };
1094            let Ok(mut inner) = inner.write() else {
1095                return;
1096            };
1097            let expired = inner
1098                .health_probe_tombstones
1099                .get(&key)
1100                .is_some_and(|tombstone| tombstone.expires_at <= Instant::now());
1101            if expired {
1102                inner.health_probe_tombstones.remove(&key);
1103            }
1104        });
1105    }
1106
1107    pub(crate) fn complete_pending_relay(
1108        &self,
1109        connection_id: ConnectionId,
1110        corr: u64,
1111        outcome: RouteBindRelayOutcome,
1112    ) -> Result<PendingRelayCompletion, ForwardingError> {
1113        let mut inner = self.write_inner()?;
1114        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1115            return Ok(PendingRelayCompletion {
1116                settled: false,
1117                abandoned: None,
1118            });
1119        };
1120        let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1121            return Ok(PendingRelayCompletion {
1122                settled: false,
1123                abandoned: None,
1124            });
1125        };
1126
1127        if Instant::now() >= pending.deadline {
1128            release_reserved_route_locked(
1129                &mut inner,
1130                pending.reservation.client_key,
1131                pending.reservation.module_key,
1132            );
1133            let abandoned = matches!(outcome, RouteBindRelayOutcome::Accepted)
1134                .then(|| abandoned_route_target(&inner, &pending.reservation))
1135                .flatten();
1136            let _ = pending
1137                .sender
1138                .send(RouteBindRelayOutcome::Rejected(ErrorBody {
1139                    code: "module_timeout".to_string(),
1140                    message: "route.bind response arrived after its daemon deadline".to_string(),
1141                    detail: None,
1142                }));
1143            return Ok(PendingRelayCompletion {
1144                settled: true,
1145                abandoned,
1146            });
1147        }
1148
1149        match outcome {
1150            // Two shapes of "the client is not there to receive this route" that
1151            // must resolve identically: its egress is already closed, or it is
1152            // marked closing (its connection loop has been asked to end, but has
1153            // not drained yet, so the sink is still open).
1154            //
1155            // Only the first used to be caught here. The second fell through to
1156            // `commit_route_locked`, which refuses a closing client with
1157            // `ConnectionClosing` -- and this function is called from the MODULE
1158            // connection's frame handler, so that refusal ended the module's
1159            // connection instead of this one client's route. One dying client's
1160            // route.open then took down a connection carrying every other
1161            // client's routes to that module.
1162            //
1163            // The remedy for both is the same, which is why they share an arm:
1164            // give back the reserved handle pair, tell the waiting route.open the
1165            // route is gone, and report the module-side channel so the caller can
1166            // send a channel-scoped GOODBYE for the binding the module just
1167            // created. Nothing here touches the module connection.
1168            RouteBindRelayOutcome::Accepted
1169                if pending.client_sink.is_closed()
1170                    || inner
1171                        .closing_connections
1172                        .contains(&pending.reservation.client_key.connection_id) =>
1173            {
1174                let reason = if pending.client_sink.is_closed() {
1175                    "client egress closed before route publication"
1176                } else {
1177                    "client connection is closing before route publication"
1178                };
1179                release_reserved_route_locked(
1180                    &mut inner,
1181                    pending.reservation.client_key,
1182                    pending.reservation.module_key,
1183                );
1184                let abandoned = pending
1185                    .relay_enqueued
1186                    .then(|| abandoned_route_target(&inner, &pending.reservation))
1187                    .flatten();
1188                let _ = pending
1189                    .sender
1190                    .send(RouteBindRelayOutcome::ModuleGone(reason.to_string()));
1191                return Ok(PendingRelayCompletion {
1192                    settled: true,
1193                    abandoned,
1194                });
1195            }
1196            // The acking endpoint was the active one when this relay was
1197            // reserved, and a swap has since promoted a candidate over it. The
1198            // module did nothing wrong: it bound the route it was asked to bind,
1199            // and it is still carrying every other client's routes until it is
1200            // drained. So this is settled here, while the pending entry still
1201            // holds the client's sender and the reservation pair: release the
1202            // pair, tell the waiting route.open to retry (it will reserve on the
1203            // promoted endpoint), and hand back the module-side channel so the
1204            // caller sends one channel-scoped GOODBYE for the binding the module
1205            // just created. Nothing here touches the module connection.
1206            //
1207            // Only membership in `superseded_endpoints` takes this arm. An
1208            // endpoint that stopped being the active one for any other reason
1209            // is STALE, not superseded, and still falls through to
1210            // `commit_route_locked`, which refuses it with `StaleModuleEndpoint`
1211            // exactly as before swaps existed.
1212            RouteBindRelayOutcome::Accepted
1213                if inner.superseded_endpoints.contains_key(&endpoint) =>
1214            {
1215                release_reserved_route_locked(
1216                    &mut inner,
1217                    pending.reservation.client_key,
1218                    pending.reservation.module_key,
1219                );
1220                // Not gated on `relay_enqueued`: an ack proves the module
1221                // received the bind, whether or not the enqueue mark was set.
1222                let abandoned = abandoned_route_target(&inner, &pending.reservation);
1223                let module_id = inner
1224                    .module_id_by_endpoint
1225                    .get(&endpoint)
1226                    .cloned()
1227                    .unwrap_or_else(|| "unknown".to_string());
1228                let _ = pending
1229                    .sender
1230                    .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
1231                        "module_reloading",
1232                        format!("module_id '{module_id}' is reloading"),
1233                    )));
1234                return Ok(PendingRelayCompletion {
1235                    settled: true,
1236                    abandoned,
1237                });
1238            }
1239            RouteBindRelayOutcome::Accepted => {
1240                let abandoned = commit_route_locked(&mut inner, pending)?;
1241                return Ok(PendingRelayCompletion {
1242                    settled: true,
1243                    abandoned,
1244                });
1245            }
1246            terminal => {
1247                release_reserved_route_locked(
1248                    &mut inner,
1249                    pending.reservation.client_key,
1250                    pending.reservation.module_key,
1251                );
1252                let _ = pending.sender.send(terminal);
1253            }
1254        }
1255        Ok(PendingRelayCompletion {
1256            settled: true,
1257            abandoned: None,
1258        })
1259    }
1260
1261    pub(crate) fn pending_module_control_op(
1262        &self,
1263        connection_id: ConnectionId,
1264        corr: u64,
1265    ) -> Result<Option<String>, ForwardingError> {
1266        let inner = self.read_inner()?;
1267        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1268            return Ok(None);
1269        };
1270        let key = (endpoint, corr);
1271        Ok(inner
1272            .pending_control_rpcs
1273            .get(&key)
1274            .map(|pending| pending.expected_op.clone())
1275            .or_else(|| {
1276                inner
1277                    .health_probe_tombstones
1278                    .get(&key)
1279                    .filter(|tombstone| tombstone.expires_at > Instant::now())
1280                    .map(|tombstone| tombstone.expected_op.clone())
1281            }))
1282    }
1283
1284    pub(crate) fn complete_module_control_rpc(
1285        &self,
1286        connection_id: ConnectionId,
1287        corr: u64,
1288        actual_op: Option<&str>,
1289        outcome: ModuleControlRpcOutcome,
1290    ) -> Result<ModuleControlRpcCompletion, ForwardingError> {
1291        let now = Instant::now();
1292        let mut inner = self.write_inner()?;
1293        let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1294            return Ok(ModuleControlRpcCompletion::Unknown);
1295        };
1296        let key = (endpoint, corr);
1297        if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
1298            if now >= pending.deadline {
1299                let late_health_answer = pending.health_probe_started_at.map(|probe_started_at| {
1300                    ModuleControlRpcCompletion::LateHealthAnswer {
1301                        module_id: inner
1302                            .module_id_by_endpoint
1303                            .get(&endpoint)
1304                            .cloned()
1305                            .unwrap_or_else(|| "unknown".to_string()),
1306                        latency: now.saturating_duration_since(probe_started_at),
1307                    }
1308                });
1309                let _ = pending
1310                    .sender
1311                    .send(ModuleControlRpcOutcome::DeadlineElapsed);
1312                return Ok(late_health_answer.unwrap_or(ModuleControlRpcCompletion::Settled));
1313            }
1314            let outcome = match actual_op {
1315                Some(actual) if actual != pending.expected_op => {
1316                    ModuleControlRpcOutcome::UnexpectedOp {
1317                        expected: pending.expected_op,
1318                        actual: actual.to_string(),
1319                    }
1320                }
1321                _ => outcome,
1322            };
1323            let _ = pending.sender.send(outcome);
1324            return Ok(ModuleControlRpcCompletion::Settled);
1325        }
1326
1327        let Some(tombstone) = inner.health_probe_tombstones.remove(&key) else {
1328            return Ok(ModuleControlRpcCompletion::Unknown);
1329        };
1330        if tombstone.expires_at <= now {
1331            return Ok(ModuleControlRpcCompletion::Unknown);
1332        }
1333        Ok(ModuleControlRpcCompletion::LateHealthAnswer {
1334            module_id: tombstone.module_id,
1335            latency: now.saturating_duration_since(tombstone.probe_started_at),
1336        })
1337    }
1338
1339    #[cfg(test)]
1340    pub(crate) fn health_probe_tombstone_count(&self) -> Result<usize, ForwardingError> {
1341        Ok(self.read_inner()?.health_probe_tombstones.len())
1342    }
1343
1344    #[cfg(test)]
1345    pub(crate) fn closing_connection_count(&self) -> Result<usize, ForwardingError> {
1346        Ok(self.read_inner()?.closing_connections.len())
1347    }
1348
1349    /// Reserved-but-uncommitted route handles, counted on both index sides, so
1350    /// a test can assert a reservation pair was actually given back.
1351    #[cfg(test)]
1352    pub(crate) fn reserved_route_count(&self) -> Result<(usize, usize), ForwardingError> {
1353        let inner = self.read_inner()?;
1354        Ok((inner.reserved_client.len(), inner.reserved_module.len()))
1355    }
1356
1357    pub(crate) fn module_endpoint_for_connection(
1358        &self,
1359        connection_id: ConnectionId,
1360    ) -> Result<Option<ModuleEndpointId>, ForwardingError> {
1361        Ok(self
1362            .read_inner()?
1363            .endpoint_by_connection
1364            .get(&connection_id)
1365            .copied())
1366    }
1367
1368    /// Looks up the module registered on a data-plane connection so route-drop
1369    /// diagnostics name the emitter instead of only reporting a daemon total.
1370    pub(crate) fn module_id_for_connection(
1371        &self,
1372        connection_id: ConnectionId,
1373    ) -> Result<Option<String>, ForwardingError> {
1374        let inner = self.read_inner()?;
1375        Ok(inner
1376            .endpoint_by_connection
1377            .get(&connection_id)
1378            .and_then(|endpoint| inner.module_id_by_endpoint.get(endpoint))
1379            .cloned())
1380    }
1381
1382    pub(crate) fn has_live_module_connection(
1383        &self,
1384        module_id: &str,
1385    ) -> Result<bool, ForwardingError> {
1386        Ok(self.read_inner()?.modules_by_id.contains_key(module_id))
1387    }
1388
1389    pub(crate) fn lookup_data_route(
1390        &self,
1391        connection_id: ConnectionId,
1392        channel: u16,
1393        epoch: u32,
1394    ) -> Result<DataRoute, ForwardingError> {
1395        let inner = self.read_inner()?;
1396        let state = if let Some(endpoint) =
1397            inner.endpoint_by_connection.get(&connection_id).copied()
1398        {
1399            let key = ModuleRouteKey { endpoint, channel };
1400            match inner.module_to_client.get(&key) {
1401                Some(route) if route.module_epoch == epoch => {
1402                    DataRouteState::Bound(Arc::clone(route))
1403                }
1404                Some(_) => DataRouteState::EpochMismatch,
1405                None if inner.reserved_module.contains_key(&key)
1406                    && inner.module_slot_epochs.get(&key).copied() == Some(epoch) =>
1407                {
1408                    DataRouteState::Reserved
1409                }
1410                None if inner.reserved_module.contains_key(&key) => DataRouteState::EpochMismatch,
1411                None => DataRouteState::Absent,
1412            }
1413        } else {
1414            let key = ClientRouteKey {
1415                connection_id,
1416                channel,
1417            };
1418            match inner.client_to_module.get(&key) {
1419                Some(route) if route.client_epoch == epoch => {
1420                    DataRouteState::Bound(Arc::clone(route))
1421                }
1422                Some(_) => DataRouteState::EpochMismatch,
1423                None if inner.reserved_client.contains_key(&key)
1424                    && inner.client_slot_epochs.get(&key).copied() == Some(epoch) =>
1425                {
1426                    DataRouteState::Reserved
1427                }
1428                None if inner.reserved_client.contains_key(&key) => DataRouteState::EpochMismatch,
1429                None => DataRouteState::Absent,
1430            }
1431        };
1432        Ok(
1433            if inner.endpoint_by_connection.contains_key(&connection_id) {
1434                DataRoute::Module(state)
1435            } else {
1436                DataRoute::Client(state)
1437            },
1438        )
1439    }
1440
1441    #[cfg(test)]
1442    pub(crate) fn inject_client_slot_epoch(
1443        &self,
1444        connection_id: ConnectionId,
1445        channel: u16,
1446        last_epoch: u32,
1447    ) {
1448        let mut inner = self.write_inner().expect("forwarding lock");
1449        inner.client_slot_epochs.insert(
1450            ClientRouteKey {
1451                connection_id,
1452                channel,
1453            },
1454            last_epoch,
1455        );
1456        inner.next_client_channel.insert(connection_id, channel);
1457    }
1458
1459    #[cfg(test)]
1460    pub(crate) fn inject_module_slot_epoch(
1461        &self,
1462        endpoint: ModuleEndpointId,
1463        channel: u16,
1464        last_epoch: u32,
1465    ) {
1466        let mut inner = self.write_inner().expect("forwarding lock");
1467        inner
1468            .module_slot_epochs
1469            .insert(ModuleRouteKey { endpoint, channel }, last_epoch);
1470        inner.next_module_channel.insert(endpoint, channel);
1471    }
1472
1473    #[cfg(test)]
1474    pub(crate) fn inject_control_corr(&self, endpoint: ModuleEndpointId, next_corr: u64) {
1475        self.write_inner()
1476            .expect("forwarding lock")
1477            .next_control_corr
1478            .insert(endpoint, next_corr);
1479    }
1480
1481    pub(crate) fn cache_status(
1482        &self,
1483        endpoint: ModuleEndpointId,
1484        module_channel: u16,
1485        module_epoch: u32,
1486        status: String,
1487    ) -> Result<bool, ForwardingError> {
1488        let mut inner = self.write_inner()?;
1489        if !inner.module_id_by_endpoint.contains_key(&endpoint) {
1490            return Err(ForwardingError::StaleModuleEndpoint);
1491        }
1492
1493        let module_key = ModuleRouteKey {
1494            endpoint,
1495            channel: module_channel,
1496        };
1497        let handle = if let Some(route) = inner.module_to_client.get(&module_key) {
1498            (route.module_epoch == module_epoch).then_some((
1499                ClientRouteKey {
1500                    connection_id: route.client_connection_id,
1501                    channel: route.client_channel,
1502                },
1503                route.client_epoch,
1504            ))
1505        } else if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
1506            (inner.module_slot_epochs.get(&module_key).copied() == Some(module_epoch)).then_some((
1507                client_key,
1508                inner
1509                    .client_slot_epochs
1510                    .get(&client_key)
1511                    .copied()
1512                    .unwrap_or(0),
1513            ))
1514        } else {
1515            None
1516        };
1517
1518        if let Some(handle) = handle {
1519            inner.status.insert(handle, status);
1520            Ok(true)
1521        } else {
1522            debug!(
1523                module_channel,
1524                module_epoch,
1525                generation = endpoint.generation,
1526                connection_id = endpoint.connection_id.get(),
1527                "dropping stale status update for module route handle"
1528            );
1529            Ok(false)
1530        }
1531    }
1532
1533    pub(crate) fn route_poll_snapshot(
1534        &self,
1535        client_connection_id: ConnectionId,
1536        client_channel: u16,
1537        client_epoch: u32,
1538    ) -> Result<RoutePollSnapshot, ForwardingError> {
1539        let inner = self.read_inner()?;
1540        let client_key = ClientRouteKey {
1541            connection_id: client_connection_id,
1542            channel: client_channel,
1543        };
1544        let Some(route) = inner.client_to_module.get(&client_key) else {
1545            return Ok(RoutePollSnapshot::Absent);
1546        };
1547        if route.client_epoch != client_epoch
1548            || !inner
1549                .module_id_by_endpoint
1550                .contains_key(&route.module_endpoint)
1551        {
1552            return Ok(RoutePollSnapshot::Absent);
1553        }
1554        Ok(RoutePollSnapshot::Bound {
1555            module_id: route.module_id.clone(),
1556            status: inner.status.get(&(client_key, client_epoch)).cloned(),
1557        })
1558    }
1559
1560    pub fn active_binding_count(&self) -> Result<usize, ForwardingError> {
1561        Ok(self.read_inner()?.client_to_module.len())
1562    }
1563
1564    /// How many distinct client connections hold at least one committed route,
1565    /// alongside the largest number of routes any single connection holds.
1566    ///
1567    /// `connected_clients` alone cannot distinguish many clients with a route
1568    /// each from one client accumulating hundreds, and those have opposite
1569    /// causes. Reading it required an out-of-band `lsof` during a live
1570    /// investigation, and the count of connections was mistaken for a count of
1571    /// client processes — which sent two of us after cleanup paths that were
1572    /// working correctly.
1573    pub fn client_route_concentration(&self) -> Result<(usize, usize), ForwardingError> {
1574        let inner = self.read_inner()?;
1575        let mut per_connection: HashMap<ConnectionId, usize> = HashMap::new();
1576        for key in inner.client_to_module.keys() {
1577            *per_connection.entry(key.connection_id).or_insert(0) += 1;
1578        }
1579        let max = per_connection.values().copied().max().unwrap_or(0);
1580        Ok((per_connection.len(), max))
1581    }
1582
1583    pub fn has_route_channel(&self, route_channel: u16) -> Result<bool, ForwardingError> {
1584        let inner = self.read_inner()?;
1585        Ok(inner
1586            .client_to_module
1587            .keys()
1588            .any(|key| key.channel == route_channel))
1589    }
1590
1591    /// Gate every provider atomically, including registrations racing shutdown.
1592    /// No supervisor lock is held while taking the forwarding lock.
1593    #[cfg(unix)]
1594    pub(crate) fn begin_daemon_drain(&self) -> Result<Vec<String>, ForwardingError> {
1595        let mut inner = self.write_inner()?;
1596        inner.daemon_draining = true;
1597        let modules = inner
1598            .modules_by_id
1599            .iter()
1600            .map(|(id, module)| (id.clone(), module.endpoint))
1601            .collect::<Vec<_>>();
1602        for (_, endpoint) in &modules {
1603            inner
1604                .draining_endpoints
1605                .insert(*endpoint, RouteCloseReason::Restart);
1606        }
1607        // Swap candidates and superseded incumbents are not routable, but they
1608        // are live endpoints that could still be sent module-control work, so
1609        // the daemon-wide gate covers them too. The returned ids are unchanged:
1610        // they name modules for the supervisor to drain, one per id.
1611        let off_slot_endpoints = inner
1612            .candidates_by_id
1613            .values()
1614            .map(|module| module.endpoint)
1615            .chain(inner.superseded_endpoints.keys().copied())
1616            .collect::<Vec<_>>();
1617        for endpoint in off_slot_endpoints {
1618            inner
1619                .draining_endpoints
1620                .insert(endpoint, RouteCloseReason::Restart);
1621        }
1622        Ok(modules.into_iter().map(|(id, _)| id).collect())
1623    }
1624
1625    /// Begin draining whatever endpoint is ACTIVE for `module_id`.
1626    ///
1627    /// After a swap's cutover the active endpoint is the promoted candidate, so
1628    /// this must not be used to drain the incumbent; use
1629    /// [`Self::begin_endpoint_drain`] with the endpoint `cutover_candidate`
1630    /// returned.
1631    pub(crate) fn begin_module_drain(
1632        &self,
1633        module_id: &str,
1634        reason: RouteCloseReason,
1635    ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
1636        let mut inner = self.write_inner()?;
1637        let Some(module) = inner.modules_by_id.get(module_id).cloned() else {
1638            return Ok(None);
1639        };
1640        Ok(Some(begin_drain_locked(
1641            &mut inner, module_id, module, reason,
1642        )))
1643    }
1644
1645    /// Begin draining one specific endpoint, whichever slot it is in.
1646    ///
1647    /// This is how a swap drains its incumbent after cutover: the incumbent's
1648    /// endpoint is captured by `cutover_candidate`, and resolving it by module id
1649    /// instead would find the promoted candidate and leave neither process
1650    /// routable. Returns `Ok(None)` when the endpoint is no longer registered.
1651    pub(crate) fn begin_endpoint_drain(
1652        &self,
1653        endpoint: ModuleEndpointId,
1654        reason: RouteCloseReason,
1655    ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
1656        let mut inner = self.write_inner()?;
1657        let Some(module) = module_connection_for_endpoint_locked(&inner, endpoint).cloned() else {
1658            return Ok(None);
1659        };
1660        let module_id = inner
1661            .module_id_by_endpoint
1662            .get(&endpoint)
1663            .cloned()
1664            .expect("an endpoint resolved to a module connection has a module id");
1665        Ok(Some(begin_drain_locked(
1666            &mut inner, &module_id, module, reason,
1667        )))
1668    }
1669}
1670
1671/// Mark `module`'s endpoint draining and settle everything still pending on it.
1672/// Shared by the by-id and by-endpoint drain entry points, which differ only in
1673/// how they find the endpoint.
1674fn begin_drain_locked(
1675    inner: &mut ForwardingInner,
1676    module_id: &str,
1677    module: ModuleConnection,
1678    reason: RouteCloseReason,
1679) -> ModuleDrainTarget {
1680    {
1681        let endpoint = module.endpoint;
1682        inner.draining_endpoints.insert(endpoint, reason);
1683
1684        let flows = inner
1685            .client_to_module
1686            .values()
1687            .filter(|route| route.module_endpoint == endpoint)
1688            .map(|route| Arc::clone(&route.flow))
1689            .collect::<Vec<_>>();
1690        let excluded_subscriptions = flows
1691            .into_iter()
1692            .map(|flow| flow.begin_drain())
1693            .fold(0u32, u32::saturating_add);
1694
1695        let pending_keys = inner
1696            .pending_relays
1697            .keys()
1698            .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
1699            .copied()
1700            .collect::<Vec<_>>();
1701        let mut abandoned_bindings = Vec::new();
1702        for key in pending_keys {
1703            let Some(pending) = inner.pending_relays.remove(&key) else {
1704                continue;
1705            };
1706            release_reserved_route_locked(
1707                inner,
1708                pending.reservation.client_key,
1709                pending.reservation.module_key,
1710            );
1711            if pending.relay_enqueued {
1712                if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
1713                    abandoned_bindings.push(target);
1714                }
1715            }
1716            let _ = pending
1717                .sender
1718                .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
1719                    "module_reloading",
1720                    format!("module_id '{module_id}' is reloading"),
1721                )));
1722        }
1723
1724        let pending_control_keys = inner
1725            .pending_control_rpcs
1726            .keys()
1727            .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
1728            .copied()
1729            .collect::<Vec<_>>();
1730        for key in pending_control_keys {
1731            if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
1732                let _ = pending
1733                    .sender
1734                    .send(ModuleControlRpcOutcome::ModuleGone(format!(
1735                        "module '{module_id}' began draining during module-control RPC"
1736                    )));
1737            }
1738        }
1739
1740        ModuleDrainTarget {
1741            endpoint,
1742            sink: module.sink,
1743            negotiated_ver: module.negotiated_ver,
1744            abandoned_bindings,
1745            excluded_subscriptions,
1746        }
1747    }
1748}
1749
1750/// What a drain that timed out was still waiting on, for the log line that
1751/// reports the timeout. Counts only: per-request ages are not tracked.
1752#[derive(Debug, Default, PartialEq, Eq)]
1753pub(crate) struct DrainHoldouts {
1754    /// Requests still counted against the drain (subscriptions flagged as
1755    /// such are excluded from the drain and not counted here).
1756    pub(crate) requests: usize,
1757    /// Routes holding at least one of those requests.
1758    pub(crate) routes: usize,
1759    /// Every route on the endpoint, for scale.
1760    pub(crate) total_routes: usize,
1761    /// The client connections holding the most requests, largest first, at
1762    /// most three: enough to name the consumer without listing every route.
1763    pub(crate) top_connections: Vec<(u64, usize)>,
1764}
1765
1766impl ForwardingTable {
1767    /// Summarise the requests one endpoint's drain is still waiting on.
1768    pub(crate) fn endpoint_drain_holdouts(
1769        &self,
1770        endpoint: ModuleEndpointId,
1771    ) -> Result<DrainHoldouts, ForwardingError> {
1772        let inner = self.read_inner()?;
1773        let mut holdouts = DrainHoldouts::default();
1774        let mut by_connection: HashMap<u64, usize> = HashMap::new();
1775        for (key, route) in &inner.client_to_module {
1776            if route.module_endpoint != endpoint {
1777                continue;
1778            }
1779            holdouts.total_routes += 1;
1780            let held = route.flow.drain_in_flight();
1781            if held == 0 {
1782                continue;
1783            }
1784            holdouts.requests += held;
1785            holdouts.routes += 1;
1786            *by_connection.entry(key.connection_id.get()).or_default() += held;
1787        }
1788        let mut connections = by_connection.into_iter().collect::<Vec<_>>();
1789        connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
1790        connections.truncate(3);
1791        holdouts.top_connections = connections;
1792        Ok(holdouts)
1793    }
1794
1795    pub(crate) fn endpoint_in_flight_count(
1796        &self,
1797        endpoint: ModuleEndpointId,
1798    ) -> Result<usize, ForwardingError> {
1799        let inner = self.read_inner()?;
1800        Ok(inner
1801            .client_to_module
1802            .values()
1803            .filter(|route| route.module_endpoint == endpoint)
1804            .map(|route| route.flow.drain_in_flight())
1805            .sum())
1806    }
1807
1808    pub(crate) fn endpoint_is_draining(
1809        &self,
1810        endpoint: ModuleEndpointId,
1811    ) -> Result<bool, ForwardingError> {
1812        Ok(self
1813            .read_inner()?
1814            .draining_endpoints
1815            .contains_key(&endpoint))
1816    }
1817
1818    pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
1819        let inner = self.read_inner()?;
1820        Ok(inner
1821            .modules_by_id
1822            .get(module_id)
1823            .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
1824    }
1825
1826    pub(crate) fn release_module_endpoint_routes(
1827        &self,
1828        endpoint: ModuleEndpointId,
1829    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1830        let mut inner = self.write_inner()?;
1831        let routes = inner
1832            .module_to_client
1833            .iter()
1834            .filter(|(module_key, _)| module_key.endpoint == endpoint)
1835            .map(|(module_key, route)| (*module_key, route.module_epoch))
1836            .collect::<Vec<_>>();
1837        let mut released = Vec::with_capacity(routes.len());
1838        for (module_key, epoch) in routes {
1839            if let RouteRelease::Removed(target) =
1840                release_module_route_locked(&mut inner, module_key, epoch)
1841            {
1842                released.push(target);
1843            }
1844        }
1845        Ok(released)
1846    }
1847
1848    /// Enumerate one endpoint's current routes without contacting the module.
1849    ///
1850    /// The read lock makes this safe while the endpoint drains: the returned
1851    /// `draining` marker describes the same table state that owns the route,
1852    /// rather than inferring liveness from a module that may be stopping.
1853    pub(crate) fn endpoint_routes(
1854        &self,
1855        endpoint: ModuleEndpointId,
1856    ) -> Result<Vec<EndpointRoute>, ForwardingError> {
1857        let inner = self.read_inner()?;
1858        Ok(endpoint_routes_locked(&inner, endpoint))
1859    }
1860
1861    /// Snapshot all live endpoint route sets under one forwarding-table read lock.
1862    pub(crate) fn route_census(
1863        &self,
1864        module_id: Option<&str>,
1865    ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
1866        let inner = self.read_inner()?;
1867        let mut endpoints = inner
1868            .modules_by_id
1869            .iter()
1870            .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
1871            .map(|(id, module)| (id.clone(), module.endpoint))
1872            .collect::<Vec<_>>();
1873        endpoints.sort_by(|left, right| left.0.cmp(&right.0));
1874        Ok(endpoints
1875            .into_iter()
1876            .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
1877            .collect())
1878    }
1879
1880    /// Snapshot only the endpoint currently routable under this module id.
1881    pub(crate) fn live_roots(
1882        &self,
1883        module_id: &str,
1884    ) -> Result<ModuleControlResponseToModule, ForwardingError> {
1885        let inner = self.read_inner()?;
1886        let endpoint = inner
1887            .modules_by_id
1888            .get(module_id)
1889            .map(|module| module.endpoint);
1890        let mut roots = BTreeMap::new();
1891        let mut unknown_root_bindings = 0;
1892        let mut total_bindings = 0;
1893        if let Some(endpoint) = endpoint {
1894            for binding in inner
1895                .module_to_client
1896                .values()
1897                .filter(|binding| binding.module_endpoint == endpoint)
1898            {
1899                total_bindings += 1;
1900                if let Some(root) = &binding.project_root {
1901                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1902                    entry.0 += 1;
1903                } else {
1904                    unknown_root_bindings += 1;
1905                }
1906            }
1907            for pending in inner
1908                .pending_relays
1909                .values()
1910                .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
1911            {
1912                total_bindings += 1;
1913                if let Some(root) = &pending.reservation.project_root {
1914                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1915                    entry.1 += 1;
1916                } else {
1917                    unknown_root_bindings += 1;
1918                }
1919            }
1920        }
1921        Ok(ModuleControlResponseToModule::LiveRoots {
1922            roots: roots
1923                .into_iter()
1924                .map(|(project_root, (bound, pending))| LiveRoot {
1925                    project_root,
1926                    bound,
1927                    pending,
1928                })
1929                .collect(),
1930            unknown_root_bindings,
1931            total_bindings,
1932        })
1933    }
1934
1935    /// True if this connection already owns committed or reserved CLIENT routes.
1936    /// A module registers (HELLO) before serving and never opens client routes, so
1937    /// a connection that has client routes must not also become a module endpoint
1938    /// — otherwise one connection holds both client and module state and cleanup
1939    /// only releases one side.
1940    pub(crate) fn connection_has_client_routes(
1941        &self,
1942        connection_id: ConnectionId,
1943    ) -> Result<bool, ForwardingError> {
1944        let inner = self.read_inner()?;
1945        let has = inner
1946            .client_to_module
1947            .keys()
1948            .any(|key| key.connection_id == connection_id)
1949            || inner
1950                .reserved_client
1951                .keys()
1952                .any(|key| key.connection_id == connection_id);
1953        Ok(has)
1954    }
1955
1956    pub(crate) fn cleanup_connection(
1957        &self,
1958        connection_id: ConnectionId,
1959    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1960        let mut inner = self.write_inner()?;
1961        inner.closing_connections.insert(connection_id);
1962        let released = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
1963            remove_module_connection_locked(&mut inner, endpoint)
1964        } else {
1965            Self::cleanup_client_connection_locked(&mut inner, connection_id)
1966        };
1967        // The closing mark refuses new work for a connection whose teardown is
1968        // still pending. This is the latest point at which lifting it is safe:
1969        // teardown has just removed every per-connection entry above, under
1970        // this same write lock, so no lookup can still find live state for the
1971        // id; and ids come from a monotonic counter that never reissues one,
1972        // so the id can never name a future connection either. Keeping the
1973        // mark past this point only grew the set by one entry per connection
1974        // for the life of the daemon.
1975        inner.closing_connections.remove(&connection_id);
1976        Ok(released)
1977    }
1978
1979    fn cleanup_client_connection_locked(
1980        inner: &mut ForwardingInner,
1981        connection_id: ConnectionId,
1982    ) -> Vec<GoodbyeTarget> {
1983        let routes = inner
1984            .client_to_module
1985            .iter()
1986            .filter(|(key, _)| key.connection_id == connection_id)
1987            .map(|(key, route)| (*key, route.client_epoch))
1988            .collect::<Vec<_>>();
1989        let mut released = Vec::with_capacity(routes.len());
1990        for (client_key, epoch) in routes {
1991            if let RouteRelease::Removed(target) =
1992                release_client_route_locked(inner, client_key, epoch)
1993            {
1994                released.push(target);
1995            }
1996        }
1997
1998        let pending_keys = inner
1999            .pending_relays
2000            .iter()
2001            .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2002            .map(|(key, _)| *key)
2003            .collect::<Vec<_>>();
2004        for key in pending_keys {
2005            let Some(pending) = inner.pending_relays.remove(&key) else {
2006                continue;
2007            };
2008            release_reserved_route_locked(
2009                inner,
2010                pending.reservation.client_key,
2011                pending.reservation.module_key,
2012            );
2013            if pending.relay_enqueued {
2014                if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2015                    released.push(target);
2016                }
2017            }
2018            let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2019                "client connection closed during route.bind relay".to_string(),
2020            ));
2021        }
2022
2023        let orphaned = inner
2024            .reserved_client
2025            .iter()
2026            .filter(|(key, _)| key.connection_id == connection_id)
2027            .map(|(client, module)| (*client, *module))
2028            .collect::<Vec<_>>();
2029        for (client_key, module_key) in orphaned {
2030            release_reserved_route_locked(inner, client_key, module_key);
2031        }
2032        inner.next_client_channel.remove(&connection_id);
2033        inner
2034            .client_slot_epochs
2035            .retain(|key, _| key.connection_id != connection_id);
2036        inner
2037            .last_published_epoch
2038            .retain(|key, _| key.connection_id != connection_id);
2039        inner
2040            .status
2041            .retain(|(key, _), _| key.connection_id != connection_id);
2042
2043        released
2044    }
2045
2046    pub(crate) fn escalate_client_delivery_failure(
2047        &self,
2048        connection_id: ConnectionId,
2049        channel: u16,
2050        expected_epoch: u32,
2051        reason: CloseReason,
2052    ) -> Result<bool, ForwardingError> {
2053        let should_close = {
2054            let mut inner = self.write_inner()?;
2055            let key = ClientRouteKey {
2056                connection_id,
2057                channel,
2058            };
2059            if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2060                false
2061            } else {
2062                inner.closing_connections.insert(connection_id);
2063                true
2064            }
2065        };
2066        if should_close {
2067            self.request_connection_close(connection_id, reason);
2068        }
2069        Ok(should_close)
2070    }
2071
2072    fn record_route_release(&self, release: &RouteRelease) {
2073        match release {
2074            RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2075            RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2076            RouteRelease::Absent => {}
2077        }
2078    }
2079
2080    fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2081        self.inner.read().map_err(|_| ForwardingError::Poisoned)
2082    }
2083
2084    fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2085        self.inner.write().map_err(|_| ForwardingError::Poisoned)
2086    }
2087
2088    fn lock_close_registry(
2089        &self,
2090    ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2091        self.close_registry
2092            .lock()
2093            .unwrap_or_else(|poisoned| poisoned.into_inner())
2094    }
2095}
2096
2097impl ForwardingInner {
2098    fn allocate_route_slots(
2099        &mut self,
2100        connection_id: ConnectionId,
2101        endpoint: ModuleEndpointId,
2102    ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2103        let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2104        let mut client_channel = client_start;
2105        let client_channel = loop {
2106            let key = ClientRouteKey {
2107                connection_id,
2108                channel: client_channel,
2109            };
2110            let eligible = !self.client_to_module.contains_key(&key)
2111                && !self.reserved_client.contains_key(&key)
2112                && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2113            if eligible {
2114                break client_channel;
2115            }
2116            client_channel = next_channel(client_channel);
2117            if client_channel == client_start {
2118                return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2119            }
2120        };
2121
2122        let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2123        let mut module_channel = module_start;
2124        let module_channel = loop {
2125            let key = ModuleRouteKey {
2126                endpoint,
2127                channel: module_channel,
2128            };
2129            let eligible = !self.module_to_client.contains_key(&key)
2130                && !self.reserved_module.contains_key(&key)
2131                && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2132            if eligible {
2133                break module_channel;
2134            }
2135            module_channel = next_channel(module_channel);
2136            if module_channel == module_start {
2137                return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2138            }
2139        };
2140
2141        let client_key = ClientRouteKey {
2142            connection_id,
2143            channel: client_channel,
2144        };
2145        let module_key = ModuleRouteKey {
2146            endpoint,
2147            channel: module_channel,
2148        };
2149        let client_epoch = self
2150            .client_slot_epochs
2151            .get(&client_key)
2152            .copied()
2153            .unwrap_or(0)
2154            + 1;
2155        let module_epoch = self
2156            .module_slot_epochs
2157            .get(&module_key)
2158            .copied()
2159            .unwrap_or(0)
2160            + 1;
2161        self.client_slot_epochs.insert(client_key, client_epoch);
2162        self.module_slot_epochs.insert(module_key, module_epoch);
2163        self.next_client_channel
2164            .insert(connection_id, next_channel(client_channel));
2165        self.next_module_channel
2166            .insert(endpoint, next_channel(module_channel));
2167        Ok((client_channel, client_epoch, module_channel, module_epoch))
2168    }
2169
2170    fn allocate_control_corr(
2171        &mut self,
2172        endpoint: ModuleEndpointId,
2173    ) -> Result<u64, ForwardingError> {
2174        let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2175        if candidate == 0 {
2176            self.closing_connections.insert(endpoint.connection_id);
2177            return Err(ForwardingError::RelayCorrelationExhausted);
2178        }
2179        self.next_control_corr.insert(
2180            endpoint,
2181            if candidate == u64::MAX {
2182                0
2183            } else {
2184                candidate + 1
2185            },
2186        );
2187        Ok(candidate)
2188    }
2189}
2190
2191fn next_channel(channel: u16) -> u16 {
2192    let next = channel.wrapping_add(1);
2193    if next == 0 {
2194        1
2195    } else {
2196        next
2197    }
2198}
2199
2200fn endpoint_routes_locked(
2201    inner: &ForwardingInner,
2202    endpoint: ModuleEndpointId,
2203) -> Vec<EndpointRoute> {
2204    let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2205    let draining = drain_reason.is_some();
2206    let mut routes = inner
2207        .module_to_client
2208        .iter()
2209        .filter(|(module_key, _)| module_key.endpoint == endpoint)
2210        .map(|(_, route)| EndpointRoute {
2211            goodbye_target: GoodbyeTarget {
2212                connection_id: route.client_connection_id,
2213                sink: route.client_sink.clone(),
2214                negotiated_ver: route.client_negotiated_ver,
2215                channel: route.client_channel,
2216                epoch: route.client_epoch,
2217                kind: GoodbyeTargetKind::Client,
2218                module_id: None,
2219            },
2220            principal: route.principal.clone(),
2221            bound_at: route.bound_at,
2222            draining,
2223            drain_reason,
2224        })
2225        .collect::<Vec<_>>();
2226    routes.sort_by_key(|route| {
2227        (
2228            route.goodbye_target.connection_id.get(),
2229            route.goodbye_target.channel,
2230            route.goodbye_target.epoch,
2231        )
2232    });
2233    routes
2234}
2235
2236fn release_reserved_route_locked(
2237    inner: &mut ForwardingInner,
2238    client_key: ClientRouteKey,
2239    module_key: ModuleRouteKey,
2240) {
2241    if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
2242        inner.reserved_client.remove(&client_key);
2243    }
2244    if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
2245        inner.reserved_module.remove(&module_key);
2246    }
2247    inner.status.retain(|(key, _), _| *key != client_key);
2248}
2249
2250fn release_client_route_locked(
2251    inner: &mut ForwardingInner,
2252    client_key: ClientRouteKey,
2253    expected_epoch: u32,
2254) -> RouteRelease {
2255    let Some(route) = inner.client_to_module.get(&client_key) else {
2256        return RouteRelease::Absent;
2257    };
2258    if route.client_epoch != expected_epoch {
2259        return RouteRelease::Stale;
2260    }
2261    let route = inner
2262        .client_to_module
2263        .remove(&client_key)
2264        .expect("route checked under the same forwarding lock");
2265    route.flow.close();
2266    inner.module_to_client.remove(&ModuleRouteKey {
2267        endpoint: route.module_endpoint,
2268        channel: route.module_channel,
2269    });
2270    inner.status.remove(&(client_key, expected_epoch));
2271    RouteRelease::Removed(GoodbyeTarget {
2272        connection_id: route.module_endpoint.connection_id,
2273        sink: route.module_sink.clone(),
2274        negotiated_ver: route.module_negotiated_ver,
2275        channel: route.module_channel,
2276        epoch: route.module_epoch,
2277        kind: GoodbyeTargetKind::Module,
2278        module_id: Some(route.module_id.clone()),
2279    })
2280}
2281
2282fn release_module_route_locked(
2283    inner: &mut ForwardingInner,
2284    module_key: ModuleRouteKey,
2285    expected_epoch: u32,
2286) -> RouteRelease {
2287    let Some(route) = inner.module_to_client.get(&module_key) else {
2288        return RouteRelease::Absent;
2289    };
2290    if route.module_epoch != expected_epoch {
2291        return RouteRelease::Stale;
2292    }
2293    let route = inner
2294        .module_to_client
2295        .remove(&module_key)
2296        .expect("route checked under the same forwarding lock");
2297    route.flow.close();
2298    let client_key = ClientRouteKey {
2299        connection_id: route.client_connection_id,
2300        channel: route.client_channel,
2301    };
2302    inner.client_to_module.remove(&client_key);
2303    inner.status.remove(&(client_key, route.client_epoch));
2304    RouteRelease::Removed(GoodbyeTarget {
2305        connection_id: route.client_connection_id,
2306        sink: route.client_sink.clone(),
2307        negotiated_ver: route.client_negotiated_ver,
2308        channel: route.client_channel,
2309        epoch: route.client_epoch,
2310        kind: GoodbyeTargetKind::Client,
2311        module_id: None,
2312    })
2313}
2314
2315fn commit_route_locked(
2316    inner: &mut ForwardingInner,
2317    pending: PendingRouteBindRelayEntry,
2318) -> Result<Option<GoodbyeTarget>, ForwardingError> {
2319    let reservation = pending.reservation;
2320    if inner
2321        .closing_connections
2322        .contains(&reservation.client_key.connection_id)
2323    {
2324        return Err(ForwardingError::ConnectionClosing {
2325            connection_id: reservation.client_key.connection_id,
2326        });
2327    }
2328    let module_id = inner
2329        .module_id_by_endpoint
2330        .get(&reservation.module_key.endpoint)
2331        .cloned()
2332        .ok_or(ForwardingError::StaleModuleEndpoint)?;
2333    if inner
2334        .draining_endpoints
2335        .contains_key(&reservation.module_key.endpoint)
2336    {
2337        return Err(ForwardingError::ModuleReloading { module_id });
2338    }
2339    if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
2340        || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
2341    {
2342        return Err(ForwardingError::UnknownReservation {
2343            client_channel: reservation.client_key.channel,
2344            module_channel: reservation.module_key.channel,
2345        });
2346    }
2347    let module = inner
2348        .modules_by_id
2349        .get(&module_id)
2350        .filter(|module| module.endpoint == reservation.module_key.endpoint)
2351        .cloned()
2352        .ok_or(ForwardingError::StaleModuleEndpoint)?;
2353    let binding = Arc::new(RouteBinding {
2354        client_connection_id: reservation.client_key.connection_id,
2355        client_sink: pending.client_sink,
2356        client_negotiated_ver: pending.client_negotiated_ver,
2357        client_channel: reservation.client_key.channel,
2358        client_epoch: reservation.client_epoch,
2359        module_id,
2360        module_endpoint: reservation.module_key.endpoint,
2361        module_sink: module.sink,
2362        module_negotiated_ver: module.negotiated_ver,
2363        module_channel: reservation.module_key.channel,
2364        module_epoch: reservation.module_epoch,
2365        principal: pending.principal,
2366        project_root: reservation.project_root.clone(),
2367        bound_at: Instant::now(),
2368        flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
2369    });
2370    inner
2371        .client_to_module
2372        .insert(reservation.client_key, Arc::clone(&binding));
2373    inner
2374        .module_to_client
2375        .insert(reservation.module_key, binding);
2376    let previous_published = inner
2377        .last_published_epoch
2378        .insert(reservation.client_key, reservation.client_epoch);
2379
2380    // OwnedPermit::send cannot fail, but its returned sender reveals a receiver
2381    // that closed after reservation and before this locked publication point.
2382    // Stamp at publication: the permit was reserved earlier, but queue residency
2383    // for the reply-write diagnosis starts when the frame actually enters the queue.
2384    let client_sender = pending.client_permit.send(crate::router::OutboundFrame {
2385        frame: pending.route_open_frame,
2386        enqueued_at: std::time::Instant::now(),
2387        flushed: None,
2388    });
2389    if client_sender.is_closed() {
2390        let abandoned = pending
2391            .relay_enqueued
2392            .then(|| abandoned_route_target(inner, &reservation))
2393            .flatten();
2394        if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
2395            route.flow.close();
2396        }
2397        inner.module_to_client.remove(&reservation.module_key);
2398        inner
2399            .status
2400            .remove(&(reservation.client_key, reservation.client_epoch));
2401        match previous_published {
2402            Some(epoch) => {
2403                inner
2404                    .last_published_epoch
2405                    .insert(reservation.client_key, epoch);
2406            }
2407            None => {
2408                inner.last_published_epoch.remove(&reservation.client_key);
2409            }
2410        }
2411        let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2412            "client egress closed during route publication".to_string(),
2413        ));
2414        return Ok(abandoned);
2415    }
2416
2417    let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
2418    Ok(None)
2419}
2420
2421/// The live module connection registered under exactly `endpoint`, whichever
2422/// slot holds it: active, swap candidate, or superseded incumbent.
2423///
2424/// Resolving by endpoint rather than by module id is what keeps an incumbent
2425/// addressable after a swap promoted a candidate over its id. An endpoint that
2426/// is in none of the three (a stale one, replaced without a promotion) resolves
2427/// to nothing, as it always has.
2428fn module_connection_for_endpoint_locked(
2429    inner: &ForwardingInner,
2430    endpoint: ModuleEndpointId,
2431) -> Option<&ModuleConnection> {
2432    let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
2433    inner
2434        .modules_by_id
2435        .get(module_id)
2436        .filter(|module| module.endpoint == endpoint)
2437        .or_else(|| {
2438            inner
2439                .candidates_by_id
2440                .get(module_id)
2441                .filter(|module| module.endpoint == endpoint)
2442        })
2443        .or_else(|| inner.superseded_endpoints.get(&endpoint))
2444}
2445
2446fn abandoned_route_target(
2447    inner: &ForwardingInner,
2448    reservation: &RouteReservation,
2449) -> Option<GoodbyeTarget> {
2450    let module_id = inner
2451        .module_id_by_endpoint
2452        .get(&reservation.module_key.endpoint)?;
2453    let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
2454    (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
2455        connection_id: module.endpoint.connection_id,
2456        sink: module.sink.clone(),
2457        negotiated_ver: module.negotiated_ver,
2458        channel: reservation.module_key.channel,
2459        epoch: reservation.module_epoch,
2460        kind: GoodbyeTargetKind::Module,
2461        module_id: Some(module_id.clone()),
2462    })
2463}
2464
2465fn remove_module_connection_locked(
2466    inner: &mut ForwardingInner,
2467    endpoint: ModuleEndpointId,
2468) -> Vec<GoodbyeTarget> {
2469    inner.draining_endpoints.remove(&endpoint);
2470    let module_id = inner.module_id_by_endpoint.remove(&endpoint);
2471    if let Some(module_id) = module_id.as_ref() {
2472        if inner
2473            .modules_by_id
2474            .get(module_id)
2475            .is_some_and(|module| module.endpoint == endpoint)
2476        {
2477            inner.modules_by_id.remove(module_id);
2478        }
2479        if inner
2480            .candidates_by_id
2481            .get(module_id)
2482            .is_some_and(|module| module.endpoint == endpoint)
2483        {
2484            inner.candidates_by_id.remove(module_id);
2485        }
2486    }
2487    inner.superseded_endpoints.remove(&endpoint);
2488    inner.endpoint_by_connection.remove(&endpoint.connection_id);
2489    inner.next_module_channel.remove(&endpoint);
2490    inner.next_control_corr.remove(&endpoint);
2491    inner
2492        .health_probe_tombstones
2493        .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
2494    inner
2495        .module_slot_epochs
2496        .retain(|key, _| key.endpoint != endpoint);
2497    let reserved_module_keys: Vec<ModuleRouteKey> = inner
2498        .reserved_module
2499        .keys()
2500        .filter(|module_key| module_key.endpoint == endpoint)
2501        .copied()
2502        .collect();
2503    for module_key in reserved_module_keys {
2504        if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2505            release_reserved_route_locked(inner, client_key, module_key);
2506        }
2507    }
2508
2509    let pending_keys: Vec<_> = inner
2510        .pending_relays
2511        .keys()
2512        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2513        .copied()
2514        .collect();
2515    let pending: Vec<_> = pending_keys
2516        .into_iter()
2517        .filter_map(|key| inner.pending_relays.remove(&key))
2518        .collect();
2519    for pending in pending {
2520        let module_label = module_id.as_deref().unwrap_or("unknown");
2521        let _ = pending
2522            .sender
2523            .send(RouteBindRelayOutcome::ModuleGone(format!(
2524                "module '{module_label}' connection closed during route.bind relay"
2525            )));
2526    }
2527
2528    let pending_control_keys: Vec<_> = inner
2529        .pending_control_rpcs
2530        .keys()
2531        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2532        .copied()
2533        .collect();
2534    let pending_control: Vec<_> = pending_control_keys
2535        .into_iter()
2536        .filter_map(|key| inner.pending_control_rpcs.remove(&key))
2537        .collect();
2538    for pending in pending_control {
2539        let module_label = module_id.as_deref().unwrap_or("unknown");
2540        let _ = pending
2541            .sender
2542            .send(ModuleControlRpcOutcome::ModuleGone(format!(
2543                "module '{module_label}' connection closed during module-control RPC"
2544            )));
2545    }
2546
2547    let module_routes = inner
2548        .module_to_client
2549        .iter()
2550        .filter(|(module_key, _)| module_key.endpoint == endpoint)
2551        .map(|(module_key, route)| (*module_key, route.module_epoch))
2552        .collect::<Vec<_>>();
2553    let mut released = Vec::with_capacity(module_routes.len());
2554    for (module_key, epoch) in module_routes {
2555        if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
2556        {
2557            released.push(target);
2558        }
2559    }
2560    released
2561}
2562
2563#[derive(Debug, Clone, Copy)]
2564struct RequestCredit {
2565    subscription: bool,
2566    excluded_from_drain: bool,
2567}
2568
2569#[derive(Debug, Default)]
2570struct CreditLedger {
2571    by_corr: HashMap<u64, Vec<RequestCredit>>,
2572}
2573
2574impl CreditLedger {
2575    fn acquire(&mut self, corr: u64, subscription: bool) {
2576        self.by_corr.entry(corr).or_default().push(RequestCredit {
2577            subscription,
2578            excluded_from_drain: false,
2579        });
2580    }
2581
2582    fn release(&mut self, corr: u64) -> bool {
2583        let Some(credits) = self.by_corr.get_mut(&corr) else {
2584            return false;
2585        };
2586        let released = credits.pop().is_some();
2587        if credits.is_empty() {
2588            self.by_corr.remove(&corr);
2589        }
2590        released
2591    }
2592
2593    fn capture_subscription_exclusions(&mut self) -> u32 {
2594        let mut excluded = 0u32;
2595        for credit in self.by_corr.values_mut().flatten() {
2596            if credit.subscription && !credit.excluded_from_drain {
2597                credit.excluded_from_drain = true;
2598                excluded = excluded.saturating_add(1);
2599            }
2600        }
2601        excluded
2602    }
2603
2604    #[cfg(test)]
2605    fn in_flight(&self) -> usize {
2606        self.by_corr.values().map(Vec::len).sum()
2607    }
2608
2609    fn drain_in_flight(&self) -> usize {
2610        self.by_corr
2611            .values()
2612            .flatten()
2613            .filter(|credit| !credit.excluded_from_drain)
2614            .count()
2615    }
2616}
2617
2618#[derive(Debug, Default)]
2619struct ChannelFlowState {
2620    closed: bool,
2621    credits: CreditLedger,
2622}
2623
2624/// Per-channel request-credit accounting shared by the client and module route halves.
2625#[derive(Debug)]
2626pub(crate) struct ChannelFlow {
2627    sem: Semaphore,
2628    window: usize,
2629    state: Mutex<ChannelFlowState>,
2630}
2631
2632impl ChannelFlow {
2633    pub(crate) fn new(window: usize) -> Self {
2634        debug_assert!(window > 0, "flow-control window must be non-zero");
2635        Self {
2636            sem: Semaphore::new(window),
2637            window,
2638            state: Mutex::new(ChannelFlowState::default()),
2639        }
2640    }
2641
2642    #[cfg(test)]
2643    pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
2644        self.acquire_tagged(0, false).await
2645    }
2646
2647    pub(crate) async fn acquire_tagged(
2648        &self,
2649        corr: u64,
2650        subscription: bool,
2651    ) -> Result<(), ChannelFlowClosed> {
2652        let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
2653        let mut state = self
2654            .state
2655            .lock()
2656            .unwrap_or_else(|poisoned| poisoned.into_inner());
2657        if state.closed {
2658            return Err(ChannelFlowClosed);
2659        }
2660        state.credits.acquire(corr, subscription);
2661        permit.forget();
2662        Ok(())
2663    }
2664
2665    #[cfg(test)]
2666    pub(crate) fn release(&self) {
2667        self.release_corr(0);
2668    }
2669
2670    pub(crate) fn release_corr(&self, corr: u64) {
2671        let released = self
2672            .state
2673            .lock()
2674            .unwrap_or_else(|poisoned| poisoned.into_inner())
2675            .credits
2676            .release(corr);
2677        if !released {
2678            // Protocol-conforming modules emit exactly one terminal per request.
2679            // This guard is a best-effort safety net against window growth, not a
2680            // security boundary against malicious peers.
2681            warn!(
2682                window = self.window,
2683                available = self.sem.available_permits(),
2684                "flow-control over-release ignored"
2685            );
2686            return;
2687        }
2688        if !self.sem.is_closed() {
2689            self.sem.add_permits(1);
2690        }
2691    }
2692
2693    #[cfg(test)]
2694    pub(crate) fn in_flight(&self) -> usize {
2695        self.state
2696            .lock()
2697            .unwrap_or_else(|poisoned| poisoned.into_inner())
2698            .credits
2699            .in_flight()
2700    }
2701
2702    pub(crate) fn drain_in_flight(&self) -> usize {
2703        self.state
2704            .lock()
2705            .unwrap_or_else(|poisoned| poisoned.into_inner())
2706            .credits
2707            .drain_in_flight()
2708    }
2709
2710    #[cfg(test)]
2711    pub(crate) fn available_permits(&self) -> usize {
2712        self.sem.available_permits()
2713    }
2714
2715    pub(crate) fn begin_drain(&self) -> u32 {
2716        let mut state = self
2717            .state
2718            .lock()
2719            .unwrap_or_else(|poisoned| poisoned.into_inner());
2720        state.closed = true;
2721        self.sem.close();
2722        state.credits.capture_subscription_exclusions()
2723    }
2724
2725    pub(crate) fn close(&self) {
2726        self.state
2727            .lock()
2728            .unwrap_or_else(|poisoned| poisoned.into_inner())
2729            .closed = true;
2730        self.sem.close();
2731    }
2732}
2733
2734#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2735pub(crate) struct ChannelFlowClosed;
2736
2737impl fmt::Display for ChannelFlowClosed {
2738    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2739        write!(f, "flow-control window closed")
2740    }
2741}
2742
2743impl Error for ChannelFlowClosed {}
2744
2745fn window_for(concurrency: &Concurrency) -> usize {
2746    match concurrency {
2747        Concurrency::Serial => 1,
2748        Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
2749        Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
2750    }
2751}
2752
2753#[derive(Debug, Clone, PartialEq, Eq)]
2754pub enum ForwardingError {
2755    NoModuleConnection,
2756    ModuleReloading {
2757        module_id: String,
2758    },
2759    StaleModuleEndpoint,
2760    UnknownReservation {
2761        client_channel: u16,
2762        module_channel: u16,
2763    },
2764    ClientRouteChannelExhausted {
2765        connection_id: ConnectionId,
2766    },
2767    ModuleRouteChannelExhausted {
2768        endpoint: ModuleEndpointId,
2769    },
2770    RelayCorrelationExhausted,
2771    ConnectionClosing {
2772        connection_id: ConnectionId,
2773    },
2774    ClientEgressClosed {
2775        connection_id: ConnectionId,
2776    },
2777    RouteOpenBuild(String),
2778    /// A swap candidate is already registered for this module id.
2779    CandidateSlotOccupied {
2780        module_id: String,
2781    },
2782    Poisoned,
2783}
2784
2785impl fmt::Display for ForwardingError {
2786    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2787        match self {
2788            Self::NoModuleConnection => write!(f, "no module connection is registered"),
2789            Self::ModuleReloading { module_id } => {
2790                write!(f, "module_id '{module_id}' is reloading")
2791            }
2792            Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
2793            Self::UnknownReservation {
2794                client_channel,
2795                module_channel,
2796            } => write!(
2797                f,
2798                "route reservation client channel {client_channel} / module channel {module_channel} was not found"
2799            ),
2800            Self::ClientRouteChannelExhausted { connection_id } => write!(
2801                f,
2802                "no client route channels are available for connection {}",
2803                connection_id.get()
2804            ),
2805            Self::ModuleRouteChannelExhausted { endpoint } => write!(
2806                f,
2807                "no module route channels are available for endpoint generation {} on connection {}",
2808                endpoint.generation,
2809                endpoint.connection_id.get()
2810            ),
2811            Self::RelayCorrelationExhausted => {
2812                write!(f, "module control correlation ids are exhausted")
2813            }
2814            Self::ConnectionClosing { connection_id } => write!(
2815                f,
2816                "connection {} is closing and cannot accept route allocation",
2817                connection_id.get()
2818            ),
2819            Self::ClientEgressClosed { connection_id } => write!(
2820                f,
2821                "client connection {} egress is closed",
2822                connection_id.get()
2823            ),
2824            Self::RouteOpenBuild(message) => {
2825                write!(f, "failed to prebuild route.open response: {message}")
2826            }
2827            Self::CandidateSlotOccupied { module_id } => write!(
2828                f,
2829                "module_id '{module_id}' already has a swap candidate registered"
2830            ),
2831            Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
2832        }
2833    }
2834}
2835
2836impl Error for ForwardingError {}
2837
2838#[cfg(test)]
2839mod tests {
2840    use std::time::Duration;
2841
2842    use super::*;
2843    use tokio::sync::mpsc;
2844
2845    #[test]
2846    fn ordinary_long_running_request_is_not_excluded_from_drain() {
2847        let mut ledger = CreditLedger::default();
2848        ledger.acquire(1, false);
2849
2850        assert_eq!(ledger.capture_subscription_exclusions(), 0);
2851        assert_eq!(ledger.drain_in_flight(), 1);
2852    }
2853
2854    #[test]
2855    fn bit_set_subscription_is_excluded_and_counted() {
2856        let mut ledger = CreditLedger::default();
2857        ledger.acquire(1, true);
2858
2859        assert_eq!(ledger.capture_subscription_exclusions(), 1);
2860        assert_eq!(ledger.drain_in_flight(), 0);
2861    }
2862
2863    #[test]
2864    fn subscription_opened_after_drain_snapshot_is_not_excluded() {
2865        let mut ledger = CreditLedger::default();
2866        ledger.acquire(1, true);
2867        assert_eq!(ledger.capture_subscription_exclusions(), 1);
2868
2869        ledger.acquire(2, true);
2870
2871        assert_eq!(ledger.drain_in_flight(), 1);
2872    }
2873
2874    #[test]
2875    fn drain_with_no_subscriptions_reports_zero_excluded() {
2876        let mut ledger = CreditLedger::default();
2877        assert_eq!(ledger.capture_subscription_exclusions(), 0);
2878    }
2879
2880    #[test]
2881    fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
2882        let forwarding = ForwardingTable::default();
2883        let module_connection = ConnectionId::new(10);
2884        let exhausted_client = ConnectionId::new(20);
2885        let second_client = ConnectionId::new(30);
2886        let (module_tx, _module_rx) = mpsc::channel(1);
2887        let endpoint = forwarding
2888            .register_module_connection(
2889                module_connection,
2890                "route-limit-provider".to_string(),
2891                1,
2892                Concurrency::ModuleManaged,
2893                FrameSink::new(module_tx),
2894            )
2895            .unwrap();
2896
2897        {
2898            let mut inner = forwarding.inner.write().unwrap();
2899            for channel in 1..=u16::MAX {
2900                inner.reserved_client.insert(
2901                    ClientRouteKey {
2902                        connection_id: exhausted_client,
2903                        channel,
2904                    },
2905                    ModuleRouteKey {
2906                        endpoint,
2907                        channel: 1,
2908                    },
2909                );
2910            }
2911        }
2912
2913        let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
2914        let err = forwarding
2915            .begin_route_bind_relay_for_test(
2916                exhausted_client,
2917                FrameSink::new(exhausted_tx),
2918                1,
2919                "route-limit-provider",
2920            )
2921            .unwrap_err();
2922        assert!(matches!(
2923            err,
2924            ForwardingError::ClientRouteChannelExhausted { connection_id }
2925                if connection_id == exhausted_client
2926        ));
2927
2928        let (second_tx, _second_rx) = mpsc::channel(1);
2929        let pending = forwarding
2930            .begin_route_bind_relay_for_test(
2931                second_client,
2932                FrameSink::new(second_tx),
2933                2,
2934                "route-limit-provider",
2935            )
2936            .unwrap();
2937        assert_eq!(pending.client_channel, 1);
2938    }
2939
2940    #[test]
2941    fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
2942        let forwarding = ForwardingTable::default();
2943        let module_connection = ConnectionId::new(40);
2944        let client = ConnectionId::new(50);
2945        let (module_tx, _module_rx) = mpsc::channel(1);
2946        forwarding
2947            .register_module_connection(
2948                module_connection,
2949                "slot-reuse-provider".to_string(),
2950                1,
2951                Concurrency::ModuleManaged,
2952                FrameSink::new(module_tx),
2953            )
2954            .unwrap();
2955
2956        let (client_tx, _client_rx) = mpsc::channel(1);
2957        let client_sink = FrameSink::new(client_tx);
2958        let mut wrapped_channel = None;
2959        for index in 0..=usize::from(u16::MAX) {
2960            let pending = forwarding
2961                .begin_route_bind_relay_for_test(
2962                    client,
2963                    client_sink.clone(),
2964                    index as u64 + 1,
2965                    "slot-reuse-provider",
2966                )
2967                .unwrap();
2968            if index == usize::from(u16::MAX) {
2969                wrapped_channel = Some(pending.module_channel);
2970            }
2971            forwarding
2972                .abort_pending_relay(
2973                    pending.endpoint,
2974                    pending.corr,
2975                    RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
2976                )
2977                .unwrap();
2978        }
2979
2980        assert_eq!(wrapped_channel, Some(1));
2981    }
2982
2983    #[test]
2984    fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
2985        let forwarding = ForwardingTable::default();
2986        let client = ConnectionId::new(60);
2987        forwarding
2988            .inner
2989            .write()
2990            .unwrap()
2991            .next_client_channel
2992            .insert(client, 41);
2993
2994        let released = forwarding.cleanup_connection(client).unwrap();
2995
2996        assert!(released.is_empty());
2997        assert!(!forwarding
2998            .inner
2999            .read()
3000            .unwrap()
3001            .next_client_channel
3002            .contains_key(&client));
3003    }
3004
3005    #[test]
3006    fn stale_module_cleanup_preserves_fast_reconnect_successor() {
3007        let forwarding = ForwardingTable::default();
3008        let module_id = "fast-reconnect-provider";
3009        let first_connection = ConnectionId::new(70);
3010        let second_connection = ConnectionId::new(80);
3011        let (first_tx, _first_rx) = mpsc::channel(1);
3012        let first_endpoint = forwarding
3013            .register_module_connection(
3014                first_connection,
3015                module_id.to_string(),
3016                1,
3017                Concurrency::ModuleManaged,
3018                FrameSink::new(first_tx),
3019            )
3020            .unwrap();
3021        let (second_tx, _second_rx) = mpsc::channel(1);
3022        let second_endpoint = forwarding
3023            .register_module_connection(
3024                second_connection,
3025                module_id.to_string(),
3026                1,
3027                Concurrency::ModuleManaged,
3028                FrameSink::new(second_tx),
3029            )
3030            .unwrap();
3031        assert_ne!(first_endpoint, second_endpoint);
3032
3033        let released = forwarding.cleanup_connection(first_connection).unwrap();
3034
3035        assert!(released.is_empty());
3036        assert_eq!(
3037            forwarding
3038                .inner
3039                .read()
3040                .unwrap()
3041                .modules_by_id
3042                .get(module_id)
3043                .map(|module| module.endpoint),
3044            Some(second_endpoint)
3045        );
3046        assert!(forwarding.has_live_module_connection(module_id).unwrap());
3047        let control_rpc = forwarding
3048            .begin_module_control_rpc_for(
3049                module_id,
3050                "health.check",
3051                Instant::now() + Duration::from_secs(1),
3052            )
3053            .unwrap();
3054        assert_eq!(control_rpc.endpoint, second_endpoint);
3055    }
3056
3057    fn route_fixture(
3058        module_id: &str,
3059    ) -> (
3060        ForwardingTable,
3061        ConnectionId,
3062        ModuleEndpointId,
3063        ConnectionId,
3064        FrameSink,
3065        mpsc::Receiver<crate::router::OutboundFrame>,
3066    ) {
3067        let forwarding = ForwardingTable::default();
3068        let module_connection = ConnectionId::new(100);
3069        let client_connection = ConnectionId::new(200);
3070        let (module_tx, _module_rx) = mpsc::channel(8);
3071        let endpoint = forwarding
3072            .register_module_connection(
3073                module_connection,
3074                module_id.to_string(),
3075                2,
3076                Concurrency::ModuleManaged,
3077                FrameSink::new(module_tx),
3078            )
3079            .unwrap();
3080        let (client_tx, client_rx) = mpsc::channel(8);
3081        (
3082            forwarding,
3083            module_connection,
3084            endpoint,
3085            client_connection,
3086            FrameSink::new(client_tx),
3087            client_rx,
3088        )
3089    }
3090
3091    #[test]
3092    #[cfg(unix)]
3093    fn daemon_drain_gates_current_and_racing_provider_registrations() {
3094        let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
3095        assert_eq!(forwarding.begin_daemon_drain().unwrap(), ["provider"]);
3096        assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
3097        assert!(matches!(
3098            forwarding.register_module_connection(
3099                ConnectionId::new(300),
3100                "late-provider".into(),
3101                2,
3102                Concurrency::ModuleManaged,
3103                sink,
3104            ),
3105            Err(ForwardingError::ConnectionClosing { .. })
3106        ));
3107    }
3108
3109    fn test_ping(corr: u64) -> Frame {
3110        Frame::build(
3111            FrameType::Ping,
3112            Flags::new(false, Priority::Passive, false),
3113            0,
3114            0,
3115            corr,
3116            Vec::new(),
3117        )
3118        .unwrap()
3119    }
3120
3121    fn begin_test_route(
3122        forwarding: &ForwardingTable,
3123        client_connection: ConnectionId,
3124        client_sink: FrameSink,
3125        corr: u64,
3126        module_id: &str,
3127    ) -> PendingRouteBindRelay {
3128        forwarding
3129            .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
3130            .unwrap()
3131    }
3132
3133    /// The drain-timeout line reports these numbers, so they must count the
3134    /// requests the drain is waiting on and no others: a route with nothing in
3135    /// flight is not a holdout, and a flagged subscription is excluded from the
3136    /// drain and so from the count.
3137    #[tokio::test]
3138    async fn drain_holdouts_count_held_requests_and_name_the_connection() {
3139        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3140            route_fixture("holdouts");
3141        let mut bound = |corr| {
3142            let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
3143            forwarding
3144                .complete_pending_relay(
3145                    module_connection,
3146                    route.corr,
3147                    RouteBindRelayOutcome::Accepted,
3148                )
3149                .unwrap();
3150            client_rx.try_recv().unwrap();
3151            match forwarding
3152                .lookup_data_route(client, route.client_channel, route.client_epoch)
3153                .unwrap()
3154            {
3155                DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3156                other => panic!("expected live route, got {other:?}"),
3157            }
3158        };
3159        let holding = bound(61);
3160        let _idle = bound(62);
3161        holding.flow.acquire_tagged(1, false).await.unwrap();
3162        holding.flow.acquire_tagged(2, false).await.unwrap();
3163        holding.flow.acquire_tagged(3, true).await.unwrap();
3164        forwarding
3165            .begin_module_drain("holdouts", RouteCloseReason::Restart)
3166            .unwrap();
3167
3168        let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
3169        assert_eq!(
3170            holdouts,
3171            DrainHoldouts {
3172                requests: 2,
3173                routes: 1,
3174                total_routes: 2,
3175                top_connections: vec![(client.get(), 2)],
3176            }
3177        );
3178    }
3179
3180    #[test]
3181    fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
3182        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3183            route_fixture("census");
3184        let pending = begin_test_route(&forwarding, client, sink, 1, "census");
3185        forwarding
3186            .complete_pending_relay(
3187                module_connection,
3188                pending.corr,
3189                RouteBindRelayOutcome::Accepted,
3190            )
3191            .unwrap();
3192
3193        let routes = forwarding.endpoint_routes(endpoint).unwrap();
3194        assert_eq!(routes.len(), 1);
3195        assert!(matches!(routes[0].principal, Principal::Direct));
3196        assert_eq!(routes[0].goodbye_target.connection_id, client);
3197        assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
3198        assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
3199        assert!(!routes[0].draining);
3200
3201        forwarding
3202            .begin_module_drain("census", RouteCloseReason::Restart)
3203            .unwrap();
3204        let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
3205        assert_eq!(draining_routes.len(), 1);
3206        assert!(draining_routes[0].draining);
3207    }
3208
3209    #[test]
3210    fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
3211        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
3212        let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
3213        assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
3214        forwarding
3215            .abort_pending_relay(
3216                first.endpoint,
3217                first.corr,
3218                RouteBindRelayOutcome::ModuleGone("abort".into()),
3219            )
3220            .unwrap();
3221        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3222        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3223
3224        let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
3225        assert_eq!(second.client_channel, first.client_channel);
3226        assert_eq!(second.module_channel, first.module_channel);
3227        assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
3228    }
3229
3230    #[test]
3231    fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
3232        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3233            route_fixture("epoch-release");
3234        let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
3235        forwarding
3236            .complete_pending_relay(
3237                module_connection,
3238                first.corr,
3239                RouteBindRelayOutcome::Accepted,
3240            )
3241            .unwrap();
3242        assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
3243        assert!(matches!(
3244            forwarding
3245                .release_client_route(client, first.client_channel, first.client_epoch)
3246                .unwrap(),
3247            RouteRelease::Removed(_)
3248        ));
3249        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3250        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3251
3252        let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
3253        forwarding
3254            .complete_pending_relay(
3255                module_connection,
3256                second.corr,
3257                RouteBindRelayOutcome::Accepted,
3258            )
3259            .unwrap();
3260        assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
3261        assert!(matches!(
3262            forwarding
3263                .release_client_route(client, second.client_channel, first.client_epoch)
3264                .unwrap(),
3265            RouteRelease::Stale
3266        ));
3267        assert!(!forwarding
3268            .cache_status(
3269                endpoint,
3270                second.module_channel,
3271                first.module_epoch,
3272                "stale".into(),
3273            )
3274            .unwrap());
3275        assert!(forwarding
3276            .cache_status(
3277                endpoint,
3278                second.module_channel,
3279                second.module_epoch,
3280                "current".into(),
3281            )
3282            .unwrap());
3283        match forwarding
3284            .route_poll_snapshot(client, second.client_channel, second.client_epoch)
3285            .unwrap()
3286        {
3287            RoutePollSnapshot::Bound { status, .. } => {
3288                assert_eq!(status.as_deref(), Some("current"));
3289            }
3290            RoutePollSnapshot::Absent => panic!("successor binding was removed"),
3291        }
3292        let counters = forwarding.counters().snapshot();
3293        assert_eq!(counters["route_released_epoch_fenced"], 1);
3294        assert_eq!(counters["route_release_stale_skipped"], 1);
3295    }
3296
3297    #[test]
3298    fn max_epoch_reservation_retires_only_that_slot() {
3299        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
3300        forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
3301        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
3302        let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
3303        assert_eq!(
3304            (final_use.client_channel, final_use.client_epoch),
3305            (7, u32::MAX)
3306        );
3307        assert_eq!(
3308            (final_use.module_channel, final_use.module_epoch),
3309            (9, u32::MAX)
3310        );
3311        forwarding
3312            .abort_pending_relay(
3313                endpoint,
3314                final_use.corr,
3315                RouteBindRelayOutcome::ModuleGone("abort".into()),
3316            )
3317            .unwrap();
3318        forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
3319        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
3320        let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
3321        assert_ne!(next.client_channel, 7);
3322        assert_ne!(next.module_channel, 9);
3323        assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
3324    }
3325
3326    #[test]
3327    fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
3328        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3329            route_fixture("corr-shared");
3330        let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
3331        assert_eq!(bind.corr, 1);
3332        forwarding
3333            .abort_pending_relay(
3334                endpoint,
3335                bind.corr,
3336                RouteBindRelayOutcome::ModuleGone("abort".into()),
3337            )
3338            .unwrap();
3339        let rpc = forwarding
3340            .begin_module_control_rpc_for(
3341                "corr-shared",
3342                "health.check",
3343                Instant::now() - Duration::from_millis(1),
3344            )
3345            .unwrap();
3346        assert_eq!(rpc.corr, 2);
3347        assert_eq!(
3348            forwarding
3349                .complete_module_control_rpc(
3350                    module_connection,
3351                    rpc.corr,
3352                    Some("health.check"),
3353                    ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
3354                        status: subc_protocol::session::HealthStatus::Ok,
3355                        detail: None,
3356                        metrics: None,
3357                    }),
3358                )
3359                .unwrap(),
3360            ModuleControlRpcCompletion::Settled
3361        );
3362        assert!(matches!(
3363            rpc.receiver.blocking_recv().unwrap(),
3364            ModuleControlRpcOutcome::DeadlineElapsed
3365        ));
3366    }
3367
3368    #[tokio::test(start_paused = true)]
3369    async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
3370        let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
3371        let probe_started_at = Instant::now();
3372        let rpc = forwarding
3373            .begin_health_probe_rpc_for(
3374                "tombstone-ttl",
3375                "health.check",
3376                probe_started_at,
3377                probe_started_at + Duration::from_secs(5),
3378            )
3379            .unwrap();
3380        assert!(forwarding
3381            .tombstone_health_probe_rpc(endpoint, rpc.corr)
3382            .unwrap());
3383        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
3384
3385        tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
3386        tokio::task::yield_now().await;
3387
3388        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
3389    }
3390
3391    #[test]
3392    fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
3393        let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
3394        let mut close = forwarding.register_connection_close(endpoint.connection_id);
3395        forwarding.inject_control_corr(endpoint, u64::MAX);
3396        let final_rpc = forwarding
3397            .begin_module_control_rpc_for(
3398                "corr-max",
3399                "health.check",
3400                Instant::now() + Duration::from_secs(1),
3401            )
3402            .unwrap();
3403        assert_eq!(final_rpc.corr, u64::MAX);
3404        forwarding
3405            .cancel_module_control_rpc(endpoint, final_rpc.corr)
3406            .unwrap();
3407        assert!(matches!(
3408            forwarding.begin_module_control_rpc_for(
3409                "corr-max",
3410                "health.check",
3411                Instant::now() + Duration::from_secs(1),
3412            ),
3413            Err(ForwardingError::RelayCorrelationExhausted)
3414        ));
3415        assert!(close.try_recv().is_ok());
3416    }
3417
3418    #[test]
3419    fn publication_epoch_controls_delivery_failure_escalation() {
3420        fn setup_successor(
3421            commit_successor: Option<bool>,
3422        ) -> (ForwardingTable, ConnectionId, u16, u32) {
3423            let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3424                route_fixture("escalation");
3425            let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
3426            forwarding
3427                .complete_pending_relay(
3428                    module_connection,
3429                    first.corr,
3430                    RouteBindRelayOutcome::Accepted,
3431                )
3432                .unwrap();
3433            client_rx.try_recv().unwrap();
3434            assert!(matches!(
3435                forwarding
3436                    .release_client_route(client, first.client_channel, first.client_epoch)
3437                    .unwrap(),
3438                RouteRelease::Removed(_)
3439            ));
3440            if let Some(commit_successor) = commit_successor {
3441                forwarding.inject_client_slot_epoch(
3442                    client,
3443                    first.client_channel,
3444                    first.client_epoch,
3445                );
3446                forwarding.inject_module_slot_epoch(
3447                    endpoint,
3448                    first.module_channel,
3449                    first.module_epoch,
3450                );
3451                let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
3452                if commit_successor {
3453                    forwarding
3454                        .complete_pending_relay(
3455                            module_connection,
3456                            successor.corr,
3457                            RouteBindRelayOutcome::Accepted,
3458                        )
3459                        .unwrap();
3460                    client_rx.try_recv().unwrap();
3461                } else {
3462                    forwarding
3463                        .abort_pending_relay(
3464                            endpoint,
3465                            successor.corr,
3466                            RouteBindRelayOutcome::ModuleGone("abort".into()),
3467                        )
3468                        .unwrap();
3469                }
3470            }
3471            (forwarding, client, first.client_channel, first.client_epoch)
3472        }
3473
3474        let (no_successor, client, channel, epoch) = setup_successor(None);
3475        let mut close = no_successor.register_connection_close(client);
3476        assert!(no_successor
3477            .escalate_client_delivery_failure(
3478                client,
3479                channel,
3480                epoch,
3481                CloseReason::new("delivery", "failed"),
3482            )
3483            .unwrap());
3484        assert!(close.try_recv().is_ok());
3485
3486        let (aborted, client, channel, epoch) = setup_successor(Some(false));
3487        let mut close = aborted.register_connection_close(client);
3488        assert!(aborted
3489            .escalate_client_delivery_failure(
3490                client,
3491                channel,
3492                epoch,
3493                CloseReason::new("delivery", "failed"),
3494            )
3495            .unwrap());
3496        assert!(close.try_recv().is_ok());
3497
3498        let (published, client, channel, epoch) = setup_successor(Some(true));
3499        let mut close = published.register_connection_close(client);
3500        assert!(!published
3501            .escalate_client_delivery_failure(
3502                client,
3503                channel,
3504                epoch,
3505                CloseReason::new("delivery", "stale failure"),
3506            )
3507            .unwrap());
3508        assert!(close.try_recv().is_err());
3509    }
3510
3511    #[test]
3512    fn route_concentration_separates_client_count_from_routes_per_client() {
3513        // The distinction this asserts is the one a bare connection count cannot
3514        // make: two connections holding one route each and one connection
3515        // holding two are the same total, and have opposite causes.
3516        let (forwarding, module_connection, _, client, sink, _client_rx) =
3517            route_fixture("concentration");
3518        assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
3519
3520        for corr in [70_u64, 71] {
3521            let pending =
3522                begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
3523            forwarding
3524                .complete_pending_relay(
3525                    module_connection,
3526                    pending.corr,
3527                    RouteBindRelayOutcome::Accepted,
3528                )
3529                .unwrap();
3530        }
3531
3532        // One connection, two routes — not two connections with a route each.
3533        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3534        assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
3535    }
3536
3537    #[test]
3538    fn cleanup_and_accepted_resolution_have_one_lock_winner() {
3539        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3540            route_fixture("cleanup-race");
3541        let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
3542        forwarding
3543            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3544            .unwrap();
3545        let released = forwarding.cleanup_connection(client).unwrap();
3546        assert_eq!(released.len(), 1);
3547        let completion = forwarding
3548            .complete_pending_relay(
3549                module_connection,
3550                pending.corr,
3551                RouteBindRelayOutcome::Accepted,
3552            )
3553            .unwrap();
3554        assert!(!completion.settled);
3555        assert!(client_rx.try_recv().is_err());
3556        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3557
3558        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3559            route_fixture("accepted-race");
3560        let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
3561        forwarding
3562            .complete_pending_relay(
3563                module_connection,
3564                pending.corr,
3565                RouteBindRelayOutcome::Accepted,
3566            )
3567            .unwrap();
3568        assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
3569        let released = forwarding.cleanup_connection(client).unwrap();
3570        assert_eq!(released.len(), 1);
3571        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3572    }
3573
3574    #[test]
3575    fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
3576        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3577            route_fixture("drain-gap");
3578        let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
3579        forwarding
3580            .complete_pending_relay(
3581                module_connection,
3582                live.corr,
3583                RouteBindRelayOutcome::Accepted,
3584            )
3585            .unwrap();
3586        client_rx.try_recv().unwrap();
3587        let binding = match forwarding
3588            .lookup_data_route(client, live.client_channel, live.client_epoch)
3589            .unwrap()
3590        {
3591            DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3592            other => panic!("expected live route, got {other:?}"),
3593        };
3594
3595        let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
3596        forwarding
3597            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3598            .unwrap();
3599        let control_rpc = forwarding
3600            .begin_module_control_rpc_for(
3601                "drain-gap",
3602                "health.check",
3603                Instant::now() + Duration::from_secs(1),
3604            )
3605            .unwrap();
3606        let target = forwarding
3607            .begin_module_drain("drain-gap", RouteCloseReason::Reload)
3608            .unwrap()
3609            .unwrap();
3610        assert!(matches!(
3611            control_rpc.receiver.blocking_recv().unwrap(),
3612            ModuleControlRpcOutcome::ModuleGone(_)
3613        ));
3614        assert_eq!(target.abandoned_bindings.len(), 1);
3615        assert!(binding.flow.sem.is_closed());
3616        assert!(
3617            !forwarding
3618                .complete_pending_relay(
3619                    module_connection,
3620                    pending.corr,
3621                    RouteBindRelayOutcome::Accepted,
3622                )
3623                .unwrap()
3624                .settled
3625        );
3626        assert!(matches!(
3627            forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
3628            Err(ForwardingError::ModuleReloading { .. })
3629        ));
3630        let released = forwarding
3631            .release_module_endpoint_routes(target.endpoint)
3632            .unwrap();
3633        assert_eq!(released.len(), 1);
3634        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3635    }
3636
3637    /// A client can be marked closing while its egress is still open: the daemon
3638    /// asks a connection to close (here through the production path, a module
3639    /// frame that its egress refused) and the connection loop tears down a moment
3640    /// later. A route.bind ack that lands inside that window is answered on the
3641    /// MODULE connection's frame handler, so resolving it must not produce an
3642    /// error -- an error there ends the module connection, and that connection
3643    /// carries every other client's routes to the module.
3644    #[test]
3645    fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
3646        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3647            route_fixture("closing-client");
3648
3649        // A published route on this client: escalate_client_delivery_failure only
3650        // marks a connection closing for a route it has already published.
3651        let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
3652        forwarding
3653            .complete_pending_relay(
3654                module_connection,
3655                live.corr,
3656                RouteBindRelayOutcome::Accepted,
3657            )
3658            .unwrap();
3659        client_rx.try_recv().unwrap();
3660
3661        // A second route.open from the same client, relayed and awaiting its ack.
3662        let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
3663        forwarding
3664            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3665            .unwrap();
3666
3667        // The window: closing, but the sink is still open.
3668        assert!(forwarding
3669            .escalate_client_delivery_failure(
3670                client,
3671                live.client_channel,
3672                live.client_epoch,
3673                CloseReason::new(
3674                    "module_to_client_delivery_failed",
3675                    "client egress refused a module frame",
3676                ),
3677            )
3678            .unwrap());
3679        assert!(!sink.is_closed());
3680
3681        let completion = forwarding
3682            .complete_pending_relay(
3683                module_connection,
3684                pending.corr,
3685                RouteBindRelayOutcome::Accepted,
3686            )
3687            .expect("a closing client must not turn a module's ack into an error");
3688
3689        assert!(completion.settled);
3690        let abandoned = completion
3691            .abandoned
3692            .expect("the module must be told to drop the binding it just created");
3693        assert_eq!(abandoned.connection_id, module_connection);
3694        assert_eq!(abandoned.channel, pending.module_channel);
3695        assert_eq!(abandoned.epoch, pending.module_epoch);
3696        assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
3697        assert!(matches!(
3698            pending.receiver.blocking_recv().unwrap(),
3699            RouteBindRelayOutcome::ModuleGone(_)
3700        ));
3701        // No route was published to a client that is on its way out, and the
3702        // reserved handle pair went back.
3703        assert!(client_rx.try_recv().is_err());
3704        assert_eq!(forwarding.active_binding_count().unwrap(), 1);
3705
3706        // The module endpoint is untouched: still live, and still able to take a
3707        // route from another client.
3708        assert!(forwarding
3709            .has_live_module_connection("closing-client")
3710            .unwrap());
3711        let cotenant = ConnectionId::new(201);
3712        let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
3713        let cotenant_route = begin_test_route(
3714            &forwarding,
3715            cotenant,
3716            FrameSink::new(cotenant_tx),
3717            62,
3718            "closing-client",
3719        );
3720        assert_eq!(cotenant_route.endpoint, endpoint);
3721        forwarding
3722            .complete_pending_relay(
3723                module_connection,
3724                cotenant_route.corr,
3725                RouteBindRelayOutcome::Accepted,
3726            )
3727            .unwrap();
3728        assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
3729        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3730    }
3731
3732    #[test]
3733    fn pending_route_permit_is_released_on_rejection_and_abort() {
3734        let forwarding = ForwardingTable::default();
3735        let module_connection = ConnectionId::new(300);
3736        let client = ConnectionId::new(301);
3737        let (module_tx, _module_rx) = mpsc::channel(1);
3738        let endpoint = forwarding
3739            .register_module_connection(
3740                module_connection,
3741                "permit".into(),
3742                2,
3743                Concurrency::ModuleManaged,
3744                FrameSink::new(module_tx),
3745            )
3746            .unwrap();
3747        let (client_tx, mut client_rx) = mpsc::channel(1);
3748        let sink = FrameSink::new(client_tx);
3749        let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
3750        assert!(sink.try_send(test_ping(999)).is_err());
3751        forwarding
3752            .complete_pending_relay(
3753                module_connection,
3754                rejected.corr,
3755                RouteBindRelayOutcome::Rejected(ErrorBody {
3756                    code: "no".into(),
3757                    message: "rejected".into(),
3758                    detail: None,
3759                }),
3760            )
3761            .unwrap();
3762        sink.try_send(test_ping(1000)).unwrap();
3763        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
3764
3765        let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
3766        assert!(sink.try_send(test_ping(1001)).is_err());
3767        forwarding
3768            .abort_pending_relay(
3769                endpoint,
3770                aborted.corr,
3771                RouteBindRelayOutcome::ModuleGone("abort".into()),
3772            )
3773            .unwrap();
3774        sink.try_send(test_ping(1002)).unwrap();
3775        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
3776
3777        let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
3778        forwarding
3779            .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
3780            .unwrap();
3781        drop(client_rx);
3782        let completion = forwarding
3783            .complete_pending_relay(
3784                module_connection,
3785                receiver_closed.corr,
3786                RouteBindRelayOutcome::Accepted,
3787            )
3788            .unwrap();
3789        assert!(completion.abandoned.is_some());
3790        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3791    }
3792
3793    /// Every connection teardown passes through `cleanup_connection`, and
3794    /// connection ids come from a monotonic counter that never hands an id out
3795    /// twice. If the closing mark survives teardown, the set grows by one entry
3796    /// per connection for the life of the daemon -- the self-watchdog alone
3797    /// reconnects once a minute.
3798    #[test]
3799    fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
3800        let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
3801            route_fixture("closing-set-leak");
3802
3803        const CONNECTIONS: u64 = 32;
3804        for index in 0..CONNECTIONS {
3805            let client = ConnectionId::new(1000 + index);
3806            let (client_tx, _client_rx) = mpsc::channel(8);
3807            let route = begin_test_route(
3808                &forwarding,
3809                client,
3810                FrameSink::new(client_tx),
3811                index + 1,
3812                "closing-set-leak",
3813            );
3814            forwarding
3815                .complete_pending_relay(
3816                    module_connection,
3817                    route.corr,
3818                    RouteBindRelayOutcome::Accepted,
3819                )
3820                .unwrap();
3821            forwarding.cleanup_connection(client).unwrap();
3822        }
3823        forwarding.cleanup_connection(module_connection).unwrap();
3824
3825        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3826    }
3827
3828    /// The closing mark exists to refuse new work for a connection that is on
3829    /// its way out but whose teardown has not run yet: the daemon asks the
3830    /// connection loop to end, and only when the loop reacts does
3831    /// `cleanup_connection` strip the connection's state. Inside that window an
3832    /// operation for the dying connection must still be refused; only after
3833    /// cleanup completes may the mark go.
3834    #[test]
3835    fn closing_connection_is_refused_new_work_until_cleanup_completes() {
3836        let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
3837            route_fixture("closing-gate");
3838
3839        // A published route: escalate_client_delivery_failure only marks a
3840        // connection closing for a route it has already published.
3841        let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
3842        forwarding
3843            .complete_pending_relay(
3844                module_connection,
3845                live.corr,
3846                RouteBindRelayOutcome::Accepted,
3847            )
3848            .unwrap();
3849        client_rx.try_recv().unwrap();
3850
3851        // Mark the connection closing through the production path without
3852        // running teardown, pinning the window open.
3853        assert!(forwarding
3854            .escalate_client_delivery_failure(
3855                client,
3856                live.client_channel,
3857                live.client_epoch,
3858                CloseReason::new(
3859                    "module_to_client_delivery_failed",
3860                    "client egress refused a module frame",
3861                ),
3862            )
3863            .unwrap());
3864        assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
3865
3866        // A late route.open for the closing client is refused ...
3867        assert!(matches!(
3868            forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
3869            Err(ForwardingError::ConnectionClosing { connection_id })
3870                if connection_id == client
3871        ));
3872        // ... and so is a late attempt to register the connection as a module.
3873        let (late_tx, _late_rx) = mpsc::channel(1);
3874        assert!(matches!(
3875            forwarding.register_module_connection(
3876                client,
3877                "late-module".into(),
3878                2,
3879                Concurrency::ModuleManaged,
3880                FrameSink::new(late_tx),
3881            ),
3882            Err(ForwardingError::ConnectionClosing { connection_id })
3883                if connection_id == client
3884        ));
3885
3886        // Teardown is the point that lifts the mark: it has just removed every
3887        // per-connection entry under the same lock, so the gate has nothing
3888        // left to protect for this id.
3889        forwarding.cleanup_connection(client).unwrap();
3890        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3891    }
3892}
3893
3894/// Blue/green swap slots: a candidate registered beside the active endpoint,
3895/// promoted by `cutover_candidate`, with the old incumbent drained by endpoint.
3896#[cfg(test)]
3897mod swap_slot_tests {
3898    use std::time::Duration;
3899
3900    use super::*;
3901    use tokio::sync::mpsc;
3902
3903    const MODULE_ID: &str = "swapped";
3904
3905    struct SwapFixture {
3906        forwarding: ForwardingTable,
3907        incumbent_connection: ConnectionId,
3908        incumbent: ModuleEndpointId,
3909        candidate_connection: ConnectionId,
3910        candidate: ModuleEndpointId,
3911        _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
3912    }
3913
3914    fn swap_fixture() -> SwapFixture {
3915        let forwarding = ForwardingTable::default();
3916        let incumbent_connection = ConnectionId::new(100);
3917        let candidate_connection = ConnectionId::new(110);
3918        let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
3919        let incumbent = forwarding
3920            .register_module_connection(
3921                incumbent_connection,
3922                MODULE_ID.to_string(),
3923                2,
3924                Concurrency::ModuleManaged,
3925                FrameSink::new(incumbent_tx),
3926            )
3927            .unwrap();
3928        let (candidate_tx, candidate_rx) = mpsc::channel(8);
3929        let candidate = forwarding
3930            .register_candidate_module_connection(
3931                candidate_connection,
3932                MODULE_ID.to_string(),
3933                2,
3934                Concurrency::ModuleManaged,
3935                FrameSink::new(candidate_tx),
3936            )
3937            .unwrap();
3938        SwapFixture {
3939            forwarding,
3940            incumbent_connection,
3941            incumbent,
3942            candidate_connection,
3943            candidate,
3944            _module_rxs: vec![incumbent_rx, candidate_rx],
3945        }
3946    }
3947
3948    fn client(
3949        raw: u64,
3950    ) -> (
3951        ConnectionId,
3952        FrameSink,
3953        mpsc::Receiver<crate::router::OutboundFrame>,
3954    ) {
3955        let (tx, rx) = mpsc::channel(8);
3956        (ConnectionId::new(raw), FrameSink::new(tx), rx)
3957    }
3958
3959    fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
3960        forwarding
3961            .read_inner()
3962            .unwrap()
3963            .client_to_module
3964            .values()
3965            .map(|route| route.module_endpoint)
3966            .collect()
3967    }
3968
3969    #[test]
3970    fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
3971        let fixture = swap_fixture();
3972        let forwarding = &fixture.forwarding;
3973        assert_ne!(fixture.incumbent, fixture.candidate);
3974
3975        // Every by-id consumer still resolves the incumbent.
3976        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
3977        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
3978        let (client_connection, client_sink, _client_rx) = client(200);
3979        let pending = forwarding
3980            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
3981            .unwrap();
3982        assert_eq!(pending.endpoint, fixture.incumbent);
3983        let rpc = forwarding
3984            .begin_module_control_rpc_for(
3985                MODULE_ID,
3986                "health.check",
3987                Instant::now() + Duration::from_secs(1),
3988            )
3989            .unwrap();
3990        assert_eq!(rpc.endpoint, fixture.incumbent);
3991        let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
3992        assert_eq!(census.len(), 1, "the census lists one endpoint per id");
3993
3994        // Connection-keyed lookups see the candidate, so its own frames resolve.
3995        assert_eq!(
3996            forwarding
3997                .module_endpoint_for_connection(fixture.candidate_connection)
3998                .unwrap(),
3999            Some(fixture.candidate)
4000        );
4001        assert_eq!(
4002            forwarding
4003                .module_id_for_connection(fixture.candidate_connection)
4004                .unwrap()
4005                .as_deref(),
4006            Some(MODULE_ID)
4007        );
4008
4009        // One candidate per id.
4010        let (other_tx, _other_rx) = mpsc::channel(1);
4011        assert_eq!(
4012            forwarding.register_candidate_module_connection(
4013                ConnectionId::new(120),
4014                MODULE_ID.to_string(),
4015                2,
4016                Concurrency::ModuleManaged,
4017                FrameSink::new(other_tx),
4018            ),
4019            Err(ForwardingError::CandidateSlotOccupied {
4020                module_id: MODULE_ID.to_string()
4021            })
4022        );
4023    }
4024
4025    /// The cutover linearization point. A relay reserved on the incumbent before
4026    /// cutover must never become a route on the incumbent, and every relay
4027    /// reserved after it must land on the promoted candidate.
4028    #[test]
4029    fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
4030        let fixture = swap_fixture();
4031        let forwarding = &fixture.forwarding;
4032        let (early_client, early_sink, _early_rx) = client(200);
4033        let mut early = forwarding
4034            .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
4035            .unwrap();
4036        assert_eq!(early.endpoint, fixture.incumbent);
4037        assert!(forwarding
4038            .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
4039            .unwrap());
4040
4041        let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4042        assert_eq!(
4043            cutover,
4044            ForwardingCutover {
4045                promoted: fixture.candidate,
4046                incumbent: Some(fixture.incumbent),
4047            }
4048        );
4049
4050        // A route.open reserved after cutover goes to the promoted candidate.
4051        let (late_client, late_sink, _late_rx) = client(201);
4052        let late = forwarding
4053            .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
4054            .unwrap();
4055        assert_eq!(
4056            late.endpoint, fixture.candidate,
4057            "a route.open after cutover was reserved on the incumbent"
4058        );
4059
4060        // The incumbent acks the early relay after cutover.
4061        let completion = forwarding
4062            .complete_pending_relay(
4063                fixture.incumbent_connection,
4064                early.corr,
4065                RouteBindRelayOutcome::Accepted,
4066            )
4067            .expect("a superseded endpoint's ack is not an error on its connection");
4068        assert!(completion.settled);
4069        assert!(
4070            !committed_endpoints(forwarding).contains(&fixture.incumbent),
4071            "a relay reserved before cutover committed a route on the incumbent"
4072        );
4073        let goodbye = completion
4074            .abandoned
4075            .expect("the incumbent is told to drop the binding it just created");
4076        assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
4077        assert_eq!(goodbye.channel, early.module_channel);
4078        assert_eq!(goodbye.epoch, early.module_epoch);
4079        assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
4080        match early.receiver.try_recv() {
4081            Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
4082            other => panic!("expected a retryable module_reloading answer, got {other:?}"),
4083        }
4084        assert!(matches!(
4085            forwarding
4086                .lookup_data_route(early_client, early.client_channel, early.client_epoch)
4087                .unwrap(),
4088            DataRoute::Client(DataRouteState::Absent)
4089        ));
4090
4091        // The late relay commits on the candidate; only its pair stays reserved
4092        // until then.
4093        assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4094        forwarding
4095            .complete_pending_relay(
4096                fixture.candidate_connection,
4097                late.corr,
4098                RouteBindRelayOutcome::Accepted,
4099            )
4100            .unwrap();
4101        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4102        assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
4103    }
4104
4105    #[test]
4106    fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
4107        let fixture = swap_fixture();
4108        let forwarding = &fixture.forwarding;
4109        // One bound route and one in-flight relay on the incumbent.
4110        let (bound_client, bound_sink, _bound_rx) = client(200);
4111        let bound = forwarding
4112            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4113            .unwrap();
4114        forwarding
4115            .complete_pending_relay(
4116                fixture.incumbent_connection,
4117                bound.corr,
4118                RouteBindRelayOutcome::Accepted,
4119            )
4120            .unwrap();
4121        let (pending_client, pending_sink, _pending_rx) = client(201);
4122        let mut in_flight = forwarding
4123            .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
4124            .unwrap();
4125        forwarding
4126            .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
4127            .unwrap();
4128
4129        let incumbent = forwarding
4130            .cutover_candidate(MODULE_ID)
4131            .unwrap()
4132            .unwrap()
4133            .incumbent
4134            .unwrap();
4135        let target = forwarding
4136            .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
4137            .unwrap()
4138            .expect("the superseded incumbent is still registered");
4139
4140        assert_eq!(target.endpoint, fixture.incumbent);
4141        assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
4142        assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
4143        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4144        assert_eq!(target.abandoned_bindings.len(), 1);
4145        assert_eq!(
4146            target.abandoned_bindings[0].channel,
4147            in_flight.module_channel
4148        );
4149        assert!(matches!(
4150            in_flight.receiver.try_recv(),
4151            Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
4152        ));
4153        assert_eq!(
4154            forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
4155            1,
4156            "the incumbent's bound route stays until its drain finishes"
4157        );
4158
4159        let (next_client, next_sink, _next_rx) = client(202);
4160        let next = forwarding
4161            .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
4162            .expect("the promoted candidate keeps accepting routes");
4163        assert_eq!(next.endpoint, fixture.candidate);
4164    }
4165
4166    /// An endpoint replaced WITHOUT a promotion (a successor registered over it
4167    /// as an ordinary active HELLO) is stale, not superseded, and its ack still
4168    /// fails the acking connection exactly as it did before swap slots existed:
4169    /// `StaleModuleEndpoint`, the reservation indexes already stripped, and the
4170    /// waiting client's sender dropped unanswered.
4171    #[test]
4172    fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
4173        let forwarding = ForwardingTable::default();
4174        let first_connection = ConnectionId::new(70);
4175        let (first_tx, _first_rx) = mpsc::channel(8);
4176        forwarding
4177            .register_module_connection(
4178                first_connection,
4179                MODULE_ID.to_string(),
4180                2,
4181                Concurrency::ModuleManaged,
4182                FrameSink::new(first_tx),
4183            )
4184            .unwrap();
4185        let (client_connection, client_sink, _client_rx) = client(200);
4186        let mut pending = forwarding
4187            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4188            .unwrap();
4189        let (second_tx, _second_rx) = mpsc::channel(8);
4190        forwarding
4191            .register_module_connection(
4192                ConnectionId::new(80),
4193                MODULE_ID.to_string(),
4194                2,
4195                Concurrency::ModuleManaged,
4196                FrameSink::new(second_tx),
4197            )
4198            .unwrap();
4199
4200        assert_eq!(
4201            forwarding
4202                .complete_pending_relay(
4203                    first_connection,
4204                    pending.corr,
4205                    RouteBindRelayOutcome::Accepted
4206                )
4207                .unwrap_err(),
4208            ForwardingError::StaleModuleEndpoint
4209        );
4210        assert!(committed_endpoints(&forwarding).is_empty());
4211        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4212        assert!(matches!(
4213            pending.receiver.try_recv(),
4214            Err(oneshot::error::TryRecvError::Closed)
4215        ));
4216    }
4217
4218    #[test]
4219    fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
4220        // A candidate whose connection drops leaves the incumbent routable.
4221        let fixture = swap_fixture();
4222        let forwarding = &fixture.forwarding;
4223        assert!(forwarding
4224            .cleanup_connection(fixture.candidate_connection)
4225            .unwrap()
4226            .is_empty());
4227        assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
4228        let (client_connection, client_sink, _client_rx) = client(200);
4229        assert_eq!(
4230            forwarding
4231                .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4232                .unwrap()
4233                .endpoint,
4234            fixture.incumbent
4235        );
4236
4237        // After a cutover, the incumbent's teardown releases its own routes and
4238        // leaves the promoted candidate in place.
4239        let fixture = swap_fixture();
4240        let forwarding = &fixture.forwarding;
4241        let (bound_client, bound_sink, _bound_rx) = client(200);
4242        let bound = forwarding
4243            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4244            .unwrap();
4245        forwarding
4246            .complete_pending_relay(
4247                fixture.incumbent_connection,
4248                bound.corr,
4249                RouteBindRelayOutcome::Accepted,
4250            )
4251            .unwrap();
4252        forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4253        let released = forwarding
4254            .cleanup_connection(fixture.incumbent_connection)
4255            .unwrap();
4256        assert_eq!(released.len(), 1);
4257        assert_eq!(released[0].connection_id, bound_client);
4258        assert!(forwarding
4259            .read_inner()
4260            .unwrap()
4261            .superseded_endpoints
4262            .is_empty());
4263        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4264        let (next_client, next_sink, _next_rx) = client(201);
4265        assert_eq!(
4266            forwarding
4267                .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
4268                .unwrap()
4269                .endpoint,
4270            fixture.candidate
4271        );
4272    }
4273}