Skip to main content

subc_daemon/
forwarding.rs

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