Skip to main content

subc_daemon/
forwarding.rs

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