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. 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    /// The held requests themselves, as the module saw them, so its own log
1765    /// can say what each one was: `(module channel, corr)`, ordered by channel
1766    /// then corr, at most [`DRAIN_HELD_REQUESTS_LISTED`]. The daemon never reads
1767    /// request bodies, so it cannot name a request's method; the module can,
1768    /// from the channel and corr. A request is released only when the module
1769    /// sends a terminal frame (Response, Error or StreamEnd) with its corr on
1770    /// its route, so every pair listed here is one the module ended, if at all,
1771    /// without sending that frame.
1772    pub(crate) held: Vec<(u16, u64)>,
1773}
1774
1775/// How many held requests the drain-timeout line lists by channel and corr.
1776pub(crate) const DRAIN_HELD_REQUESTS_LISTED: usize = 32;
1777
1778impl ForwardingTable {
1779    /// Summarise the requests one endpoint's drain is still waiting on.
1780    pub(crate) fn endpoint_drain_holdouts(
1781        &self,
1782        endpoint: ModuleEndpointId,
1783    ) -> Result<DrainHoldouts, ForwardingError> {
1784        let inner = self.read_inner()?;
1785        let mut holdouts = DrainHoldouts::default();
1786        let mut by_connection: HashMap<u64, usize> = HashMap::new();
1787        for (key, route) in &inner.client_to_module {
1788            if route.module_endpoint != endpoint {
1789                continue;
1790            }
1791            holdouts.total_routes += 1;
1792            let held = route.flow.drain_in_flight();
1793            if held == 0 {
1794                continue;
1795            }
1796            holdouts.requests += held;
1797            holdouts.routes += 1;
1798            *by_connection.entry(key.connection_id.get()).or_default() += held;
1799            holdouts.held.extend(
1800                route
1801                    .flow
1802                    .drain_held_corrs()
1803                    .into_iter()
1804                    .map(|corr| (route.module_channel, corr)),
1805            );
1806        }
1807        holdouts.held.sort_unstable();
1808        holdouts.held.truncate(DRAIN_HELD_REQUESTS_LISTED);
1809        let mut connections = by_connection.into_iter().collect::<Vec<_>>();
1810        connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
1811        connections.truncate(3);
1812        holdouts.top_connections = connections;
1813        Ok(holdouts)
1814    }
1815
1816    pub(crate) fn endpoint_in_flight_count(
1817        &self,
1818        endpoint: ModuleEndpointId,
1819    ) -> Result<usize, ForwardingError> {
1820        let inner = self.read_inner()?;
1821        Ok(inner
1822            .client_to_module
1823            .values()
1824            .filter(|route| route.module_endpoint == endpoint)
1825            .map(|route| route.flow.drain_in_flight())
1826            .sum())
1827    }
1828
1829    pub(crate) fn endpoint_is_draining(
1830        &self,
1831        endpoint: ModuleEndpointId,
1832    ) -> Result<bool, ForwardingError> {
1833        Ok(self
1834            .read_inner()?
1835            .draining_endpoints
1836            .contains_key(&endpoint))
1837    }
1838
1839    pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
1840        let inner = self.read_inner()?;
1841        Ok(inner
1842            .modules_by_id
1843            .get(module_id)
1844            .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
1845    }
1846
1847    pub(crate) fn release_module_endpoint_routes(
1848        &self,
1849        endpoint: ModuleEndpointId,
1850    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1851        let mut inner = self.write_inner()?;
1852        let routes = inner
1853            .module_to_client
1854            .iter()
1855            .filter(|(module_key, _)| module_key.endpoint == endpoint)
1856            .map(|(module_key, route)| (*module_key, route.module_epoch))
1857            .collect::<Vec<_>>();
1858        let mut released = Vec::with_capacity(routes.len());
1859        for (module_key, epoch) in routes {
1860            if let RouteRelease::Removed(target) =
1861                release_module_route_locked(&mut inner, module_key, epoch)
1862            {
1863                released.push(target);
1864            }
1865        }
1866        Ok(released)
1867    }
1868
1869    /// Enumerate one endpoint's current routes without contacting the module.
1870    ///
1871    /// The read lock makes this safe while the endpoint drains: the returned
1872    /// `draining` marker describes the same table state that owns the route,
1873    /// rather than inferring liveness from a module that may be stopping.
1874    pub(crate) fn endpoint_routes(
1875        &self,
1876        endpoint: ModuleEndpointId,
1877    ) -> Result<Vec<EndpointRoute>, ForwardingError> {
1878        let inner = self.read_inner()?;
1879        Ok(endpoint_routes_locked(&inner, endpoint))
1880    }
1881
1882    /// Snapshot all live endpoint route sets under one forwarding-table read lock.
1883    pub(crate) fn route_census(
1884        &self,
1885        module_id: Option<&str>,
1886    ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
1887        let inner = self.read_inner()?;
1888        let mut endpoints = inner
1889            .modules_by_id
1890            .iter()
1891            .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
1892            .map(|(id, module)| (id.clone(), module.endpoint))
1893            .collect::<Vec<_>>();
1894        endpoints.sort_by(|left, right| left.0.cmp(&right.0));
1895        Ok(endpoints
1896            .into_iter()
1897            .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
1898            .collect())
1899    }
1900
1901    /// Snapshot only the endpoint currently routable under this module id.
1902    pub(crate) fn live_roots(
1903        &self,
1904        module_id: &str,
1905    ) -> Result<ModuleControlResponseToModule, ForwardingError> {
1906        let inner = self.read_inner()?;
1907        let endpoint = inner
1908            .modules_by_id
1909            .get(module_id)
1910            .map(|module| module.endpoint);
1911        let mut roots = BTreeMap::new();
1912        let mut unknown_root_bindings = 0;
1913        let mut total_bindings = 0;
1914        if let Some(endpoint) = endpoint {
1915            for binding in inner
1916                .module_to_client
1917                .values()
1918                .filter(|binding| binding.module_endpoint == endpoint)
1919            {
1920                total_bindings += 1;
1921                if let Some(root) = &binding.project_root {
1922                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1923                    entry.0 += 1;
1924                } else {
1925                    unknown_root_bindings += 1;
1926                }
1927            }
1928            for pending in inner
1929                .pending_relays
1930                .values()
1931                .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
1932            {
1933                total_bindings += 1;
1934                if let Some(root) = &pending.reservation.project_root {
1935                    let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1936                    entry.1 += 1;
1937                } else {
1938                    unknown_root_bindings += 1;
1939                }
1940            }
1941        }
1942        Ok(ModuleControlResponseToModule::LiveRoots {
1943            roots: roots
1944                .into_iter()
1945                .map(|(project_root, (bound, pending))| LiveRoot {
1946                    project_root,
1947                    bound,
1948                    pending,
1949                })
1950                .collect(),
1951            unknown_root_bindings,
1952            total_bindings,
1953        })
1954    }
1955
1956    /// True if this connection already owns committed or reserved CLIENT routes.
1957    /// A module registers (HELLO) before serving and never opens client routes, so
1958    /// a connection that has client routes must not also become a module endpoint
1959    /// — otherwise one connection holds both client and module state and cleanup
1960    /// only releases one side.
1961    pub(crate) fn connection_has_client_routes(
1962        &self,
1963        connection_id: ConnectionId,
1964    ) -> Result<bool, ForwardingError> {
1965        let inner = self.read_inner()?;
1966        let has = inner
1967            .client_to_module
1968            .keys()
1969            .any(|key| key.connection_id == connection_id)
1970            || inner
1971                .reserved_client
1972                .keys()
1973                .any(|key| key.connection_id == connection_id);
1974        Ok(has)
1975    }
1976
1977    pub(crate) fn cleanup_connection(
1978        &self,
1979        connection_id: ConnectionId,
1980    ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1981        let mut inner = self.write_inner()?;
1982        inner.closing_connections.insert(connection_id);
1983        let released = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
1984            remove_module_connection_locked(&mut inner, endpoint)
1985        } else {
1986            Self::cleanup_client_connection_locked(&mut inner, connection_id)
1987        };
1988        // The closing mark refuses new work for a connection whose teardown is
1989        // still pending. This is the latest point at which lifting it is safe:
1990        // teardown has just removed every per-connection entry above, under
1991        // this same write lock, so no lookup can still find live state for the
1992        // id; and ids come from a monotonic counter that never reissues one,
1993        // so the id can never name a future connection either. Keeping the
1994        // mark past this point only grew the set by one entry per connection
1995        // for the life of the daemon.
1996        inner.closing_connections.remove(&connection_id);
1997        Ok(released)
1998    }
1999
2000    fn cleanup_client_connection_locked(
2001        inner: &mut ForwardingInner,
2002        connection_id: ConnectionId,
2003    ) -> Vec<GoodbyeTarget> {
2004        let routes = inner
2005            .client_to_module
2006            .iter()
2007            .filter(|(key, _)| key.connection_id == connection_id)
2008            .map(|(key, route)| (*key, route.client_epoch))
2009            .collect::<Vec<_>>();
2010        let mut released = Vec::with_capacity(routes.len());
2011        for (client_key, epoch) in routes {
2012            if let RouteRelease::Removed(target) =
2013                release_client_route_locked(inner, client_key, epoch)
2014            {
2015                released.push(target);
2016            }
2017        }
2018
2019        let pending_keys = inner
2020            .pending_relays
2021            .iter()
2022            .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2023            .map(|(key, _)| *key)
2024            .collect::<Vec<_>>();
2025        for key in pending_keys {
2026            let Some(pending) = inner.pending_relays.remove(&key) else {
2027                continue;
2028            };
2029            release_reserved_route_locked(
2030                inner,
2031                pending.reservation.client_key,
2032                pending.reservation.module_key,
2033            );
2034            if pending.relay_enqueued {
2035                if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2036                    released.push(target);
2037                }
2038            }
2039            let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2040                "client connection closed during route.bind relay".to_string(),
2041            ));
2042        }
2043
2044        let orphaned = inner
2045            .reserved_client
2046            .iter()
2047            .filter(|(key, _)| key.connection_id == connection_id)
2048            .map(|(client, module)| (*client, *module))
2049            .collect::<Vec<_>>();
2050        for (client_key, module_key) in orphaned {
2051            release_reserved_route_locked(inner, client_key, module_key);
2052        }
2053        inner.next_client_channel.remove(&connection_id);
2054        inner
2055            .client_slot_epochs
2056            .retain(|key, _| key.connection_id != connection_id);
2057        inner
2058            .last_published_epoch
2059            .retain(|key, _| key.connection_id != connection_id);
2060        inner
2061            .status
2062            .retain(|(key, _), _| key.connection_id != connection_id);
2063
2064        released
2065    }
2066
2067    pub(crate) fn escalate_client_delivery_failure(
2068        &self,
2069        connection_id: ConnectionId,
2070        channel: u16,
2071        expected_epoch: u32,
2072        reason: CloseReason,
2073    ) -> Result<bool, ForwardingError> {
2074        let should_close = {
2075            let mut inner = self.write_inner()?;
2076            let key = ClientRouteKey {
2077                connection_id,
2078                channel,
2079            };
2080            if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2081                false
2082            } else {
2083                inner.closing_connections.insert(connection_id);
2084                true
2085            }
2086        };
2087        if should_close {
2088            self.request_connection_close(connection_id, reason);
2089        }
2090        Ok(should_close)
2091    }
2092
2093    fn record_route_release(&self, release: &RouteRelease) {
2094        match release {
2095            RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2096            RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2097            RouteRelease::Absent => {}
2098        }
2099    }
2100
2101    fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2102        self.inner.read().map_err(|_| ForwardingError::Poisoned)
2103    }
2104
2105    fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2106        self.inner.write().map_err(|_| ForwardingError::Poisoned)
2107    }
2108
2109    fn lock_close_registry(
2110        &self,
2111    ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2112        self.close_registry
2113            .lock()
2114            .unwrap_or_else(|poisoned| poisoned.into_inner())
2115    }
2116}
2117
2118impl ForwardingInner {
2119    fn allocate_route_slots(
2120        &mut self,
2121        connection_id: ConnectionId,
2122        endpoint: ModuleEndpointId,
2123    ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2124        let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2125        let mut client_channel = client_start;
2126        let client_channel = loop {
2127            let key = ClientRouteKey {
2128                connection_id,
2129                channel: client_channel,
2130            };
2131            let eligible = !self.client_to_module.contains_key(&key)
2132                && !self.reserved_client.contains_key(&key)
2133                && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2134            if eligible {
2135                break client_channel;
2136            }
2137            client_channel = next_channel(client_channel);
2138            if client_channel == client_start {
2139                return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2140            }
2141        };
2142
2143        let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2144        let mut module_channel = module_start;
2145        let module_channel = loop {
2146            let key = ModuleRouteKey {
2147                endpoint,
2148                channel: module_channel,
2149            };
2150            let eligible = !self.module_to_client.contains_key(&key)
2151                && !self.reserved_module.contains_key(&key)
2152                && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2153            if eligible {
2154                break module_channel;
2155            }
2156            module_channel = next_channel(module_channel);
2157            if module_channel == module_start {
2158                return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2159            }
2160        };
2161
2162        let client_key = ClientRouteKey {
2163            connection_id,
2164            channel: client_channel,
2165        };
2166        let module_key = ModuleRouteKey {
2167            endpoint,
2168            channel: module_channel,
2169        };
2170        let client_epoch = self
2171            .client_slot_epochs
2172            .get(&client_key)
2173            .copied()
2174            .unwrap_or(0)
2175            + 1;
2176        let module_epoch = self
2177            .module_slot_epochs
2178            .get(&module_key)
2179            .copied()
2180            .unwrap_or(0)
2181            + 1;
2182        self.client_slot_epochs.insert(client_key, client_epoch);
2183        self.module_slot_epochs.insert(module_key, module_epoch);
2184        self.next_client_channel
2185            .insert(connection_id, next_channel(client_channel));
2186        self.next_module_channel
2187            .insert(endpoint, next_channel(module_channel));
2188        Ok((client_channel, client_epoch, module_channel, module_epoch))
2189    }
2190
2191    fn allocate_control_corr(
2192        &mut self,
2193        endpoint: ModuleEndpointId,
2194    ) -> Result<u64, ForwardingError> {
2195        let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2196        if candidate == 0 {
2197            self.closing_connections.insert(endpoint.connection_id);
2198            return Err(ForwardingError::RelayCorrelationExhausted);
2199        }
2200        self.next_control_corr.insert(
2201            endpoint,
2202            if candidate == u64::MAX {
2203                0
2204            } else {
2205                candidate + 1
2206            },
2207        );
2208        Ok(candidate)
2209    }
2210}
2211
2212fn next_channel(channel: u16) -> u16 {
2213    let next = channel.wrapping_add(1);
2214    if next == 0 {
2215        1
2216    } else {
2217        next
2218    }
2219}
2220
2221fn endpoint_routes_locked(
2222    inner: &ForwardingInner,
2223    endpoint: ModuleEndpointId,
2224) -> Vec<EndpointRoute> {
2225    let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2226    let draining = drain_reason.is_some();
2227    let mut routes = inner
2228        .module_to_client
2229        .iter()
2230        .filter(|(module_key, _)| module_key.endpoint == endpoint)
2231        .map(|(_, route)| EndpointRoute {
2232            goodbye_target: GoodbyeTarget {
2233                connection_id: route.client_connection_id,
2234                sink: route.client_sink.clone(),
2235                negotiated_ver: route.client_negotiated_ver,
2236                channel: route.client_channel,
2237                epoch: route.client_epoch,
2238                kind: GoodbyeTargetKind::Client,
2239                module_id: None,
2240            },
2241            principal: route.principal.clone(),
2242            bound_at: route.bound_at,
2243            draining,
2244            drain_reason,
2245        })
2246        .collect::<Vec<_>>();
2247    routes.sort_by_key(|route| {
2248        (
2249            route.goodbye_target.connection_id.get(),
2250            route.goodbye_target.channel,
2251            route.goodbye_target.epoch,
2252        )
2253    });
2254    routes
2255}
2256
2257fn release_reserved_route_locked(
2258    inner: &mut ForwardingInner,
2259    client_key: ClientRouteKey,
2260    module_key: ModuleRouteKey,
2261) {
2262    if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
2263        inner.reserved_client.remove(&client_key);
2264    }
2265    if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
2266        inner.reserved_module.remove(&module_key);
2267    }
2268    inner.status.retain(|(key, _), _| *key != client_key);
2269}
2270
2271fn release_client_route_locked(
2272    inner: &mut ForwardingInner,
2273    client_key: ClientRouteKey,
2274    expected_epoch: u32,
2275) -> RouteRelease {
2276    let Some(route) = inner.client_to_module.get(&client_key) else {
2277        return RouteRelease::Absent;
2278    };
2279    if route.client_epoch != expected_epoch {
2280        return RouteRelease::Stale;
2281    }
2282    let route = inner
2283        .client_to_module
2284        .remove(&client_key)
2285        .expect("route checked under the same forwarding lock");
2286    route.flow.close();
2287    inner.module_to_client.remove(&ModuleRouteKey {
2288        endpoint: route.module_endpoint,
2289        channel: route.module_channel,
2290    });
2291    inner.status.remove(&(client_key, expected_epoch));
2292    RouteRelease::Removed(GoodbyeTarget {
2293        connection_id: route.module_endpoint.connection_id,
2294        sink: route.module_sink.clone(),
2295        negotiated_ver: route.module_negotiated_ver,
2296        channel: route.module_channel,
2297        epoch: route.module_epoch,
2298        kind: GoodbyeTargetKind::Module,
2299        module_id: Some(route.module_id.clone()),
2300    })
2301}
2302
2303fn release_module_route_locked(
2304    inner: &mut ForwardingInner,
2305    module_key: ModuleRouteKey,
2306    expected_epoch: u32,
2307) -> RouteRelease {
2308    let Some(route) = inner.module_to_client.get(&module_key) else {
2309        return RouteRelease::Absent;
2310    };
2311    if route.module_epoch != expected_epoch {
2312        return RouteRelease::Stale;
2313    }
2314    let route = inner
2315        .module_to_client
2316        .remove(&module_key)
2317        .expect("route checked under the same forwarding lock");
2318    route.flow.close();
2319    let client_key = ClientRouteKey {
2320        connection_id: route.client_connection_id,
2321        channel: route.client_channel,
2322    };
2323    inner.client_to_module.remove(&client_key);
2324    inner.status.remove(&(client_key, route.client_epoch));
2325    RouteRelease::Removed(GoodbyeTarget {
2326        connection_id: route.client_connection_id,
2327        sink: route.client_sink.clone(),
2328        negotiated_ver: route.client_negotiated_ver,
2329        channel: route.client_channel,
2330        epoch: route.client_epoch,
2331        kind: GoodbyeTargetKind::Client,
2332        module_id: None,
2333    })
2334}
2335
2336fn commit_route_locked(
2337    inner: &mut ForwardingInner,
2338    pending: PendingRouteBindRelayEntry,
2339) -> Result<Option<GoodbyeTarget>, ForwardingError> {
2340    let reservation = pending.reservation;
2341    if inner
2342        .closing_connections
2343        .contains(&reservation.client_key.connection_id)
2344    {
2345        return Err(ForwardingError::ConnectionClosing {
2346            connection_id: reservation.client_key.connection_id,
2347        });
2348    }
2349    let module_id = inner
2350        .module_id_by_endpoint
2351        .get(&reservation.module_key.endpoint)
2352        .cloned()
2353        .ok_or(ForwardingError::StaleModuleEndpoint)?;
2354    if inner
2355        .draining_endpoints
2356        .contains_key(&reservation.module_key.endpoint)
2357    {
2358        return Err(ForwardingError::ModuleReloading { module_id });
2359    }
2360    if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
2361        || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
2362    {
2363        return Err(ForwardingError::UnknownReservation {
2364            client_channel: reservation.client_key.channel,
2365            module_channel: reservation.module_key.channel,
2366        });
2367    }
2368    let module = inner
2369        .modules_by_id
2370        .get(&module_id)
2371        .filter(|module| module.endpoint == reservation.module_key.endpoint)
2372        .cloned()
2373        .ok_or(ForwardingError::StaleModuleEndpoint)?;
2374    let binding = Arc::new(RouteBinding {
2375        client_connection_id: reservation.client_key.connection_id,
2376        client_sink: pending.client_sink,
2377        client_negotiated_ver: pending.client_negotiated_ver,
2378        client_channel: reservation.client_key.channel,
2379        client_epoch: reservation.client_epoch,
2380        module_id,
2381        module_endpoint: reservation.module_key.endpoint,
2382        module_sink: module.sink,
2383        module_negotiated_ver: module.negotiated_ver,
2384        module_channel: reservation.module_key.channel,
2385        module_epoch: reservation.module_epoch,
2386        principal: pending.principal,
2387        project_root: reservation.project_root.clone(),
2388        bound_at: Instant::now(),
2389        flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
2390    });
2391    inner
2392        .client_to_module
2393        .insert(reservation.client_key, Arc::clone(&binding));
2394    inner
2395        .module_to_client
2396        .insert(reservation.module_key, binding);
2397    let previous_published = inner
2398        .last_published_epoch
2399        .insert(reservation.client_key, reservation.client_epoch);
2400
2401    // OwnedPermit::send cannot fail, but its returned sender reveals a receiver
2402    // that closed after reservation and before this locked publication point.
2403    // Stamp at publication: the permit was reserved earlier, but queue residency
2404    // for the reply-write diagnosis starts when the frame actually enters the queue.
2405    let client_sender = pending.client_permit.send(crate::router::OutboundFrame {
2406        frame: pending.route_open_frame,
2407        enqueued_at: std::time::Instant::now(),
2408        flushed: None,
2409    });
2410    if client_sender.is_closed() {
2411        let abandoned = pending
2412            .relay_enqueued
2413            .then(|| abandoned_route_target(inner, &reservation))
2414            .flatten();
2415        if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
2416            route.flow.close();
2417        }
2418        inner.module_to_client.remove(&reservation.module_key);
2419        inner
2420            .status
2421            .remove(&(reservation.client_key, reservation.client_epoch));
2422        match previous_published {
2423            Some(epoch) => {
2424                inner
2425                    .last_published_epoch
2426                    .insert(reservation.client_key, epoch);
2427            }
2428            None => {
2429                inner.last_published_epoch.remove(&reservation.client_key);
2430            }
2431        }
2432        let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2433            "client egress closed during route publication".to_string(),
2434        ));
2435        return Ok(abandoned);
2436    }
2437
2438    let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
2439    Ok(None)
2440}
2441
2442/// The live module connection registered under exactly `endpoint`, whichever
2443/// slot holds it: active, swap candidate, or superseded incumbent.
2444///
2445/// Resolving by endpoint rather than by module id is what keeps an incumbent
2446/// addressable after a swap promoted a candidate over its id. An endpoint that
2447/// is in none of the three (a stale one, replaced without a promotion) resolves
2448/// to nothing, as it always has.
2449fn module_connection_for_endpoint_locked(
2450    inner: &ForwardingInner,
2451    endpoint: ModuleEndpointId,
2452) -> Option<&ModuleConnection> {
2453    let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
2454    inner
2455        .modules_by_id
2456        .get(module_id)
2457        .filter(|module| module.endpoint == endpoint)
2458        .or_else(|| {
2459            inner
2460                .candidates_by_id
2461                .get(module_id)
2462                .filter(|module| module.endpoint == endpoint)
2463        })
2464        .or_else(|| inner.superseded_endpoints.get(&endpoint))
2465}
2466
2467fn abandoned_route_target(
2468    inner: &ForwardingInner,
2469    reservation: &RouteReservation,
2470) -> Option<GoodbyeTarget> {
2471    let module_id = inner
2472        .module_id_by_endpoint
2473        .get(&reservation.module_key.endpoint)?;
2474    let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
2475    (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
2476        connection_id: module.endpoint.connection_id,
2477        sink: module.sink.clone(),
2478        negotiated_ver: module.negotiated_ver,
2479        channel: reservation.module_key.channel,
2480        epoch: reservation.module_epoch,
2481        kind: GoodbyeTargetKind::Module,
2482        module_id: Some(module_id.clone()),
2483    })
2484}
2485
2486fn remove_module_connection_locked(
2487    inner: &mut ForwardingInner,
2488    endpoint: ModuleEndpointId,
2489) -> Vec<GoodbyeTarget> {
2490    inner.draining_endpoints.remove(&endpoint);
2491    let module_id = inner.module_id_by_endpoint.remove(&endpoint);
2492    if let Some(module_id) = module_id.as_ref() {
2493        if inner
2494            .modules_by_id
2495            .get(module_id)
2496            .is_some_and(|module| module.endpoint == endpoint)
2497        {
2498            inner.modules_by_id.remove(module_id);
2499        }
2500        if inner
2501            .candidates_by_id
2502            .get(module_id)
2503            .is_some_and(|module| module.endpoint == endpoint)
2504        {
2505            inner.candidates_by_id.remove(module_id);
2506        }
2507    }
2508    inner.superseded_endpoints.remove(&endpoint);
2509    inner.endpoint_by_connection.remove(&endpoint.connection_id);
2510    inner.next_module_channel.remove(&endpoint);
2511    inner.next_control_corr.remove(&endpoint);
2512    inner
2513        .health_probe_tombstones
2514        .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
2515    inner
2516        .module_slot_epochs
2517        .retain(|key, _| key.endpoint != endpoint);
2518    let reserved_module_keys: Vec<ModuleRouteKey> = inner
2519        .reserved_module
2520        .keys()
2521        .filter(|module_key| module_key.endpoint == endpoint)
2522        .copied()
2523        .collect();
2524    for module_key in reserved_module_keys {
2525        if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2526            release_reserved_route_locked(inner, client_key, module_key);
2527        }
2528    }
2529
2530    let pending_keys: Vec<_> = inner
2531        .pending_relays
2532        .keys()
2533        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2534        .copied()
2535        .collect();
2536    let pending: Vec<_> = pending_keys
2537        .into_iter()
2538        .filter_map(|key| inner.pending_relays.remove(&key))
2539        .collect();
2540    for pending in pending {
2541        let module_label = module_id.as_deref().unwrap_or("unknown");
2542        let _ = pending
2543            .sender
2544            .send(RouteBindRelayOutcome::ModuleGone(format!(
2545                "module '{module_label}' connection closed during route.bind relay"
2546            )));
2547    }
2548
2549    let pending_control_keys: Vec<_> = inner
2550        .pending_control_rpcs
2551        .keys()
2552        .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2553        .copied()
2554        .collect();
2555    let pending_control: Vec<_> = pending_control_keys
2556        .into_iter()
2557        .filter_map(|key| inner.pending_control_rpcs.remove(&key))
2558        .collect();
2559    for pending in pending_control {
2560        let module_label = module_id.as_deref().unwrap_or("unknown");
2561        let _ = pending
2562            .sender
2563            .send(ModuleControlRpcOutcome::ModuleGone(format!(
2564                "module '{module_label}' connection closed during module-control RPC"
2565            )));
2566    }
2567
2568    let module_routes = inner
2569        .module_to_client
2570        .iter()
2571        .filter(|(module_key, _)| module_key.endpoint == endpoint)
2572        .map(|(module_key, route)| (*module_key, route.module_epoch))
2573        .collect::<Vec<_>>();
2574    let mut released = Vec::with_capacity(module_routes.len());
2575    for (module_key, epoch) in module_routes {
2576        if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
2577        {
2578            released.push(target);
2579        }
2580    }
2581    released
2582}
2583
2584#[derive(Debug, Clone, Copy)]
2585struct RequestCredit {
2586    subscription: bool,
2587    excluded_from_drain: bool,
2588}
2589
2590#[derive(Debug, Default)]
2591struct CreditLedger {
2592    by_corr: HashMap<u64, Vec<RequestCredit>>,
2593}
2594
2595impl CreditLedger {
2596    fn acquire(&mut self, corr: u64, subscription: bool) {
2597        self.by_corr.entry(corr).or_default().push(RequestCredit {
2598            subscription,
2599            excluded_from_drain: false,
2600        });
2601    }
2602
2603    fn release(&mut self, corr: u64) -> bool {
2604        let Some(credits) = self.by_corr.get_mut(&corr) else {
2605            return false;
2606        };
2607        let released = credits.pop().is_some();
2608        if credits.is_empty() {
2609            self.by_corr.remove(&corr);
2610        }
2611        released
2612    }
2613
2614    fn capture_subscription_exclusions(&mut self) -> u32 {
2615        let mut excluded = 0u32;
2616        for credit in self.by_corr.values_mut().flatten() {
2617            if credit.subscription && !credit.excluded_from_drain {
2618                credit.excluded_from_drain = true;
2619                excluded = excluded.saturating_add(1);
2620            }
2621        }
2622        excluded
2623    }
2624
2625    #[cfg(test)]
2626    fn in_flight(&self) -> usize {
2627        self.by_corr.values().map(Vec::len).sum()
2628    }
2629
2630    fn drain_in_flight(&self) -> usize {
2631        self.by_corr
2632            .values()
2633            .flatten()
2634            .filter(|credit| !credit.excluded_from_drain)
2635            .count()
2636    }
2637
2638    /// The correlation ids of the requests a drain still waits on, one entry
2639    /// per held credit, in ascending order.
2640    fn drain_held_corrs(&self) -> Vec<u64> {
2641        let mut corrs = self
2642            .by_corr
2643            .iter()
2644            .flat_map(|(corr, credits)| {
2645                credits
2646                    .iter()
2647                    .filter(|credit| !credit.excluded_from_drain)
2648                    .map(move |_| *corr)
2649            })
2650            .collect::<Vec<_>>();
2651        corrs.sort_unstable();
2652        corrs
2653    }
2654}
2655
2656#[derive(Debug, Default)]
2657struct ChannelFlowState {
2658    closed: bool,
2659    credits: CreditLedger,
2660}
2661
2662/// Per-channel request-credit accounting shared by the client and module route halves.
2663#[derive(Debug)]
2664pub(crate) struct ChannelFlow {
2665    sem: Semaphore,
2666    window: usize,
2667    state: Mutex<ChannelFlowState>,
2668}
2669
2670impl ChannelFlow {
2671    pub(crate) fn new(window: usize) -> Self {
2672        debug_assert!(window > 0, "flow-control window must be non-zero");
2673        Self {
2674            sem: Semaphore::new(window),
2675            window,
2676            state: Mutex::new(ChannelFlowState::default()),
2677        }
2678    }
2679
2680    #[cfg(test)]
2681    pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
2682        self.acquire_tagged(0, false).await
2683    }
2684
2685    pub(crate) async fn acquire_tagged(
2686        &self,
2687        corr: u64,
2688        subscription: bool,
2689    ) -> Result<(), ChannelFlowClosed> {
2690        let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
2691        let mut state = self
2692            .state
2693            .lock()
2694            .unwrap_or_else(|poisoned| poisoned.into_inner());
2695        if state.closed {
2696            return Err(ChannelFlowClosed);
2697        }
2698        state.credits.acquire(corr, subscription);
2699        permit.forget();
2700        Ok(())
2701    }
2702
2703    #[cfg(test)]
2704    pub(crate) fn release(&self) {
2705        self.release_corr(0);
2706    }
2707
2708    pub(crate) fn release_corr(&self, corr: u64) {
2709        let released = self
2710            .state
2711            .lock()
2712            .unwrap_or_else(|poisoned| poisoned.into_inner())
2713            .credits
2714            .release(corr);
2715        if !released {
2716            // Protocol-conforming modules emit exactly one terminal per request.
2717            // This guard is a best-effort safety net against window growth, not a
2718            // security boundary against malicious peers.
2719            warn!(
2720                window = self.window,
2721                available = self.sem.available_permits(),
2722                "flow-control over-release ignored"
2723            );
2724            return;
2725        }
2726        if !self.sem.is_closed() {
2727            self.sem.add_permits(1);
2728        }
2729    }
2730
2731    #[cfg(test)]
2732    pub(crate) fn in_flight(&self) -> usize {
2733        self.state
2734            .lock()
2735            .unwrap_or_else(|poisoned| poisoned.into_inner())
2736            .credits
2737            .in_flight()
2738    }
2739
2740    pub(crate) fn drain_in_flight(&self) -> usize {
2741        self.state
2742            .lock()
2743            .unwrap_or_else(|poisoned| poisoned.into_inner())
2744            .credits
2745            .drain_in_flight()
2746    }
2747
2748    pub(crate) fn drain_held_corrs(&self) -> Vec<u64> {
2749        self.state
2750            .lock()
2751            .unwrap_or_else(|poisoned| poisoned.into_inner())
2752            .credits
2753            .drain_held_corrs()
2754    }
2755
2756    #[cfg(test)]
2757    pub(crate) fn available_permits(&self) -> usize {
2758        self.sem.available_permits()
2759    }
2760
2761    pub(crate) fn begin_drain(&self) -> u32 {
2762        let mut state = self
2763            .state
2764            .lock()
2765            .unwrap_or_else(|poisoned| poisoned.into_inner());
2766        state.closed = true;
2767        self.sem.close();
2768        state.credits.capture_subscription_exclusions()
2769    }
2770
2771    pub(crate) fn close(&self) {
2772        self.state
2773            .lock()
2774            .unwrap_or_else(|poisoned| poisoned.into_inner())
2775            .closed = true;
2776        self.sem.close();
2777    }
2778}
2779
2780#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2781pub(crate) struct ChannelFlowClosed;
2782
2783impl fmt::Display for ChannelFlowClosed {
2784    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2785        write!(f, "flow-control window closed")
2786    }
2787}
2788
2789impl Error for ChannelFlowClosed {}
2790
2791fn window_for(concurrency: &Concurrency) -> usize {
2792    match concurrency {
2793        Concurrency::Serial => 1,
2794        Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
2795        Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
2796    }
2797}
2798
2799#[derive(Debug, Clone, PartialEq, Eq)]
2800pub enum ForwardingError {
2801    NoModuleConnection,
2802    ModuleReloading {
2803        module_id: String,
2804    },
2805    StaleModuleEndpoint,
2806    UnknownReservation {
2807        client_channel: u16,
2808        module_channel: u16,
2809    },
2810    ClientRouteChannelExhausted {
2811        connection_id: ConnectionId,
2812    },
2813    ModuleRouteChannelExhausted {
2814        endpoint: ModuleEndpointId,
2815    },
2816    RelayCorrelationExhausted,
2817    ConnectionClosing {
2818        connection_id: ConnectionId,
2819    },
2820    ClientEgressClosed {
2821        connection_id: ConnectionId,
2822    },
2823    RouteOpenBuild(String),
2824    /// A swap candidate is already registered for this module id.
2825    CandidateSlotOccupied {
2826        module_id: String,
2827    },
2828    Poisoned,
2829}
2830
2831impl fmt::Display for ForwardingError {
2832    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2833        match self {
2834            Self::NoModuleConnection => write!(f, "no module connection is registered"),
2835            Self::ModuleReloading { module_id } => {
2836                write!(f, "module_id '{module_id}' is reloading")
2837            }
2838            Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
2839            Self::UnknownReservation {
2840                client_channel,
2841                module_channel,
2842            } => write!(
2843                f,
2844                "route reservation client channel {client_channel} / module channel {module_channel} was not found"
2845            ),
2846            Self::ClientRouteChannelExhausted { connection_id } => write!(
2847                f,
2848                "no client route channels are available for connection {}",
2849                connection_id.get()
2850            ),
2851            Self::ModuleRouteChannelExhausted { endpoint } => write!(
2852                f,
2853                "no module route channels are available for endpoint generation {} on connection {}",
2854                endpoint.generation,
2855                endpoint.connection_id.get()
2856            ),
2857            Self::RelayCorrelationExhausted => {
2858                write!(f, "module control correlation ids are exhausted")
2859            }
2860            Self::ConnectionClosing { connection_id } => write!(
2861                f,
2862                "connection {} is closing and cannot accept route allocation",
2863                connection_id.get()
2864            ),
2865            Self::ClientEgressClosed { connection_id } => write!(
2866                f,
2867                "client connection {} egress is closed",
2868                connection_id.get()
2869            ),
2870            Self::RouteOpenBuild(message) => {
2871                write!(f, "failed to prebuild route.open response: {message}")
2872            }
2873            Self::CandidateSlotOccupied { module_id } => write!(
2874                f,
2875                "module_id '{module_id}' already has a swap candidate registered"
2876            ),
2877            Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
2878        }
2879    }
2880}
2881
2882impl Error for ForwardingError {}
2883
2884#[cfg(test)]
2885mod tests {
2886    use std::time::Duration;
2887
2888    use super::*;
2889    use tokio::sync::mpsc;
2890
2891    #[test]
2892    fn ordinary_long_running_request_is_not_excluded_from_drain() {
2893        let mut ledger = CreditLedger::default();
2894        ledger.acquire(1, false);
2895
2896        assert_eq!(ledger.capture_subscription_exclusions(), 0);
2897        assert_eq!(ledger.drain_in_flight(), 1);
2898    }
2899
2900    #[test]
2901    fn bit_set_subscription_is_excluded_and_counted() {
2902        let mut ledger = CreditLedger::default();
2903        ledger.acquire(1, true);
2904
2905        assert_eq!(ledger.capture_subscription_exclusions(), 1);
2906        assert_eq!(ledger.drain_in_flight(), 0);
2907    }
2908
2909    #[test]
2910    fn subscription_opened_after_drain_snapshot_is_not_excluded() {
2911        let mut ledger = CreditLedger::default();
2912        ledger.acquire(1, true);
2913        assert_eq!(ledger.capture_subscription_exclusions(), 1);
2914
2915        ledger.acquire(2, true);
2916
2917        assert_eq!(ledger.drain_in_flight(), 1);
2918    }
2919
2920    #[test]
2921    fn drain_with_no_subscriptions_reports_zero_excluded() {
2922        let mut ledger = CreditLedger::default();
2923        assert_eq!(ledger.capture_subscription_exclusions(), 0);
2924    }
2925
2926    #[test]
2927    fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
2928        let forwarding = ForwardingTable::default();
2929        let module_connection = ConnectionId::new(10);
2930        let exhausted_client = ConnectionId::new(20);
2931        let second_client = ConnectionId::new(30);
2932        let (module_tx, _module_rx) = mpsc::channel(1);
2933        let endpoint = forwarding
2934            .register_module_connection(
2935                module_connection,
2936                "route-limit-provider".to_string(),
2937                1,
2938                Concurrency::ModuleManaged,
2939                FrameSink::new(module_tx),
2940            )
2941            .unwrap();
2942
2943        {
2944            let mut inner = forwarding.inner.write().unwrap();
2945            for channel in 1..=u16::MAX {
2946                inner.reserved_client.insert(
2947                    ClientRouteKey {
2948                        connection_id: exhausted_client,
2949                        channel,
2950                    },
2951                    ModuleRouteKey {
2952                        endpoint,
2953                        channel: 1,
2954                    },
2955                );
2956            }
2957        }
2958
2959        let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
2960        let err = forwarding
2961            .begin_route_bind_relay_for_test(
2962                exhausted_client,
2963                FrameSink::new(exhausted_tx),
2964                1,
2965                "route-limit-provider",
2966            )
2967            .unwrap_err();
2968        assert!(matches!(
2969            err,
2970            ForwardingError::ClientRouteChannelExhausted { connection_id }
2971                if connection_id == exhausted_client
2972        ));
2973
2974        let (second_tx, _second_rx) = mpsc::channel(1);
2975        let pending = forwarding
2976            .begin_route_bind_relay_for_test(
2977                second_client,
2978                FrameSink::new(second_tx),
2979                2,
2980                "route-limit-provider",
2981            )
2982            .unwrap();
2983        assert_eq!(pending.client_channel, 1);
2984    }
2985
2986    #[test]
2987    fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
2988        let forwarding = ForwardingTable::default();
2989        let module_connection = ConnectionId::new(40);
2990        let client = ConnectionId::new(50);
2991        let (module_tx, _module_rx) = mpsc::channel(1);
2992        forwarding
2993            .register_module_connection(
2994                module_connection,
2995                "slot-reuse-provider".to_string(),
2996                1,
2997                Concurrency::ModuleManaged,
2998                FrameSink::new(module_tx),
2999            )
3000            .unwrap();
3001
3002        let (client_tx, _client_rx) = mpsc::channel(1);
3003        let client_sink = FrameSink::new(client_tx);
3004        let mut wrapped_channel = None;
3005        for index in 0..=usize::from(u16::MAX) {
3006            let pending = forwarding
3007                .begin_route_bind_relay_for_test(
3008                    client,
3009                    client_sink.clone(),
3010                    index as u64 + 1,
3011                    "slot-reuse-provider",
3012                )
3013                .unwrap();
3014            if index == usize::from(u16::MAX) {
3015                wrapped_channel = Some(pending.module_channel);
3016            }
3017            forwarding
3018                .abort_pending_relay(
3019                    pending.endpoint,
3020                    pending.corr,
3021                    RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
3022                )
3023                .unwrap();
3024        }
3025
3026        assert_eq!(wrapped_channel, Some(1));
3027    }
3028
3029    #[test]
3030    fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
3031        let forwarding = ForwardingTable::default();
3032        let client = ConnectionId::new(60);
3033        forwarding
3034            .inner
3035            .write()
3036            .unwrap()
3037            .next_client_channel
3038            .insert(client, 41);
3039
3040        let released = forwarding.cleanup_connection(client).unwrap();
3041
3042        assert!(released.is_empty());
3043        assert!(!forwarding
3044            .inner
3045            .read()
3046            .unwrap()
3047            .next_client_channel
3048            .contains_key(&client));
3049    }
3050
3051    #[test]
3052    fn stale_module_cleanup_preserves_fast_reconnect_successor() {
3053        let forwarding = ForwardingTable::default();
3054        let module_id = "fast-reconnect-provider";
3055        let first_connection = ConnectionId::new(70);
3056        let second_connection = ConnectionId::new(80);
3057        let (first_tx, _first_rx) = mpsc::channel(1);
3058        let first_endpoint = forwarding
3059            .register_module_connection(
3060                first_connection,
3061                module_id.to_string(),
3062                1,
3063                Concurrency::ModuleManaged,
3064                FrameSink::new(first_tx),
3065            )
3066            .unwrap();
3067        let (second_tx, _second_rx) = mpsc::channel(1);
3068        let second_endpoint = forwarding
3069            .register_module_connection(
3070                second_connection,
3071                module_id.to_string(),
3072                1,
3073                Concurrency::ModuleManaged,
3074                FrameSink::new(second_tx),
3075            )
3076            .unwrap();
3077        assert_ne!(first_endpoint, second_endpoint);
3078
3079        let released = forwarding.cleanup_connection(first_connection).unwrap();
3080
3081        assert!(released.is_empty());
3082        assert_eq!(
3083            forwarding
3084                .inner
3085                .read()
3086                .unwrap()
3087                .modules_by_id
3088                .get(module_id)
3089                .map(|module| module.endpoint),
3090            Some(second_endpoint)
3091        );
3092        assert!(forwarding.has_live_module_connection(module_id).unwrap());
3093        let control_rpc = forwarding
3094            .begin_module_control_rpc_for(
3095                module_id,
3096                "health.check",
3097                Instant::now() + Duration::from_secs(1),
3098            )
3099            .unwrap();
3100        assert_eq!(control_rpc.endpoint, second_endpoint);
3101    }
3102
3103    fn route_fixture(
3104        module_id: &str,
3105    ) -> (
3106        ForwardingTable,
3107        ConnectionId,
3108        ModuleEndpointId,
3109        ConnectionId,
3110        FrameSink,
3111        mpsc::Receiver<crate::router::OutboundFrame>,
3112    ) {
3113        let forwarding = ForwardingTable::default();
3114        let module_connection = ConnectionId::new(100);
3115        let client_connection = ConnectionId::new(200);
3116        let (module_tx, _module_rx) = mpsc::channel(8);
3117        let endpoint = forwarding
3118            .register_module_connection(
3119                module_connection,
3120                module_id.to_string(),
3121                2,
3122                Concurrency::ModuleManaged,
3123                FrameSink::new(module_tx),
3124            )
3125            .unwrap();
3126        let (client_tx, client_rx) = mpsc::channel(8);
3127        (
3128            forwarding,
3129            module_connection,
3130            endpoint,
3131            client_connection,
3132            FrameSink::new(client_tx),
3133            client_rx,
3134        )
3135    }
3136
3137    #[test]
3138    #[cfg(unix)]
3139    fn daemon_drain_gates_current_and_racing_provider_registrations() {
3140        let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
3141        assert_eq!(forwarding.begin_daemon_drain().unwrap(), ["provider"]);
3142        assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
3143        assert!(matches!(
3144            forwarding.register_module_connection(
3145                ConnectionId::new(300),
3146                "late-provider".into(),
3147                2,
3148                Concurrency::ModuleManaged,
3149                sink,
3150            ),
3151            Err(ForwardingError::ConnectionClosing { .. })
3152        ));
3153    }
3154
3155    fn test_ping(corr: u64) -> Frame {
3156        Frame::build(
3157            FrameType::Ping,
3158            Flags::new(false, Priority::Passive, false),
3159            0,
3160            0,
3161            corr,
3162            Vec::new(),
3163        )
3164        .unwrap()
3165    }
3166
3167    fn begin_test_route(
3168        forwarding: &ForwardingTable,
3169        client_connection: ConnectionId,
3170        client_sink: FrameSink,
3171        corr: u64,
3172        module_id: &str,
3173    ) -> PendingRouteBindRelay {
3174        forwarding
3175            .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
3176            .unwrap()
3177    }
3178
3179    /// The drain-timeout line reports these numbers, so they must count the
3180    /// requests the drain is waiting on and no others: a route with nothing in
3181    /// flight is not a holdout, and a flagged subscription is excluded from the
3182    /// drain and so from the count.
3183    #[tokio::test]
3184    async fn drain_holdouts_count_held_requests_and_name_the_connection() {
3185        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3186            route_fixture("holdouts");
3187        let mut bound = |corr| {
3188            let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
3189            forwarding
3190                .complete_pending_relay(
3191                    module_connection,
3192                    route.corr,
3193                    RouteBindRelayOutcome::Accepted,
3194                )
3195                .unwrap();
3196            client_rx.try_recv().unwrap();
3197            match forwarding
3198                .lookup_data_route(client, route.client_channel, route.client_epoch)
3199                .unwrap()
3200            {
3201                DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3202                other => panic!("expected live route, got {other:?}"),
3203            }
3204        };
3205        let holding = bound(61);
3206        let _idle = bound(62);
3207        holding.flow.acquire_tagged(7, false).await.unwrap();
3208        holding.flow.acquire_tagged(2, false).await.unwrap();
3209        holding.flow.acquire_tagged(3, true).await.unwrap();
3210        forwarding
3211            .begin_module_drain("holdouts", RouteCloseReason::Restart)
3212            .unwrap();
3213
3214        let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
3215        assert_eq!(
3216            holdouts,
3217            DrainHoldouts {
3218                requests: 2,
3219                routes: 1,
3220                total_routes: 2,
3221                top_connections: vec![(client.get(), 2)],
3222                // The held requests by the module's channel and corr, ascending,
3223                // without the excluded subscription (corr 3).
3224                held: vec![(holding.module_channel, 2), (holding.module_channel, 7)],
3225            }
3226        );
3227    }
3228
3229    #[test]
3230    fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
3231        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3232            route_fixture("census");
3233        let pending = begin_test_route(&forwarding, client, sink, 1, "census");
3234        forwarding
3235            .complete_pending_relay(
3236                module_connection,
3237                pending.corr,
3238                RouteBindRelayOutcome::Accepted,
3239            )
3240            .unwrap();
3241
3242        let routes = forwarding.endpoint_routes(endpoint).unwrap();
3243        assert_eq!(routes.len(), 1);
3244        assert!(matches!(routes[0].principal, Principal::Direct));
3245        assert_eq!(routes[0].goodbye_target.connection_id, client);
3246        assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
3247        assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
3248        assert!(!routes[0].draining);
3249
3250        forwarding
3251            .begin_module_drain("census", RouteCloseReason::Restart)
3252            .unwrap();
3253        let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
3254        assert_eq!(draining_routes.len(), 1);
3255        assert!(draining_routes[0].draining);
3256    }
3257
3258    #[test]
3259    fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
3260        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
3261        let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
3262        assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
3263        forwarding
3264            .abort_pending_relay(
3265                first.endpoint,
3266                first.corr,
3267                RouteBindRelayOutcome::ModuleGone("abort".into()),
3268            )
3269            .unwrap();
3270        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3271        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3272
3273        let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
3274        assert_eq!(second.client_channel, first.client_channel);
3275        assert_eq!(second.module_channel, first.module_channel);
3276        assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
3277    }
3278
3279    #[test]
3280    fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
3281        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3282            route_fixture("epoch-release");
3283        let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
3284        forwarding
3285            .complete_pending_relay(
3286                module_connection,
3287                first.corr,
3288                RouteBindRelayOutcome::Accepted,
3289            )
3290            .unwrap();
3291        assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
3292        assert!(matches!(
3293            forwarding
3294                .release_client_route(client, first.client_channel, first.client_epoch)
3295                .unwrap(),
3296            RouteRelease::Removed(_)
3297        ));
3298        forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3299        forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3300
3301        let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
3302        forwarding
3303            .complete_pending_relay(
3304                module_connection,
3305                second.corr,
3306                RouteBindRelayOutcome::Accepted,
3307            )
3308            .unwrap();
3309        assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
3310        assert!(matches!(
3311            forwarding
3312                .release_client_route(client, second.client_channel, first.client_epoch)
3313                .unwrap(),
3314            RouteRelease::Stale
3315        ));
3316        assert!(!forwarding
3317            .cache_status(
3318                endpoint,
3319                second.module_channel,
3320                first.module_epoch,
3321                "stale".into(),
3322            )
3323            .unwrap());
3324        assert!(forwarding
3325            .cache_status(
3326                endpoint,
3327                second.module_channel,
3328                second.module_epoch,
3329                "current".into(),
3330            )
3331            .unwrap());
3332        match forwarding
3333            .route_poll_snapshot(client, second.client_channel, second.client_epoch)
3334            .unwrap()
3335        {
3336            RoutePollSnapshot::Bound { status, .. } => {
3337                assert_eq!(status.as_deref(), Some("current"));
3338            }
3339            RoutePollSnapshot::Absent => panic!("successor binding was removed"),
3340        }
3341        let counters = forwarding.counters().snapshot();
3342        assert_eq!(counters["route_released_epoch_fenced"], 1);
3343        assert_eq!(counters["route_release_stale_skipped"], 1);
3344    }
3345
3346    #[test]
3347    fn max_epoch_reservation_retires_only_that_slot() {
3348        let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
3349        forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
3350        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
3351        let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
3352        assert_eq!(
3353            (final_use.client_channel, final_use.client_epoch),
3354            (7, u32::MAX)
3355        );
3356        assert_eq!(
3357            (final_use.module_channel, final_use.module_epoch),
3358            (9, u32::MAX)
3359        );
3360        forwarding
3361            .abort_pending_relay(
3362                endpoint,
3363                final_use.corr,
3364                RouteBindRelayOutcome::ModuleGone("abort".into()),
3365            )
3366            .unwrap();
3367        forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
3368        forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
3369        let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
3370        assert_ne!(next.client_channel, 7);
3371        assert_ne!(next.module_channel, 9);
3372        assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
3373    }
3374
3375    #[test]
3376    fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
3377        let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3378            route_fixture("corr-shared");
3379        let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
3380        assert_eq!(bind.corr, 1);
3381        forwarding
3382            .abort_pending_relay(
3383                endpoint,
3384                bind.corr,
3385                RouteBindRelayOutcome::ModuleGone("abort".into()),
3386            )
3387            .unwrap();
3388        let rpc = forwarding
3389            .begin_module_control_rpc_for(
3390                "corr-shared",
3391                "health.check",
3392                Instant::now() - Duration::from_millis(1),
3393            )
3394            .unwrap();
3395        assert_eq!(rpc.corr, 2);
3396        assert_eq!(
3397            forwarding
3398                .complete_module_control_rpc(
3399                    module_connection,
3400                    rpc.corr,
3401                    Some("health.check"),
3402                    ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
3403                        status: subc_protocol::session::HealthStatus::Ok,
3404                        detail: None,
3405                        metrics: None,
3406                    }),
3407                )
3408                .unwrap(),
3409            ModuleControlRpcCompletion::Settled
3410        );
3411        assert!(matches!(
3412            rpc.receiver.blocking_recv().unwrap(),
3413            ModuleControlRpcOutcome::DeadlineElapsed
3414        ));
3415    }
3416
3417    #[tokio::test(start_paused = true)]
3418    async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
3419        let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
3420        let probe_started_at = Instant::now();
3421        let rpc = forwarding
3422            .begin_health_probe_rpc_for(
3423                "tombstone-ttl",
3424                "health.check",
3425                probe_started_at,
3426                probe_started_at + Duration::from_secs(5),
3427            )
3428            .unwrap();
3429        assert!(forwarding
3430            .tombstone_health_probe_rpc(endpoint, rpc.corr)
3431            .unwrap());
3432        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
3433
3434        tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
3435        tokio::task::yield_now().await;
3436
3437        assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
3438    }
3439
3440    #[test]
3441    fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
3442        let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
3443        let mut close = forwarding.register_connection_close(endpoint.connection_id);
3444        forwarding.inject_control_corr(endpoint, u64::MAX);
3445        let final_rpc = forwarding
3446            .begin_module_control_rpc_for(
3447                "corr-max",
3448                "health.check",
3449                Instant::now() + Duration::from_secs(1),
3450            )
3451            .unwrap();
3452        assert_eq!(final_rpc.corr, u64::MAX);
3453        forwarding
3454            .cancel_module_control_rpc(endpoint, final_rpc.corr)
3455            .unwrap();
3456        assert!(matches!(
3457            forwarding.begin_module_control_rpc_for(
3458                "corr-max",
3459                "health.check",
3460                Instant::now() + Duration::from_secs(1),
3461            ),
3462            Err(ForwardingError::RelayCorrelationExhausted)
3463        ));
3464        assert!(close.try_recv().is_ok());
3465    }
3466
3467    #[test]
3468    fn publication_epoch_controls_delivery_failure_escalation() {
3469        fn setup_successor(
3470            commit_successor: Option<bool>,
3471        ) -> (ForwardingTable, ConnectionId, u16, u32) {
3472            let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3473                route_fixture("escalation");
3474            let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
3475            forwarding
3476                .complete_pending_relay(
3477                    module_connection,
3478                    first.corr,
3479                    RouteBindRelayOutcome::Accepted,
3480                )
3481                .unwrap();
3482            client_rx.try_recv().unwrap();
3483            assert!(matches!(
3484                forwarding
3485                    .release_client_route(client, first.client_channel, first.client_epoch)
3486                    .unwrap(),
3487                RouteRelease::Removed(_)
3488            ));
3489            if let Some(commit_successor) = commit_successor {
3490                forwarding.inject_client_slot_epoch(
3491                    client,
3492                    first.client_channel,
3493                    first.client_epoch,
3494                );
3495                forwarding.inject_module_slot_epoch(
3496                    endpoint,
3497                    first.module_channel,
3498                    first.module_epoch,
3499                );
3500                let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
3501                if commit_successor {
3502                    forwarding
3503                        .complete_pending_relay(
3504                            module_connection,
3505                            successor.corr,
3506                            RouteBindRelayOutcome::Accepted,
3507                        )
3508                        .unwrap();
3509                    client_rx.try_recv().unwrap();
3510                } else {
3511                    forwarding
3512                        .abort_pending_relay(
3513                            endpoint,
3514                            successor.corr,
3515                            RouteBindRelayOutcome::ModuleGone("abort".into()),
3516                        )
3517                        .unwrap();
3518                }
3519            }
3520            (forwarding, client, first.client_channel, first.client_epoch)
3521        }
3522
3523        let (no_successor, client, channel, epoch) = setup_successor(None);
3524        let mut close = no_successor.register_connection_close(client);
3525        assert!(no_successor
3526            .escalate_client_delivery_failure(
3527                client,
3528                channel,
3529                epoch,
3530                CloseReason::new("delivery", "failed"),
3531            )
3532            .unwrap());
3533        assert!(close.try_recv().is_ok());
3534
3535        let (aborted, client, channel, epoch) = setup_successor(Some(false));
3536        let mut close = aborted.register_connection_close(client);
3537        assert!(aborted
3538            .escalate_client_delivery_failure(
3539                client,
3540                channel,
3541                epoch,
3542                CloseReason::new("delivery", "failed"),
3543            )
3544            .unwrap());
3545        assert!(close.try_recv().is_ok());
3546
3547        let (published, client, channel, epoch) = setup_successor(Some(true));
3548        let mut close = published.register_connection_close(client);
3549        assert!(!published
3550            .escalate_client_delivery_failure(
3551                client,
3552                channel,
3553                epoch,
3554                CloseReason::new("delivery", "stale failure"),
3555            )
3556            .unwrap());
3557        assert!(close.try_recv().is_err());
3558    }
3559
3560    #[test]
3561    fn route_concentration_separates_client_count_from_routes_per_client() {
3562        // The distinction this asserts is the one a bare connection count cannot
3563        // make: two connections holding one route each and one connection
3564        // holding two are the same total, and have opposite causes.
3565        let (forwarding, module_connection, _, client, sink, _client_rx) =
3566            route_fixture("concentration");
3567        assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
3568
3569        for corr in [70_u64, 71] {
3570            let pending =
3571                begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
3572            forwarding
3573                .complete_pending_relay(
3574                    module_connection,
3575                    pending.corr,
3576                    RouteBindRelayOutcome::Accepted,
3577                )
3578                .unwrap();
3579        }
3580
3581        // One connection, two routes — not two connections with a route each.
3582        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3583        assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
3584    }
3585
3586    #[test]
3587    fn cleanup_and_accepted_resolution_have_one_lock_winner() {
3588        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3589            route_fixture("cleanup-race");
3590        let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
3591        forwarding
3592            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3593            .unwrap();
3594        let released = forwarding.cleanup_connection(client).unwrap();
3595        assert_eq!(released.len(), 1);
3596        let completion = forwarding
3597            .complete_pending_relay(
3598                module_connection,
3599                pending.corr,
3600                RouteBindRelayOutcome::Accepted,
3601            )
3602            .unwrap();
3603        assert!(!completion.settled);
3604        assert!(client_rx.try_recv().is_err());
3605        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3606
3607        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3608            route_fixture("accepted-race");
3609        let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
3610        forwarding
3611            .complete_pending_relay(
3612                module_connection,
3613                pending.corr,
3614                RouteBindRelayOutcome::Accepted,
3615            )
3616            .unwrap();
3617        assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
3618        let released = forwarding.cleanup_connection(client).unwrap();
3619        assert_eq!(released.len(), 1);
3620        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3621    }
3622
3623    #[test]
3624    fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
3625        let (forwarding, module_connection, _, client, sink, mut client_rx) =
3626            route_fixture("drain-gap");
3627        let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
3628        forwarding
3629            .complete_pending_relay(
3630                module_connection,
3631                live.corr,
3632                RouteBindRelayOutcome::Accepted,
3633            )
3634            .unwrap();
3635        client_rx.try_recv().unwrap();
3636        let binding = match forwarding
3637            .lookup_data_route(client, live.client_channel, live.client_epoch)
3638            .unwrap()
3639        {
3640            DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3641            other => panic!("expected live route, got {other:?}"),
3642        };
3643
3644        let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
3645        forwarding
3646            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3647            .unwrap();
3648        let control_rpc = forwarding
3649            .begin_module_control_rpc_for(
3650                "drain-gap",
3651                "health.check",
3652                Instant::now() + Duration::from_secs(1),
3653            )
3654            .unwrap();
3655        let target = forwarding
3656            .begin_module_drain("drain-gap", RouteCloseReason::Reload)
3657            .unwrap()
3658            .unwrap();
3659        assert!(matches!(
3660            control_rpc.receiver.blocking_recv().unwrap(),
3661            ModuleControlRpcOutcome::ModuleGone(_)
3662        ));
3663        assert_eq!(target.abandoned_bindings.len(), 1);
3664        assert!(binding.flow.sem.is_closed());
3665        assert!(
3666            !forwarding
3667                .complete_pending_relay(
3668                    module_connection,
3669                    pending.corr,
3670                    RouteBindRelayOutcome::Accepted,
3671                )
3672                .unwrap()
3673                .settled
3674        );
3675        assert!(matches!(
3676            forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
3677            Err(ForwardingError::ModuleReloading { .. })
3678        ));
3679        let released = forwarding
3680            .release_module_endpoint_routes(target.endpoint)
3681            .unwrap();
3682        assert_eq!(released.len(), 1);
3683        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3684    }
3685
3686    /// A client can be marked closing while its egress is still open: the daemon
3687    /// asks a connection to close (here through the production path, a module
3688    /// frame that its egress refused) and the connection loop tears down a moment
3689    /// later. A route.bind ack that lands inside that window is answered on the
3690    /// MODULE connection's frame handler, so resolving it must not produce an
3691    /// error -- an error there ends the module connection, and that connection
3692    /// carries every other client's routes to the module.
3693    #[test]
3694    fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
3695        let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3696            route_fixture("closing-client");
3697
3698        // A published route on this client: escalate_client_delivery_failure only
3699        // marks a connection closing for a route it has already published.
3700        let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
3701        forwarding
3702            .complete_pending_relay(
3703                module_connection,
3704                live.corr,
3705                RouteBindRelayOutcome::Accepted,
3706            )
3707            .unwrap();
3708        client_rx.try_recv().unwrap();
3709
3710        // A second route.open from the same client, relayed and awaiting its ack.
3711        let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
3712        forwarding
3713            .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3714            .unwrap();
3715
3716        // The window: closing, but the sink is still open.
3717        assert!(forwarding
3718            .escalate_client_delivery_failure(
3719                client,
3720                live.client_channel,
3721                live.client_epoch,
3722                CloseReason::new(
3723                    "module_to_client_delivery_failed",
3724                    "client egress refused a module frame",
3725                ),
3726            )
3727            .unwrap());
3728        assert!(!sink.is_closed());
3729
3730        let completion = forwarding
3731            .complete_pending_relay(
3732                module_connection,
3733                pending.corr,
3734                RouteBindRelayOutcome::Accepted,
3735            )
3736            .expect("a closing client must not turn a module's ack into an error");
3737
3738        assert!(completion.settled);
3739        let abandoned = completion
3740            .abandoned
3741            .expect("the module must be told to drop the binding it just created");
3742        assert_eq!(abandoned.connection_id, module_connection);
3743        assert_eq!(abandoned.channel, pending.module_channel);
3744        assert_eq!(abandoned.epoch, pending.module_epoch);
3745        assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
3746        assert!(matches!(
3747            pending.receiver.blocking_recv().unwrap(),
3748            RouteBindRelayOutcome::ModuleGone(_)
3749        ));
3750        // No route was published to a client that is on its way out, and the
3751        // reserved handle pair went back.
3752        assert!(client_rx.try_recv().is_err());
3753        assert_eq!(forwarding.active_binding_count().unwrap(), 1);
3754
3755        // The module endpoint is untouched: still live, and still able to take a
3756        // route from another client.
3757        assert!(forwarding
3758            .has_live_module_connection("closing-client")
3759            .unwrap());
3760        let cotenant = ConnectionId::new(201);
3761        let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
3762        let cotenant_route = begin_test_route(
3763            &forwarding,
3764            cotenant,
3765            FrameSink::new(cotenant_tx),
3766            62,
3767            "closing-client",
3768        );
3769        assert_eq!(cotenant_route.endpoint, endpoint);
3770        forwarding
3771            .complete_pending_relay(
3772                module_connection,
3773                cotenant_route.corr,
3774                RouteBindRelayOutcome::Accepted,
3775            )
3776            .unwrap();
3777        assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
3778        assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3779    }
3780
3781    #[test]
3782    fn pending_route_permit_is_released_on_rejection_and_abort() {
3783        let forwarding = ForwardingTable::default();
3784        let module_connection = ConnectionId::new(300);
3785        let client = ConnectionId::new(301);
3786        let (module_tx, _module_rx) = mpsc::channel(1);
3787        let endpoint = forwarding
3788            .register_module_connection(
3789                module_connection,
3790                "permit".into(),
3791                2,
3792                Concurrency::ModuleManaged,
3793                FrameSink::new(module_tx),
3794            )
3795            .unwrap();
3796        let (client_tx, mut client_rx) = mpsc::channel(1);
3797        let sink = FrameSink::new(client_tx);
3798        let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
3799        assert!(sink.try_send(test_ping(999)).is_err());
3800        forwarding
3801            .complete_pending_relay(
3802                module_connection,
3803                rejected.corr,
3804                RouteBindRelayOutcome::Rejected(ErrorBody {
3805                    code: "no".into(),
3806                    message: "rejected".into(),
3807                    detail: None,
3808                }),
3809            )
3810            .unwrap();
3811        sink.try_send(test_ping(1000)).unwrap();
3812        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
3813
3814        let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
3815        assert!(sink.try_send(test_ping(1001)).is_err());
3816        forwarding
3817            .abort_pending_relay(
3818                endpoint,
3819                aborted.corr,
3820                RouteBindRelayOutcome::ModuleGone("abort".into()),
3821            )
3822            .unwrap();
3823        sink.try_send(test_ping(1002)).unwrap();
3824        assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
3825
3826        let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
3827        forwarding
3828            .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
3829            .unwrap();
3830        drop(client_rx);
3831        let completion = forwarding
3832            .complete_pending_relay(
3833                module_connection,
3834                receiver_closed.corr,
3835                RouteBindRelayOutcome::Accepted,
3836            )
3837            .unwrap();
3838        assert!(completion.abandoned.is_some());
3839        assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3840    }
3841
3842    /// Every connection teardown passes through `cleanup_connection`, and
3843    /// connection ids come from a monotonic counter that never hands an id out
3844    /// twice. If the closing mark survives teardown, the set grows by one entry
3845    /// per connection for the life of the daemon -- the self-watchdog alone
3846    /// reconnects once a minute.
3847    #[test]
3848    fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
3849        let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
3850            route_fixture("closing-set-leak");
3851
3852        const CONNECTIONS: u64 = 32;
3853        for index in 0..CONNECTIONS {
3854            let client = ConnectionId::new(1000 + index);
3855            let (client_tx, _client_rx) = mpsc::channel(8);
3856            let route = begin_test_route(
3857                &forwarding,
3858                client,
3859                FrameSink::new(client_tx),
3860                index + 1,
3861                "closing-set-leak",
3862            );
3863            forwarding
3864                .complete_pending_relay(
3865                    module_connection,
3866                    route.corr,
3867                    RouteBindRelayOutcome::Accepted,
3868                )
3869                .unwrap();
3870            forwarding.cleanup_connection(client).unwrap();
3871        }
3872        forwarding.cleanup_connection(module_connection).unwrap();
3873
3874        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3875    }
3876
3877    /// The closing mark exists to refuse new work for a connection that is on
3878    /// its way out but whose teardown has not run yet: the daemon asks the
3879    /// connection loop to end, and only when the loop reacts does
3880    /// `cleanup_connection` strip the connection's state. Inside that window an
3881    /// operation for the dying connection must still be refused; only after
3882    /// cleanup completes may the mark go.
3883    #[test]
3884    fn closing_connection_is_refused_new_work_until_cleanup_completes() {
3885        let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
3886            route_fixture("closing-gate");
3887
3888        // A published route: escalate_client_delivery_failure only marks a
3889        // connection closing for a route it has already published.
3890        let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
3891        forwarding
3892            .complete_pending_relay(
3893                module_connection,
3894                live.corr,
3895                RouteBindRelayOutcome::Accepted,
3896            )
3897            .unwrap();
3898        client_rx.try_recv().unwrap();
3899
3900        // Mark the connection closing through the production path without
3901        // running teardown, pinning the window open.
3902        assert!(forwarding
3903            .escalate_client_delivery_failure(
3904                client,
3905                live.client_channel,
3906                live.client_epoch,
3907                CloseReason::new(
3908                    "module_to_client_delivery_failed",
3909                    "client egress refused a module frame",
3910                ),
3911            )
3912            .unwrap());
3913        assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
3914
3915        // A late route.open for the closing client is refused ...
3916        assert!(matches!(
3917            forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
3918            Err(ForwardingError::ConnectionClosing { connection_id })
3919                if connection_id == client
3920        ));
3921        // ... and so is a late attempt to register the connection as a module.
3922        let (late_tx, _late_rx) = mpsc::channel(1);
3923        assert!(matches!(
3924            forwarding.register_module_connection(
3925                client,
3926                "late-module".into(),
3927                2,
3928                Concurrency::ModuleManaged,
3929                FrameSink::new(late_tx),
3930            ),
3931            Err(ForwardingError::ConnectionClosing { connection_id })
3932                if connection_id == client
3933        ));
3934
3935        // Teardown is the point that lifts the mark: it has just removed every
3936        // per-connection entry under the same lock, so the gate has nothing
3937        // left to protect for this id.
3938        forwarding.cleanup_connection(client).unwrap();
3939        assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3940    }
3941}
3942
3943/// Blue/green swap slots: a candidate registered beside the active endpoint,
3944/// promoted by `cutover_candidate`, with the old incumbent drained by endpoint.
3945#[cfg(test)]
3946mod swap_slot_tests {
3947    use std::time::Duration;
3948
3949    use super::*;
3950    use tokio::sync::mpsc;
3951
3952    const MODULE_ID: &str = "swapped";
3953
3954    struct SwapFixture {
3955        forwarding: ForwardingTable,
3956        incumbent_connection: ConnectionId,
3957        incumbent: ModuleEndpointId,
3958        candidate_connection: ConnectionId,
3959        candidate: ModuleEndpointId,
3960        _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
3961    }
3962
3963    fn swap_fixture() -> SwapFixture {
3964        let forwarding = ForwardingTable::default();
3965        let incumbent_connection = ConnectionId::new(100);
3966        let candidate_connection = ConnectionId::new(110);
3967        let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
3968        let incumbent = forwarding
3969            .register_module_connection(
3970                incumbent_connection,
3971                MODULE_ID.to_string(),
3972                2,
3973                Concurrency::ModuleManaged,
3974                FrameSink::new(incumbent_tx),
3975            )
3976            .unwrap();
3977        let (candidate_tx, candidate_rx) = mpsc::channel(8);
3978        let candidate = forwarding
3979            .register_candidate_module_connection(
3980                candidate_connection,
3981                MODULE_ID.to_string(),
3982                2,
3983                Concurrency::ModuleManaged,
3984                FrameSink::new(candidate_tx),
3985            )
3986            .unwrap();
3987        SwapFixture {
3988            forwarding,
3989            incumbent_connection,
3990            incumbent,
3991            candidate_connection,
3992            candidate,
3993            _module_rxs: vec![incumbent_rx, candidate_rx],
3994        }
3995    }
3996
3997    fn client(
3998        raw: u64,
3999    ) -> (
4000        ConnectionId,
4001        FrameSink,
4002        mpsc::Receiver<crate::router::OutboundFrame>,
4003    ) {
4004        let (tx, rx) = mpsc::channel(8);
4005        (ConnectionId::new(raw), FrameSink::new(tx), rx)
4006    }
4007
4008    fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
4009        forwarding
4010            .read_inner()
4011            .unwrap()
4012            .client_to_module
4013            .values()
4014            .map(|route| route.module_endpoint)
4015            .collect()
4016    }
4017
4018    #[test]
4019    fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
4020        let fixture = swap_fixture();
4021        let forwarding = &fixture.forwarding;
4022        assert_ne!(fixture.incumbent, fixture.candidate);
4023
4024        // Every by-id consumer still resolves the incumbent.
4025        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4026        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4027        let (client_connection, client_sink, _client_rx) = client(200);
4028        let pending = forwarding
4029            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4030            .unwrap();
4031        assert_eq!(pending.endpoint, fixture.incumbent);
4032        let rpc = forwarding
4033            .begin_module_control_rpc_for(
4034                MODULE_ID,
4035                "health.check",
4036                Instant::now() + Duration::from_secs(1),
4037            )
4038            .unwrap();
4039        assert_eq!(rpc.endpoint, fixture.incumbent);
4040        let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
4041        assert_eq!(census.len(), 1, "the census lists one endpoint per id");
4042
4043        // Connection-keyed lookups see the candidate, so its own frames resolve.
4044        assert_eq!(
4045            forwarding
4046                .module_endpoint_for_connection(fixture.candidate_connection)
4047                .unwrap(),
4048            Some(fixture.candidate)
4049        );
4050        assert_eq!(
4051            forwarding
4052                .module_id_for_connection(fixture.candidate_connection)
4053                .unwrap()
4054                .as_deref(),
4055            Some(MODULE_ID)
4056        );
4057
4058        // One candidate per id.
4059        let (other_tx, _other_rx) = mpsc::channel(1);
4060        assert_eq!(
4061            forwarding.register_candidate_module_connection(
4062                ConnectionId::new(120),
4063                MODULE_ID.to_string(),
4064                2,
4065                Concurrency::ModuleManaged,
4066                FrameSink::new(other_tx),
4067            ),
4068            Err(ForwardingError::CandidateSlotOccupied {
4069                module_id: MODULE_ID.to_string()
4070            })
4071        );
4072    }
4073
4074    /// The cutover linearization point. A relay reserved on the incumbent before
4075    /// cutover must never become a route on the incumbent, and every relay
4076    /// reserved after it must land on the promoted candidate.
4077    #[test]
4078    fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
4079        let fixture = swap_fixture();
4080        let forwarding = &fixture.forwarding;
4081        let (early_client, early_sink, _early_rx) = client(200);
4082        let mut early = forwarding
4083            .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
4084            .unwrap();
4085        assert_eq!(early.endpoint, fixture.incumbent);
4086        assert!(forwarding
4087            .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
4088            .unwrap());
4089
4090        let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4091        assert_eq!(
4092            cutover,
4093            ForwardingCutover {
4094                promoted: fixture.candidate,
4095                incumbent: Some(fixture.incumbent),
4096            }
4097        );
4098
4099        // A route.open reserved after cutover goes to the promoted candidate.
4100        let (late_client, late_sink, _late_rx) = client(201);
4101        let late = forwarding
4102            .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
4103            .unwrap();
4104        assert_eq!(
4105            late.endpoint, fixture.candidate,
4106            "a route.open after cutover was reserved on the incumbent"
4107        );
4108
4109        // The incumbent acks the early relay after cutover.
4110        let completion = forwarding
4111            .complete_pending_relay(
4112                fixture.incumbent_connection,
4113                early.corr,
4114                RouteBindRelayOutcome::Accepted,
4115            )
4116            .expect("a superseded endpoint's ack is not an error on its connection");
4117        assert!(completion.settled);
4118        assert!(
4119            !committed_endpoints(forwarding).contains(&fixture.incumbent),
4120            "a relay reserved before cutover committed a route on the incumbent"
4121        );
4122        let goodbye = completion
4123            .abandoned
4124            .expect("the incumbent is told to drop the binding it just created");
4125        assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
4126        assert_eq!(goodbye.channel, early.module_channel);
4127        assert_eq!(goodbye.epoch, early.module_epoch);
4128        assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
4129        match early.receiver.try_recv() {
4130            Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
4131            other => panic!("expected a retryable module_reloading answer, got {other:?}"),
4132        }
4133        assert!(matches!(
4134            forwarding
4135                .lookup_data_route(early_client, early.client_channel, early.client_epoch)
4136                .unwrap(),
4137            DataRoute::Client(DataRouteState::Absent)
4138        ));
4139
4140        // The late relay commits on the candidate; only its pair stays reserved
4141        // until then.
4142        assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4143        forwarding
4144            .complete_pending_relay(
4145                fixture.candidate_connection,
4146                late.corr,
4147                RouteBindRelayOutcome::Accepted,
4148            )
4149            .unwrap();
4150        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4151        assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
4152    }
4153
4154    #[test]
4155    fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
4156        let fixture = swap_fixture();
4157        let forwarding = &fixture.forwarding;
4158        // One bound route and one in-flight relay on the incumbent.
4159        let (bound_client, bound_sink, _bound_rx) = client(200);
4160        let bound = forwarding
4161            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4162            .unwrap();
4163        forwarding
4164            .complete_pending_relay(
4165                fixture.incumbent_connection,
4166                bound.corr,
4167                RouteBindRelayOutcome::Accepted,
4168            )
4169            .unwrap();
4170        let (pending_client, pending_sink, _pending_rx) = client(201);
4171        let mut in_flight = forwarding
4172            .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
4173            .unwrap();
4174        forwarding
4175            .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
4176            .unwrap();
4177
4178        let incumbent = forwarding
4179            .cutover_candidate(MODULE_ID)
4180            .unwrap()
4181            .unwrap()
4182            .incumbent
4183            .unwrap();
4184        let target = forwarding
4185            .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
4186            .unwrap()
4187            .expect("the superseded incumbent is still registered");
4188
4189        assert_eq!(target.endpoint, fixture.incumbent);
4190        assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
4191        assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
4192        assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4193        assert_eq!(target.abandoned_bindings.len(), 1);
4194        assert_eq!(
4195            target.abandoned_bindings[0].channel,
4196            in_flight.module_channel
4197        );
4198        assert!(matches!(
4199            in_flight.receiver.try_recv(),
4200            Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
4201        ));
4202        assert_eq!(
4203            forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
4204            1,
4205            "the incumbent's bound route stays until its drain finishes"
4206        );
4207
4208        let (next_client, next_sink, _next_rx) = client(202);
4209        let next = forwarding
4210            .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
4211            .expect("the promoted candidate keeps accepting routes");
4212        assert_eq!(next.endpoint, fixture.candidate);
4213    }
4214
4215    /// An endpoint replaced WITHOUT a promotion (a successor registered over it
4216    /// as an ordinary active HELLO) is stale, not superseded, and its ack still
4217    /// fails the acking connection exactly as it did before swap slots existed:
4218    /// `StaleModuleEndpoint`, the reservation indexes already stripped, and the
4219    /// waiting client's sender dropped unanswered.
4220    #[test]
4221    fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
4222        let forwarding = ForwardingTable::default();
4223        let first_connection = ConnectionId::new(70);
4224        let (first_tx, _first_rx) = mpsc::channel(8);
4225        forwarding
4226            .register_module_connection(
4227                first_connection,
4228                MODULE_ID.to_string(),
4229                2,
4230                Concurrency::ModuleManaged,
4231                FrameSink::new(first_tx),
4232            )
4233            .unwrap();
4234        let (client_connection, client_sink, _client_rx) = client(200);
4235        let mut pending = forwarding
4236            .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4237            .unwrap();
4238        let (second_tx, _second_rx) = mpsc::channel(8);
4239        forwarding
4240            .register_module_connection(
4241                ConnectionId::new(80),
4242                MODULE_ID.to_string(),
4243                2,
4244                Concurrency::ModuleManaged,
4245                FrameSink::new(second_tx),
4246            )
4247            .unwrap();
4248
4249        assert_eq!(
4250            forwarding
4251                .complete_pending_relay(
4252                    first_connection,
4253                    pending.corr,
4254                    RouteBindRelayOutcome::Accepted
4255                )
4256                .unwrap_err(),
4257            ForwardingError::StaleModuleEndpoint
4258        );
4259        assert!(committed_endpoints(&forwarding).is_empty());
4260        assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4261        assert!(matches!(
4262            pending.receiver.try_recv(),
4263            Err(oneshot::error::TryRecvError::Closed)
4264        ));
4265    }
4266
4267    #[test]
4268    fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
4269        // A candidate whose connection drops leaves the incumbent routable.
4270        let fixture = swap_fixture();
4271        let forwarding = &fixture.forwarding;
4272        assert!(forwarding
4273            .cleanup_connection(fixture.candidate_connection)
4274            .unwrap()
4275            .is_empty());
4276        assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
4277        let (client_connection, client_sink, _client_rx) = client(200);
4278        assert_eq!(
4279            forwarding
4280                .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4281                .unwrap()
4282                .endpoint,
4283            fixture.incumbent
4284        );
4285
4286        // After a cutover, the incumbent's teardown releases its own routes and
4287        // leaves the promoted candidate in place.
4288        let fixture = swap_fixture();
4289        let forwarding = &fixture.forwarding;
4290        let (bound_client, bound_sink, _bound_rx) = client(200);
4291        let bound = forwarding
4292            .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4293            .unwrap();
4294        forwarding
4295            .complete_pending_relay(
4296                fixture.incumbent_connection,
4297                bound.corr,
4298                RouteBindRelayOutcome::Accepted,
4299            )
4300            .unwrap();
4301        forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4302        let released = forwarding
4303            .cleanup_connection(fixture.incumbent_connection)
4304            .unwrap();
4305        assert_eq!(released.len(), 1);
4306        assert_eq!(released[0].connection_id, bound_client);
4307        assert!(forwarding
4308            .read_inner()
4309            .unwrap()
4310            .superseded_endpoints
4311            .is_empty());
4312        assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4313        let (next_client, next_sink, _next_rx) = client(201);
4314        assert_eq!(
4315            forwarding
4316                .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
4317                .unwrap()
4318                .endpoint,
4319            fixture.candidate
4320        );
4321    }
4322}