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