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