Skip to main content

subc_daemon/
forwarding.rs

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