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