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
158pub(crate) fn 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 if let Some(scope) = &route.scope {
3073 scope.release_bind();
3074 }
3075 inner.module_to_client.remove(&ModuleRouteKey {
3076 endpoint: route.module_endpoint,
3077 channel: route.module_channel,
3078 });
3079 inner.operator_confirms.route_closed(ModuleRouteKey {
3080 endpoint: route.module_endpoint,
3081 channel: route.module_channel,
3082 });
3083 inner.status.remove(&(client_key, expected_epoch));
3084 RouteRelease::Removed(GoodbyeTarget {
3085 connection_id: route.module_endpoint.connection_id,
3086 sink: route.module_sink.clone(),
3087 negotiated_ver: route.module_negotiated_ver,
3088 channel: route.module_channel,
3089 epoch: route.module_epoch,
3090 kind: GoodbyeTargetKind::Module,
3091 module_id: Some(route.module_id.clone()),
3092 })
3093}
3094
3095fn release_module_route_locked(
3096 inner: &mut ForwardingInner,
3097 module_key: ModuleRouteKey,
3098 expected_epoch: u32,
3099) -> RouteRelease {
3100 let Some(route) = inner.module_to_client.get(&module_key) else {
3101 return RouteRelease::Absent;
3102 };
3103 if route.module_epoch != expected_epoch {
3104 return RouteRelease::Stale;
3105 }
3106 let route = inner
3107 .module_to_client
3108 .remove(&module_key)
3109 .expect("route checked under the same forwarding lock");
3110 inner.operator_confirms.route_closed(module_key);
3111 route.flow.close();
3112 if let Some(scope) = &route.scope {
3113 scope.release_bind();
3114 }
3115 let client_key = ClientRouteKey {
3116 connection_id: route.client_connection_id,
3117 channel: route.client_channel,
3118 };
3119 inner.client_to_module.remove(&client_key);
3120 inner.status.remove(&(client_key, route.client_epoch));
3121 RouteRelease::Removed(GoodbyeTarget {
3122 connection_id: route.client_connection_id,
3123 sink: route.client_sink.clone(),
3124 negotiated_ver: route.client_negotiated_ver,
3125 channel: route.client_channel,
3126 epoch: route.client_epoch,
3127 kind: GoodbyeTargetKind::Client,
3128 module_id: Some(route.module_id.clone()),
3129 })
3130}
3131
3132fn commit_route_locked(
3133 inner: &mut ForwardingInner,
3134 pending: PendingRouteBindRelayEntry,
3135) -> Result<Option<GoodbyeTarget>, ForwardingError> {
3136 let reservation = pending.reservation;
3137 if inner
3138 .closing_connections
3139 .contains(&reservation.client_key.connection_id)
3140 {
3141 return Err(ForwardingError::ConnectionClosing {
3142 connection_id: reservation.client_key.connection_id,
3143 });
3144 }
3145 let module_id = inner
3146 .module_id_by_endpoint
3147 .get(&reservation.module_key.endpoint)
3148 .cloned()
3149 .ok_or(ForwardingError::StaleModuleEndpoint)?;
3150 if inner
3151 .draining_endpoints
3152 .contains_key(&reservation.module_key.endpoint)
3153 {
3154 return Err(ForwardingError::ModuleReloading { module_id });
3155 }
3156 if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
3157 || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
3158 {
3159 return Err(ForwardingError::UnknownReservation {
3160 client_channel: reservation.client_key.channel,
3161 module_channel: reservation.module_key.channel,
3162 });
3163 }
3164 let module = inner
3165 .modules_by_id
3166 .get(&module_id)
3167 .filter(|module| module.endpoint == reservation.module_key.endpoint)
3168 .cloned()
3169 .ok_or(ForwardingError::StaleModuleEndpoint)?;
3170 let binding = Arc::new(RouteBinding {
3171 client_connection_id: reservation.client_key.connection_id,
3172 client_sink: pending.client_sink,
3173 client_negotiated_ver: pending.client_negotiated_ver,
3174 client_channel: reservation.client_key.channel,
3175 client_epoch: reservation.client_epoch,
3176 module_id,
3177 module_endpoint: reservation.module_key.endpoint,
3178 module_sink: module.sink,
3179 module_negotiated_ver: module.negotiated_ver,
3180 module_channel: reservation.module_key.channel,
3181 module_epoch: reservation.module_epoch,
3182 principal: pending.principal,
3183 project_root: reservation.project_root.clone(),
3184 bound_at: Instant::now(),
3185 flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
3186 scope: pending.scope,
3187 });
3188 inner
3189 .client_to_module
3190 .insert(reservation.client_key, Arc::clone(&binding));
3191 inner
3192 .module_to_client
3193 .insert(reservation.module_key, Arc::clone(&binding));
3194 let previous_published = inner
3195 .last_published_epoch
3196 .insert(reservation.client_key, reservation.client_epoch);
3197
3198 let client_writer_closed = pending.client_permit.send(pending.route_open_frame);
3203 if client_writer_closed {
3204 let abandoned = abandoned_route_target(inner, &reservation);
3208 if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
3209 route.flow.close();
3210 }
3211 inner.module_to_client.remove(&reservation.module_key);
3212 inner
3213 .status
3214 .remove(&(reservation.client_key, reservation.client_epoch));
3215 match previous_published {
3216 Some(epoch) => {
3217 inner
3218 .last_published_epoch
3219 .insert(reservation.client_key, epoch);
3220 }
3221 None => {
3222 inner.last_published_epoch.remove(&reservation.client_key);
3223 }
3224 }
3225 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
3226 "client egress closed during route publication".to_string(),
3227 ));
3228 return Ok(abandoned);
3229 }
3230
3231 if let Some(scope) = &binding.scope {
3232 scope.commit_presence();
3233 }
3234 let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
3235 Ok(None)
3236}
3237
3238fn module_connection_for_endpoint_locked(
3246 inner: &ForwardingInner,
3247 endpoint: ModuleEndpointId,
3248) -> Option<&ModuleConnection> {
3249 let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
3250 inner
3251 .modules_by_id
3252 .get(module_id)
3253 .filter(|module| module.endpoint == endpoint)
3254 .or_else(|| {
3255 inner
3256 .candidates_by_id
3257 .get(module_id)
3258 .filter(|module| module.endpoint == endpoint)
3259 })
3260 .or_else(|| inner.superseded_endpoints.get(&endpoint))
3261}
3262
3263fn abandoned_route_target(
3264 inner: &ForwardingInner,
3265 reservation: &RouteReservation,
3266) -> Option<GoodbyeTarget> {
3267 let module_id = inner
3268 .module_id_by_endpoint
3269 .get(&reservation.module_key.endpoint)?;
3270 let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
3271 (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
3272 connection_id: module.endpoint.connection_id,
3273 sink: module.sink.clone(),
3274 negotiated_ver: module.negotiated_ver,
3275 channel: reservation.module_key.channel,
3276 epoch: reservation.module_epoch,
3277 kind: GoodbyeTargetKind::Module,
3278 module_id: Some(module_id.clone()),
3279 })
3280}
3281
3282fn enqueue_hello_ack_locked(
3287 sink: &FrameSink,
3288 connection_id: ConnectionId,
3289 hello_ack: Option<Frame>,
3290) -> Result<(), ForwardingError> {
3291 let Some(hello_ack) = hello_ack else {
3292 return Ok(());
3293 };
3294 sink.try_send(hello_ack)
3295 .map_err(|_| ForwardingError::ModuleEgressUnavailable { connection_id })
3296}
3297
3298fn remove_module_connection_locked(
3299 inner: &mut ForwardingInner,
3300 endpoint: ModuleEndpointId,
3301) -> ConnectionCleanup {
3302 inner.operator_confirms.module_closed(endpoint);
3305 inner.draining_endpoints.remove(&endpoint);
3306 inner.daemon_drain_results.remove(&endpoint);
3307 let module_id = inner.module_id_by_endpoint.remove(&endpoint);
3308 if let Some(module_id) = module_id.as_ref() {
3309 if inner
3310 .modules_by_id
3311 .get(module_id)
3312 .is_some_and(|module| module.endpoint == endpoint)
3313 {
3314 inner.modules_by_id.remove(module_id);
3315 }
3316 if inner
3317 .candidates_by_id
3318 .get(module_id)
3319 .is_some_and(|module| module.endpoint == endpoint)
3320 {
3321 inner.candidates_by_id.remove(module_id);
3322 }
3323 }
3324 inner.superseded_endpoints.remove(&endpoint);
3325 inner.endpoint_by_connection.remove(&endpoint.connection_id);
3326 inner.next_module_channel.remove(&endpoint);
3327 inner.next_control_corr.remove(&endpoint);
3328 inner
3329 .health_probe_tombstones
3330 .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
3331 inner
3332 .module_slot_epochs
3333 .retain(|key, _| key.endpoint != endpoint);
3334 let reserved_module_keys: Vec<ModuleRouteKey> = inner
3335 .reserved_module
3336 .keys()
3337 .filter(|module_key| module_key.endpoint == endpoint)
3338 .copied()
3339 .collect();
3340 for module_key in reserved_module_keys {
3341 if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
3342 release_reserved_route_locked(inner, client_key, module_key);
3343 }
3344 }
3345
3346 let pending_keys: Vec<_> = inner
3347 .pending_relays
3348 .keys()
3349 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
3350 .copied()
3351 .collect();
3352 let pending: Vec<_> = pending_keys
3353 .into_iter()
3354 .filter_map(|key| {
3355 inner
3356 .pending_relays
3357 .remove(&key)
3358 .map(|pending| (key, pending))
3359 })
3360 .collect();
3361 let abandoned_relays = u32::try_from(pending.len()).unwrap_or(u32::MAX);
3362 for (key, pending) in pending {
3363 if !pending.relay_enqueued {
3367 inner.relays_released_before_enqueue.insert(key);
3368 }
3369 let module_label = module_id.as_deref().unwrap_or("unknown");
3370 let _ = pending
3371 .sender
3372 .send(RouteBindRelayOutcome::ModuleGone(format!(
3373 "module '{module_label}' connection closed during route.bind relay"
3374 )));
3375 }
3376
3377 let pending_control_keys: Vec<_> = inner
3378 .pending_control_rpcs
3379 .keys()
3380 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
3381 .copied()
3382 .collect();
3383 let pending_control: Vec<_> = pending_control_keys
3384 .into_iter()
3385 .filter_map(|key| inner.pending_control_rpcs.remove(&key))
3386 .collect();
3387 for pending in pending_control {
3388 let module_label = module_id.as_deref().unwrap_or("unknown");
3389 let _ = pending
3390 .sender
3391 .send(ModuleControlRpcOutcome::ModuleGone(format!(
3392 "module '{module_label}' connection closed during module-control RPC"
3393 )));
3394 }
3395
3396 let module_routes = inner
3397 .module_to_client
3398 .iter()
3399 .filter(|(module_key, _)| module_key.endpoint == endpoint)
3400 .map(|(module_key, route)| (*module_key, route.module_epoch))
3401 .collect::<Vec<_>>();
3402 let mut released = Vec::with_capacity(module_routes.len());
3403 for (module_key, epoch) in module_routes {
3404 if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
3405 {
3406 released.push(target);
3407 }
3408 }
3409 ConnectionCleanup {
3410 released,
3411 abandoned_relays,
3412 }
3413}
3414
3415#[derive(Debug, Clone, Copy)]
3416struct RequestCredit {
3417 subscription: bool,
3418 excluded_from_drain: bool,
3419}
3420
3421#[derive(Debug, Default)]
3422struct CreditLedger {
3423 by_corr: HashMap<u64, Vec<RequestCredit>>,
3424}
3425
3426impl CreditLedger {
3427 fn acquire(&mut self, corr: u64, subscription: bool) {
3428 self.by_corr.entry(corr).or_default().push(RequestCredit {
3429 subscription,
3430 excluded_from_drain: false,
3431 });
3432 }
3433
3434 fn release(&mut self, corr: u64) -> bool {
3435 let Some(credits) = self.by_corr.get_mut(&corr) else {
3436 return false;
3437 };
3438 let released = credits.pop().is_some();
3439 if credits.is_empty() {
3440 self.by_corr.remove(&corr);
3441 }
3442 released
3443 }
3444
3445 fn capture_subscription_exclusions(&mut self) -> u32 {
3446 let mut excluded = 0u32;
3447 for credit in self.by_corr.values_mut().flatten() {
3448 if credit.subscription && !credit.excluded_from_drain {
3449 credit.excluded_from_drain = true;
3450 excluded = excluded.saturating_add(1);
3451 }
3452 }
3453 excluded
3454 }
3455
3456 #[cfg(test)]
3457 fn in_flight(&self) -> usize {
3458 self.by_corr.values().map(Vec::len).sum()
3459 }
3460
3461 fn drain_in_flight(&self) -> usize {
3462 self.by_corr
3463 .values()
3464 .flatten()
3465 .filter(|credit| !credit.excluded_from_drain)
3466 .count()
3467 }
3468
3469 fn drain_held_corrs(&self) -> Vec<u64> {
3472 let mut corrs = self
3473 .by_corr
3474 .iter()
3475 .flat_map(|(corr, credits)| {
3476 credits
3477 .iter()
3478 .filter(|credit| !credit.excluded_from_drain)
3479 .map(move |_| *corr)
3480 })
3481 .collect::<Vec<_>>();
3482 corrs.sort_unstable();
3483 corrs
3484 }
3485}
3486
3487#[derive(Debug, Default)]
3488struct ChannelFlowState {
3489 closed: bool,
3490 credits: CreditLedger,
3491}
3492
3493#[derive(Debug)]
3495pub(crate) struct ChannelFlow {
3496 sem: Semaphore,
3497 window: usize,
3498 state: Mutex<ChannelFlowState>,
3499}
3500
3501impl ChannelFlow {
3502 pub(crate) fn new(window: usize) -> Self {
3503 debug_assert!(window > 0, "flow-control window must be non-zero");
3504 Self {
3505 sem: Semaphore::new(window),
3506 window,
3507 state: Mutex::new(ChannelFlowState::default()),
3508 }
3509 }
3510
3511 #[cfg(test)]
3512 pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
3513 self.acquire_tagged(0, false).await
3514 }
3515
3516 pub(crate) async fn acquire_tagged(
3517 &self,
3518 corr: u64,
3519 subscription: bool,
3520 ) -> Result<(), ChannelFlowClosed> {
3521 let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
3522 let mut state = self
3523 .state
3524 .lock()
3525 .unwrap_or_else(|poisoned| poisoned.into_inner());
3526 if state.closed {
3527 return Err(ChannelFlowClosed);
3528 }
3529 state.credits.acquire(corr, subscription);
3530 permit.forget();
3531 Ok(())
3532 }
3533
3534 #[cfg(test)]
3535 pub(crate) fn release(&self) {
3536 self.release_corr(0);
3537 }
3538
3539 pub(crate) fn release_corr(&self, corr: u64) {
3540 let released = self
3541 .state
3542 .lock()
3543 .unwrap_or_else(|poisoned| poisoned.into_inner())
3544 .credits
3545 .release(corr);
3546 if !released {
3547 warn!(
3551 window = self.window,
3552 available = self.sem.available_permits(),
3553 "flow-control over-release ignored"
3554 );
3555 return;
3556 }
3557 if !self.sem.is_closed() {
3558 self.sem.add_permits(1);
3559 }
3560 }
3561
3562 #[cfg(test)]
3563 pub(crate) fn in_flight(&self) -> usize {
3564 self.state
3565 .lock()
3566 .unwrap_or_else(|poisoned| poisoned.into_inner())
3567 .credits
3568 .in_flight()
3569 }
3570
3571 pub(crate) fn drain_in_flight(&self) -> usize {
3572 self.state
3573 .lock()
3574 .unwrap_or_else(|poisoned| poisoned.into_inner())
3575 .credits
3576 .drain_in_flight()
3577 }
3578
3579 pub(crate) fn drain_held_corrs(&self) -> Vec<u64> {
3580 self.state
3581 .lock()
3582 .unwrap_or_else(|poisoned| poisoned.into_inner())
3583 .credits
3584 .drain_held_corrs()
3585 }
3586
3587 #[cfg(test)]
3588 pub(crate) fn available_permits(&self) -> usize {
3589 self.sem.available_permits()
3590 }
3591
3592 pub(crate) fn begin_drain(&self) -> u32 {
3593 let mut state = self
3594 .state
3595 .lock()
3596 .unwrap_or_else(|poisoned| poisoned.into_inner());
3597 state.closed = true;
3598 self.sem.close();
3599 state.credits.capture_subscription_exclusions()
3600 }
3601
3602 pub(crate) fn close(&self) {
3603 self.state
3604 .lock()
3605 .unwrap_or_else(|poisoned| poisoned.into_inner())
3606 .closed = true;
3607 self.sem.close();
3608 }
3609}
3610
3611#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3612pub(crate) struct ChannelFlowClosed;
3613
3614impl fmt::Display for ChannelFlowClosed {
3615 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3616 write!(f, "flow-control window closed")
3617 }
3618}
3619
3620impl Error for ChannelFlowClosed {}
3621
3622fn window_for(concurrency: &Concurrency) -> usize {
3623 match concurrency {
3624 Concurrency::Serial => 1,
3625 Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
3626 Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
3627 }
3628}
3629
3630#[derive(Debug, Clone, PartialEq, Eq)]
3631pub enum ForwardingError {
3632 ConnectionRoleConflict {
3635 connection_id: ConnectionId,
3636 },
3637 NoModuleConnection,
3638 ModuleReloading {
3639 module_id: String,
3640 },
3641 StaleModuleEndpoint,
3642 UnknownReservation {
3643 client_channel: u16,
3644 module_channel: u16,
3645 },
3646 ClientRouteChannelExhausted {
3647 connection_id: ConnectionId,
3648 },
3649 ModuleRouteChannelExhausted {
3650 endpoint: ModuleEndpointId,
3651 },
3652 RelayCorrelationExhausted,
3653 ConnectionClosing {
3654 connection_id: ConnectionId,
3655 },
3656 ClientEgressClosed {
3657 connection_id: ConnectionId,
3658 },
3659 RouteOpenBuild(String),
3660 CandidateSlotOccupied {
3662 module_id: String,
3663 },
3664 ModuleEgressUnavailable {
3667 connection_id: ConnectionId,
3668 },
3669 Poisoned,
3670}
3671
3672impl fmt::Display for ForwardingError {
3673 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3674 match self {
3675 Self::ConnectionRoleConflict { connection_id } => write!(f, "connection {} already holds an incompatible role", connection_id.get()),
3676 Self::NoModuleConnection => write!(f, "no module connection is registered"),
3677 Self::ModuleReloading { module_id } => {
3678 write!(f, "module_id '{module_id}' is reloading")
3679 }
3680 Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
3681 Self::UnknownReservation {
3682 client_channel,
3683 module_channel,
3684 } => write!(
3685 f,
3686 "route reservation client channel {client_channel} / module channel {module_channel} was not found"
3687 ),
3688 Self::ClientRouteChannelExhausted { connection_id } => write!(
3689 f,
3690 "no client route channels are available for connection {}",
3691 connection_id.get()
3692 ),
3693 Self::ModuleRouteChannelExhausted { endpoint } => write!(
3694 f,
3695 "no module route channels are available for endpoint generation {} on connection {}",
3696 endpoint.generation,
3697 endpoint.connection_id.get()
3698 ),
3699 Self::RelayCorrelationExhausted => {
3700 write!(f, "module control correlation ids are exhausted")
3701 }
3702 Self::ConnectionClosing { connection_id } => write!(
3703 f,
3704 "connection {} is closing and cannot accept route allocation",
3705 connection_id.get()
3706 ),
3707 Self::ClientEgressClosed { connection_id } => write!(
3708 f,
3709 "client connection {} egress is closed",
3710 connection_id.get()
3711 ),
3712 Self::RouteOpenBuild(message) => {
3713 write!(f, "failed to prebuild route.open response: {message}")
3714 }
3715 Self::CandidateSlotOccupied { module_id } => write!(
3716 f,
3717 "module_id '{module_id}' already has a swap candidate registered"
3718 ),
3719 Self::ModuleEgressUnavailable { connection_id } => write!(
3720 f,
3721 "module connection {} egress is unavailable; HELLO_ACK could not be queued",
3722 connection_id.get()
3723 ),
3724 Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
3725 }
3726 }
3727}
3728
3729impl Error for ForwardingError {}
3730
3731#[cfg(test)]
3732mod tests {
3733 use std::time::Duration;
3734
3735 use super::*;
3736 use tokio::sync::mpsc;
3737
3738 #[test]
3739 fn ordinary_long_running_request_is_not_excluded_from_drain() {
3740 let mut ledger = CreditLedger::default();
3741 ledger.acquire(1, false);
3742
3743 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3744 assert_eq!(ledger.drain_in_flight(), 1);
3745 }
3746
3747 #[test]
3748 fn bit_set_subscription_is_excluded_and_counted() {
3749 let mut ledger = CreditLedger::default();
3750 ledger.acquire(1, true);
3751
3752 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3753 assert_eq!(ledger.drain_in_flight(), 0);
3754 }
3755
3756 #[test]
3757 fn subscription_opened_after_drain_snapshot_is_not_excluded() {
3758 let mut ledger = CreditLedger::default();
3759 ledger.acquire(1, true);
3760 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3761
3762 ledger.acquire(2, true);
3763
3764 assert_eq!(ledger.drain_in_flight(), 1);
3765 }
3766
3767 #[test]
3768 fn drain_with_no_subscriptions_reports_zero_excluded() {
3769 let mut ledger = CreditLedger::default();
3770 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3771 }
3772
3773 fn test_hello_ack(corr: u64) -> Frame {
3774 Frame::build(
3775 FrameType::HelloAck,
3776 Flags::new(false, Priority::Passive, false),
3777 0,
3778 0,
3779 corr,
3780 Vec::new(),
3781 )
3782 .unwrap()
3783 }
3784
3785 #[test]
3789 fn acked_registration_that_cannot_queue_its_hello_ack_inserts_nothing() {
3790 let forwarding = ForwardingTable::default();
3791
3792 let (closed_tx, closed_rx) = mpsc::channel(8);
3793 drop(closed_rx);
3794 let closed = ConnectionId::new(1);
3795 assert_eq!(
3796 forwarding.register_module_connection_acked(
3797 closed,
3798 "closed".to_string(),
3799 2,
3800 Concurrency::ModuleManaged,
3801 FrameSink::new(closed_tx),
3802 test_hello_ack(1),
3803 ),
3804 Err(ForwardingError::ModuleEgressUnavailable {
3805 connection_id: closed
3806 })
3807 );
3808
3809 let (full_tx, _full_rx) = mpsc::channel(1);
3810 let full_sink = FrameSink::new(full_tx);
3811 full_sink.try_send(test_hello_ack(99)).unwrap();
3812 let full = ConnectionId::new(2);
3813 assert_eq!(
3814 forwarding.register_module_connection_acked(
3815 full,
3816 "full".to_string(),
3817 2,
3818 Concurrency::ModuleManaged,
3819 full_sink.clone(),
3820 test_hello_ack(2),
3821 ),
3822 Err(ForwardingError::ModuleEgressUnavailable {
3823 connection_id: full
3824 })
3825 );
3826 assert_eq!(
3827 forwarding.register_candidate_module_connection_acked(
3828 full,
3829 "full".to_string(),
3830 2,
3831 Concurrency::ModuleManaged,
3832 full_sink,
3833 test_hello_ack(3),
3834 ),
3835 Err(ForwardingError::ModuleEgressUnavailable {
3836 connection_id: full
3837 })
3838 );
3839
3840 for (connection, module_id) in [(closed, "closed"), (full, "full")] {
3841 assert_eq!(
3842 forwarding
3843 .module_endpoint_for_connection(connection)
3844 .unwrap(),
3845 None
3846 );
3847 let (client_tx, _client_rx) = mpsc::channel(8);
3848 assert_eq!(
3849 forwarding
3850 .begin_route_bind_relay_for_test(
3851 ConnectionId::new(50),
3852 FrameSink::new(client_tx),
3853 1,
3854 module_id,
3855 )
3856 .err(),
3857 Some(ForwardingError::NoModuleConnection)
3858 );
3859 }
3860 assert!(forwarding.read_inner().unwrap().candidates_by_id.is_empty());
3861 }
3862
3863 #[test]
3866 fn acked_registration_queues_the_hello_ack_first() {
3867 let forwarding = ForwardingTable::default();
3868 let (active_tx, mut active_rx) = mpsc::channel(8);
3869 forwarding
3870 .register_module_connection_acked(
3871 ConnectionId::new(1),
3872 "acked".to_string(),
3873 2,
3874 Concurrency::ModuleManaged,
3875 FrameSink::new(active_tx),
3876 test_hello_ack(11),
3877 )
3878 .unwrap();
3879 let (candidate_tx, mut candidate_rx) = mpsc::channel(8);
3880 forwarding
3881 .register_candidate_module_connection_acked(
3882 ConnectionId::new(2),
3883 "acked".to_string(),
3884 2,
3885 Concurrency::ModuleManaged,
3886 FrameSink::new(candidate_tx),
3887 test_hello_ack(12),
3888 )
3889 .unwrap();
3890
3891 let active_first = active_rx.try_recv().unwrap().frame;
3892 assert_eq!(active_first.header.ty, FrameType::HelloAck);
3893 assert_eq!(active_first.header.corr, 11);
3894 let candidate_first = candidate_rx.try_recv().unwrap().frame;
3895 assert_eq!(candidate_first.header.ty, FrameType::HelloAck);
3896 assert_eq!(candidate_first.header.corr, 12);
3897 }
3898
3899 #[test]
3900 fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
3901 let forwarding = ForwardingTable::default();
3902 let module_connection = ConnectionId::new(10);
3903 let exhausted_client = ConnectionId::new(20);
3904 let second_client = ConnectionId::new(30);
3905 let (module_tx, _module_rx) = mpsc::channel(1);
3906 let endpoint = forwarding
3907 .register_module_connection(
3908 module_connection,
3909 "route-limit-provider".to_string(),
3910 1,
3911 Concurrency::ModuleManaged,
3912 FrameSink::new(module_tx),
3913 )
3914 .unwrap();
3915
3916 {
3917 let mut inner = forwarding.inner.write().unwrap();
3918 for channel in 1..=u16::MAX {
3919 inner.reserved_client.insert(
3920 ClientRouteKey {
3921 connection_id: exhausted_client,
3922 channel,
3923 },
3924 ModuleRouteKey {
3925 endpoint,
3926 channel: 1,
3927 },
3928 );
3929 }
3930 }
3931
3932 let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
3933 let err = forwarding
3934 .begin_route_bind_relay_for_test(
3935 exhausted_client,
3936 FrameSink::new(exhausted_tx),
3937 1,
3938 "route-limit-provider",
3939 )
3940 .unwrap_err();
3941 assert!(matches!(
3942 err,
3943 ForwardingError::ClientRouteChannelExhausted { connection_id }
3944 if connection_id == exhausted_client
3945 ));
3946
3947 let (second_tx, _second_rx) = mpsc::channel(1);
3948 let pending = forwarding
3949 .begin_route_bind_relay_for_test(
3950 second_client,
3951 FrameSink::new(second_tx),
3952 2,
3953 "route-limit-provider",
3954 )
3955 .unwrap();
3956 assert_eq!(pending.client_channel, 1);
3957 }
3958
3959 #[test]
3960 fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
3961 let forwarding = ForwardingTable::default();
3962 let module_connection = ConnectionId::new(40);
3963 let client = ConnectionId::new(50);
3964 let (module_tx, _module_rx) = mpsc::channel(1);
3965 forwarding
3966 .register_module_connection(
3967 module_connection,
3968 "slot-reuse-provider".to_string(),
3969 1,
3970 Concurrency::ModuleManaged,
3971 FrameSink::new(module_tx),
3972 )
3973 .unwrap();
3974
3975 let (client_tx, _client_rx) = mpsc::channel(1);
3976 let client_sink = FrameSink::new(client_tx);
3977 let mut wrapped_channel = None;
3978 for index in 0..=usize::from(u16::MAX) {
3979 let pending = forwarding
3980 .begin_route_bind_relay_for_test(
3981 client,
3982 client_sink.clone(),
3983 index as u64 + 1,
3984 "slot-reuse-provider",
3985 )
3986 .unwrap();
3987 if index == usize::from(u16::MAX) {
3988 wrapped_channel = Some(pending.module_channel);
3989 }
3990 forwarding
3991 .abort_pending_relay(
3992 pending.endpoint,
3993 pending.corr,
3994 RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
3995 )
3996 .unwrap();
3997 }
3998
3999 assert_eq!(wrapped_channel, Some(1));
4000 }
4001
4002 #[test]
4003 fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
4004 let forwarding = ForwardingTable::default();
4005 let client = ConnectionId::new(60);
4006 forwarding
4007 .inner
4008 .write()
4009 .unwrap()
4010 .next_client_channel
4011 .insert(client, 41);
4012
4013 let released = forwarding.cleanup_connection(client).unwrap();
4014
4015 assert!(released.is_empty());
4016 assert!(!forwarding
4017 .inner
4018 .read()
4019 .unwrap()
4020 .next_client_channel
4021 .contains_key(&client));
4022 }
4023
4024 #[test]
4025 fn stale_module_cleanup_preserves_fast_reconnect_successor() {
4026 let forwarding = ForwardingTable::default();
4027 let module_id = "fast-reconnect-provider";
4028 let first_connection = ConnectionId::new(70);
4029 let second_connection = ConnectionId::new(80);
4030 let (first_tx, _first_rx) = mpsc::channel(1);
4031 let first_endpoint = forwarding
4032 .register_module_connection(
4033 first_connection,
4034 module_id.to_string(),
4035 1,
4036 Concurrency::ModuleManaged,
4037 FrameSink::new(first_tx),
4038 )
4039 .unwrap();
4040 let (second_tx, _second_rx) = mpsc::channel(1);
4041 let second_endpoint = forwarding
4042 .register_module_connection(
4043 second_connection,
4044 module_id.to_string(),
4045 1,
4046 Concurrency::ModuleManaged,
4047 FrameSink::new(second_tx),
4048 )
4049 .unwrap();
4050 assert_ne!(first_endpoint, second_endpoint);
4051
4052 let released = forwarding.cleanup_connection(first_connection).unwrap();
4053
4054 assert!(released.is_empty());
4055 assert_eq!(
4056 forwarding
4057 .inner
4058 .read()
4059 .unwrap()
4060 .modules_by_id
4061 .get(module_id)
4062 .map(|module| module.endpoint),
4063 Some(second_endpoint)
4064 );
4065 assert!(forwarding.has_live_module_connection(module_id).unwrap());
4066 let control_rpc = forwarding
4067 .begin_module_control_rpc_for(
4068 module_id,
4069 "health.check",
4070 Instant::now() + Duration::from_secs(1),
4071 )
4072 .unwrap();
4073 assert_eq!(control_rpc.endpoint, second_endpoint);
4074 }
4075
4076 fn route_fixture(
4077 module_id: &str,
4078 ) -> (
4079 ForwardingTable,
4080 ConnectionId,
4081 ModuleEndpointId,
4082 ConnectionId,
4083 FrameSink,
4084 mpsc::Receiver<crate::router::OutboundFrame>,
4085 ) {
4086 let forwarding = ForwardingTable::default();
4087 let module_connection = ConnectionId::new(100);
4088 let client_connection = ConnectionId::new(200);
4089 let (module_tx, _module_rx) = mpsc::channel(8);
4090 let endpoint = forwarding
4091 .register_module_connection(
4092 module_connection,
4093 module_id.to_string(),
4094 2,
4095 Concurrency::ModuleManaged,
4096 FrameSink::new(module_tx),
4097 )
4098 .unwrap();
4099 let (client_tx, client_rx) = mpsc::channel(8);
4100 (
4101 forwarding,
4102 module_connection,
4103 endpoint,
4104 client_connection,
4105 FrameSink::new(client_tx),
4106 client_rx,
4107 )
4108 }
4109
4110 #[test]
4111 #[cfg(unix)]
4112 fn daemon_drain_gates_current_and_racing_provider_registrations() {
4113 let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
4114 assert_eq!(
4115 forwarding.begin_daemon_drain().unwrap(),
4116 [("provider".to_string(), endpoint)]
4117 );
4118 assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
4119 assert!(matches!(
4120 forwarding.register_module_connection(
4121 ConnectionId::new(300),
4122 "late-provider".into(),
4123 2,
4124 Concurrency::ModuleManaged,
4125 sink,
4126 ),
4127 Err(ForwardingError::ConnectionClosing { .. })
4128 ));
4129 }
4130
4131 #[test]
4132 #[cfg(unix)]
4133 fn daemon_drain_results_are_endpoint_fenced_and_removed_on_cleanup() {
4134 let (forwarding, connection, endpoint, _, _, _) = route_fixture("provider");
4135 assert_eq!(
4136 forwarding
4137 .daemon_drain_result_for_connection(connection)
4138 .unwrap(),
4139 None
4140 );
4141 forwarding
4142 .record_daemon_drain_result(endpoint, true)
4143 .unwrap();
4144 forwarding.begin_daemon_drain().unwrap();
4145 assert_eq!(
4146 forwarding
4147 .daemon_drain_result_for_connection(connection)
4148 .unwrap(),
4149 Some(false)
4150 );
4151 let replacement = ModuleEndpointId {
4152 generation: endpoint.generation + 1,
4153 ..endpoint
4154 };
4155 forwarding
4156 .record_daemon_drain_result(replacement, true)
4157 .unwrap();
4158 assert_eq!(
4159 forwarding
4160 .daemon_drain_result_for_connection(connection)
4161 .unwrap(),
4162 Some(false)
4163 );
4164 forwarding
4165 .record_daemon_drain_result(endpoint, true)
4166 .unwrap();
4167 assert_eq!(
4168 forwarding
4169 .daemon_drain_result_for_connection(connection)
4170 .unwrap(),
4171 Some(true)
4172 );
4173 assert_eq!(
4174 forwarding
4175 .daemon_drain_result_for_connection(ConnectionId::new(300))
4176 .unwrap(),
4177 Some(false)
4178 );
4179 forwarding.cleanup_connection(connection).unwrap();
4180 forwarding
4181 .record_daemon_drain_result(endpoint, true)
4182 .unwrap();
4183 assert_eq!(
4184 forwarding
4185 .daemon_drain_result_for_connection(connection)
4186 .unwrap(),
4187 Some(false)
4188 );
4189 assert!(forwarding
4190 .read_inner()
4191 .unwrap()
4192 .daemon_drain_results
4193 .is_empty());
4194 }
4195
4196 #[test]
4197 #[cfg(unix)]
4198 fn late_bind_ack_during_daemon_drain_settles_without_leaking_or_closing_module() {
4199 let (forwarding, module_connection, endpoint, client_connection, sink, mut rx) =
4200 route_fixture("provider");
4201 let mut pending =
4202 begin_test_route(&forwarding, client_connection, sink.clone(), 1, "provider");
4203 assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4204 forwarding.begin_daemon_drain().unwrap();
4205 let completion = forwarding
4206 .complete_pending_relay(
4207 module_connection,
4208 pending.corr,
4209 RouteBindRelayOutcome::Accepted,
4210 )
4211 .expect("draining admission is not a fatal module error");
4212 assert!(completion.settled);
4213 let abandoned = completion
4214 .abandoned
4215 .expect("late accepted bind needs a route GOODBYE");
4216 assert_eq!(abandoned.channel, pending.module_channel);
4217 assert!(matches!(pending.receiver.try_recv().unwrap(),
4218 RouteBindRelayOutcome::Rejected(error) if error.code == "module_reloading"));
4219 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4220 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4221 assert!(
4222 rx.try_recv().is_err(),
4223 "no route.open may be published during drain"
4224 );
4225 assert_eq!(
4226 forwarding
4227 .module_endpoint_for_connection(module_connection)
4228 .unwrap(),
4229 Some(endpoint)
4230 );
4231 assert!(
4232 !forwarding
4233 .complete_pending_relay(
4234 module_connection,
4235 pending.corr,
4236 RouteBindRelayOutcome::Accepted,
4237 )
4238 .unwrap()
4239 .settled
4240 );
4241 }
4242
4243 fn test_ping(corr: u64) -> Frame {
4244 Frame::build(
4245 FrameType::Ping,
4246 Flags::new(false, Priority::Passive, false),
4247 0,
4248 0,
4249 corr,
4250 Vec::new(),
4251 )
4252 .unwrap()
4253 }
4254
4255 fn begin_test_route(
4256 forwarding: &ForwardingTable,
4257 client_connection: ConnectionId,
4258 client_sink: FrameSink,
4259 corr: u64,
4260 module_id: &str,
4261 ) -> PendingRouteBindRelay {
4262 forwarding
4263 .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
4264 .unwrap()
4265 }
4266
4267 #[test]
4268 fn module_registration_and_client_reservation_are_mutually_exclusive_in_both_orders() {
4269 for candidate in [false, true] {
4270 for register_first in [false, true] {
4271 let (forwarding, module_connection, _, client_connection, sink, _rx) =
4272 route_fixture("target");
4273 let register = || {
4274 if candidate {
4275 forwarding.register_candidate_module_connection(
4276 client_connection,
4277 "source".into(),
4278 2,
4279 Concurrency::ModuleManaged,
4280 sink.clone(),
4281 )
4282 } else {
4283 forwarding.register_module_connection(
4284 client_connection,
4285 "source".into(),
4286 2,
4287 Concurrency::ModuleManaged,
4288 sink.clone(),
4289 )
4290 }
4291 };
4292 if register_first {
4293 register().unwrap();
4294 assert!(
4295 forwarding
4296 .begin_route_bind_relay_for_test(
4297 client_connection,
4298 sink.clone(),
4299 1,
4300 "target",
4301 )
4302 .is_err(),
4303 "a deferred route.open cannot reserve after HELLO"
4304 );
4305 assert!(forwarding
4306 .cleanup_connection(client_connection)
4307 .unwrap()
4308 .is_empty());
4309 } else {
4310 let pending =
4311 begin_test_route(&forwarding, client_connection, sink.clone(), 1, "target");
4312 assert!(
4313 register().is_err(),
4314 "HELLO cannot register after a route reservation"
4315 );
4316 forwarding
4317 .complete_pending_relay(
4318 module_connection,
4319 pending.corr,
4320 RouteBindRelayOutcome::Accepted,
4321 )
4322 .unwrap();
4323 assert_eq!(
4324 forwarding
4325 .cleanup_connection(client_connection)
4326 .unwrap()
4327 .len(),
4328 1
4329 );
4330 }
4331 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4332 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4333 }
4334 }
4335 }
4336
4337 #[tokio::test]
4343 async fn pending_route_open_completes_behind_queued_data_frames() {
4344 assert_eq!(
4345 crate::server::MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
4346 8,
4347 "the per-connection pending route.open limit is its own constant"
4348 );
4349 let (forwarding, module_connection, _endpoint, client, _unused_sink, _unused_rx) =
4350 route_fixture("open-behind-data");
4351 let (sink, mut client_rx) = crate::server::connection_egress();
4352 const DATA_FRAMES: usize = 1_000;
4353 let data = |corr: u64| {
4354 Frame::build(
4355 FrameType::StreamData,
4356 Flags::new(false, Priority::Interactive, false),
4357 9,
4358 1,
4359 corr,
4360 vec![b'x'; 200],
4361 )
4362 .unwrap()
4363 };
4364 for corr in 0..DATA_FRAMES as u64 {
4365 sink.try_send(data(corr)).unwrap();
4366 }
4367 let data_bytes = DATA_FRAMES * (subc_protocol::HEADER_LEN + 200);
4368 assert_eq!(sink.backlog().queued_bytes, data_bytes);
4369
4370 let pending = tokio::time::timeout(
4371 Duration::from_secs(5),
4372 forwarding.begin_route_bind_relay_for(
4373 client,
4374 sink.clone(),
4375 subc_protocol::PROTOCOL_VERSION,
4376 4_242,
4377 "open-behind-data",
4378 Principal::Direct,
4379 None,
4380 None,
4381 Instant::now() + Duration::from_secs(60),
4382 ),
4383 )
4384 .await
4385 .expect("reserving the route.open slot must not wait behind data frames")
4386 .unwrap();
4387 forwarding
4388 .complete_pending_relay(
4389 module_connection,
4390 pending.corr,
4391 RouteBindRelayOutcome::Accepted,
4392 )
4393 .unwrap();
4394
4395 let backlog = sink.backlog();
4396 assert_eq!(backlog.queued_frames, DATA_FRAMES + 1);
4397 assert!(
4398 backlog.queued_bytes > data_bytes,
4399 "the route.open response must be counted in queued bytes: {backlog:?}"
4400 );
4401 for corr in 0..DATA_FRAMES as u64 {
4402 assert_eq!(client_rx.recv().await.unwrap().header.corr, corr);
4403 }
4404 let open = client_rx.recv().await.unwrap();
4405 assert_eq!(open.header.corr, 4_242);
4406 assert_eq!(open.header.ty, FrameType::Response);
4407 drop(open);
4408 assert_eq!(sink.backlog().queued_bytes, 0);
4409 assert_eq!(sink.backlog().queued_frames, 0);
4410 }
4411
4412 #[tokio::test]
4417 async fn drain_holdouts_count_held_requests_and_name_the_connection() {
4418 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4419 route_fixture("holdouts");
4420 let mut bound = |corr| {
4421 let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
4422 forwarding
4423 .complete_pending_relay(
4424 module_connection,
4425 route.corr,
4426 RouteBindRelayOutcome::Accepted,
4427 )
4428 .unwrap();
4429 client_rx.try_recv().unwrap();
4430 match forwarding
4431 .lookup_data_route(client, route.client_channel, route.client_epoch)
4432 .unwrap()
4433 {
4434 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4435 other => panic!("expected live route, got {other:?}"),
4436 }
4437 };
4438 let holding = bound(61);
4439 let _idle = bound(62);
4440 holding.flow.acquire_tagged(7, false).await.unwrap();
4441 holding.flow.acquire_tagged(2, false).await.unwrap();
4442 holding.flow.acquire_tagged(3, true).await.unwrap();
4443 forwarding
4444 .begin_module_drain("holdouts", RouteCloseReason::Restart)
4445 .unwrap();
4446
4447 let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
4448 assert_eq!(
4449 holdouts,
4450 DrainHoldouts {
4451 requests: 2,
4452 routes: 1,
4453 total_routes: 2,
4454 top_connections: vec![(client.get(), 2)],
4455 held: vec![(holding.module_channel, 2), (holding.module_channel, 7)],
4458 }
4459 );
4460 }
4461
4462 #[test]
4463 fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
4464 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
4465 route_fixture("census");
4466 let pending = begin_test_route(&forwarding, client, sink, 1, "census");
4467 forwarding
4468 .complete_pending_relay(
4469 module_connection,
4470 pending.corr,
4471 RouteBindRelayOutcome::Accepted,
4472 )
4473 .unwrap();
4474
4475 let routes = forwarding.endpoint_routes(endpoint).unwrap();
4476 assert_eq!(routes.len(), 1);
4477 assert!(matches!(routes[0].principal, Principal::Direct));
4478 assert_eq!(routes[0].goodbye_target.connection_id, client);
4479 assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
4480 assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
4481 assert!(!routes[0].draining);
4482
4483 forwarding
4484 .begin_module_drain("census", RouteCloseReason::Restart)
4485 .unwrap();
4486 let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
4487 assert_eq!(draining_routes.len(), 1);
4488 assert!(draining_routes[0].draining);
4489 }
4490
4491 #[test]
4492 fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
4493 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
4494 let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
4495 assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
4496 forwarding
4497 .abort_pending_relay(
4498 first.endpoint,
4499 first.corr,
4500 RouteBindRelayOutcome::ModuleGone("abort".into()),
4501 )
4502 .unwrap();
4503 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
4504 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
4505
4506 let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
4507 assert_eq!(second.client_channel, first.client_channel);
4508 assert_eq!(second.module_channel, first.module_channel);
4509 assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
4510 }
4511
4512 #[test]
4513 fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
4514 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4515 route_fixture("epoch-release");
4516 let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
4517 forwarding
4518 .complete_pending_relay(
4519 module_connection,
4520 first.corr,
4521 RouteBindRelayOutcome::Accepted,
4522 )
4523 .unwrap();
4524 assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
4525 assert!(matches!(
4526 forwarding
4527 .release_client_route(client, first.client_channel, first.client_epoch)
4528 .unwrap(),
4529 RouteRelease::Removed(_)
4530 ));
4531 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
4532 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
4533
4534 let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
4535 forwarding
4536 .complete_pending_relay(
4537 module_connection,
4538 second.corr,
4539 RouteBindRelayOutcome::Accepted,
4540 )
4541 .unwrap();
4542 assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
4543 assert!(matches!(
4544 forwarding
4545 .release_client_route(client, second.client_channel, first.client_epoch)
4546 .unwrap(),
4547 RouteRelease::Stale
4548 ));
4549 assert!(!forwarding
4550 .cache_status(
4551 endpoint,
4552 second.module_channel,
4553 first.module_epoch,
4554 "stale".into(),
4555 )
4556 .unwrap());
4557 assert!(forwarding
4558 .cache_status(
4559 endpoint,
4560 second.module_channel,
4561 second.module_epoch,
4562 "current".into(),
4563 )
4564 .unwrap());
4565 match forwarding
4566 .route_poll_snapshot(client, second.client_channel, second.client_epoch)
4567 .unwrap()
4568 {
4569 RoutePollSnapshot::Bound { status, .. } => {
4570 assert_eq!(status.as_deref(), Some("current"));
4571 }
4572 RoutePollSnapshot::Absent => panic!("successor binding was removed"),
4573 }
4574 let counters = forwarding.counters().snapshot();
4575 assert_eq!(counters["route_released_epoch_fenced"], 1);
4576 assert_eq!(counters["route_release_stale_skipped"], 1);
4577 }
4578
4579 #[test]
4580 fn max_epoch_reservation_retires_only_that_slot() {
4581 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
4582 forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
4583 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
4584 let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
4585 assert_eq!(
4586 (final_use.client_channel, final_use.client_epoch),
4587 (7, u32::MAX)
4588 );
4589 assert_eq!(
4590 (final_use.module_channel, final_use.module_epoch),
4591 (9, u32::MAX)
4592 );
4593 forwarding
4594 .abort_pending_relay(
4595 endpoint,
4596 final_use.corr,
4597 RouteBindRelayOutcome::ModuleGone("abort".into()),
4598 )
4599 .unwrap();
4600 forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
4601 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
4602 let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
4603 assert_ne!(next.client_channel, 7);
4604 assert_ne!(next.module_channel, 9);
4605 assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
4606 }
4607
4608 #[test]
4609 fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
4610 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
4611 route_fixture("corr-shared");
4612 let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
4613 assert_eq!(bind.corr, 1);
4614 forwarding
4615 .abort_pending_relay(
4616 endpoint,
4617 bind.corr,
4618 RouteBindRelayOutcome::ModuleGone("abort".into()),
4619 )
4620 .unwrap();
4621 let rpc = forwarding
4622 .begin_module_control_rpc_for(
4623 "corr-shared",
4624 "health.check",
4625 Instant::now() - Duration::from_millis(1),
4626 )
4627 .unwrap();
4628 assert_eq!(rpc.corr, 2);
4629 assert_eq!(
4630 forwarding
4631 .complete_module_control_rpc(
4632 module_connection,
4633 rpc.corr,
4634 Some("health.check"),
4635 ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
4636 status: subc_protocol::session::HealthStatus::Ok,
4637 detail: None,
4638 metrics: None,
4639 }),
4640 )
4641 .unwrap(),
4642 ModuleControlRpcCompletion::Settled
4643 );
4644 assert!(matches!(
4645 rpc.receiver.blocking_recv().unwrap(),
4646 ModuleControlRpcOutcome::DeadlineElapsed
4647 ));
4648 }
4649
4650 #[tokio::test(start_paused = true)]
4651 async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
4652 let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
4653 let probe_started_at = Instant::now();
4654 let rpc = forwarding
4655 .begin_health_probe_rpc_for(
4656 "tombstone-ttl",
4657 "health.check",
4658 probe_started_at,
4659 probe_started_at + Duration::from_secs(5),
4660 )
4661 .unwrap();
4662 assert!(forwarding
4663 .tombstone_health_probe_rpc(endpoint, rpc.corr)
4664 .unwrap());
4665 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
4666
4667 tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
4668 tokio::task::yield_now().await;
4669
4670 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
4671 }
4672
4673 #[test]
4674 fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
4675 let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
4676 let mut close = forwarding.register_connection_close(endpoint.connection_id);
4677 forwarding.inject_control_corr(endpoint, u64::MAX);
4678 let final_rpc = forwarding
4679 .begin_module_control_rpc_for(
4680 "corr-max",
4681 "health.check",
4682 Instant::now() + Duration::from_secs(1),
4683 )
4684 .unwrap();
4685 assert_eq!(final_rpc.corr, u64::MAX);
4686 forwarding
4687 .cancel_module_control_rpc(endpoint, final_rpc.corr)
4688 .unwrap();
4689 assert!(matches!(
4690 forwarding.begin_module_control_rpc_for(
4691 "corr-max",
4692 "health.check",
4693 Instant::now() + Duration::from_secs(1),
4694 ),
4695 Err(ForwardingError::RelayCorrelationExhausted)
4696 ));
4697 assert!(close.try_recv().is_ok());
4698 }
4699
4700 #[test]
4701 fn publication_epoch_controls_delivery_failure_escalation() {
4702 fn setup_successor(
4703 commit_successor: Option<bool>,
4704 ) -> (ForwardingTable, ConnectionId, u16, u32) {
4705 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4706 route_fixture("escalation");
4707 let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
4708 forwarding
4709 .complete_pending_relay(
4710 module_connection,
4711 first.corr,
4712 RouteBindRelayOutcome::Accepted,
4713 )
4714 .unwrap();
4715 client_rx.try_recv().unwrap();
4716 assert!(matches!(
4717 forwarding
4718 .release_client_route(client, first.client_channel, first.client_epoch)
4719 .unwrap(),
4720 RouteRelease::Removed(_)
4721 ));
4722 if let Some(commit_successor) = commit_successor {
4723 forwarding.inject_client_slot_epoch(
4724 client,
4725 first.client_channel,
4726 first.client_epoch,
4727 );
4728 forwarding.inject_module_slot_epoch(
4729 endpoint,
4730 first.module_channel,
4731 first.module_epoch,
4732 );
4733 let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
4734 if commit_successor {
4735 forwarding
4736 .complete_pending_relay(
4737 module_connection,
4738 successor.corr,
4739 RouteBindRelayOutcome::Accepted,
4740 )
4741 .unwrap();
4742 client_rx.try_recv().unwrap();
4743 } else {
4744 forwarding
4745 .abort_pending_relay(
4746 endpoint,
4747 successor.corr,
4748 RouteBindRelayOutcome::ModuleGone("abort".into()),
4749 )
4750 .unwrap();
4751 }
4752 }
4753 (forwarding, client, first.client_channel, first.client_epoch)
4754 }
4755
4756 let probe_sink = FrameSink::new(mpsc::channel(1).0);
4757 let (no_successor, client, channel, epoch) = setup_successor(None);
4758 let mut close = no_successor.register_connection_close(client);
4759 assert!(no_successor
4760 .escalate_client_delivery_failure(
4761 client,
4762 channel,
4763 epoch,
4764 CloseReason::new("delivery", "failed"),
4765 UndeliveredFrame {
4766 module_id: None,
4767 sink: &probe_sink,
4768 },
4769 )
4770 .unwrap());
4771 assert!(close.try_recv().is_ok());
4772
4773 let (aborted, client, channel, epoch) = setup_successor(Some(false));
4774 let mut close = aborted.register_connection_close(client);
4775 assert!(aborted
4776 .escalate_client_delivery_failure(
4777 client,
4778 channel,
4779 epoch,
4780 CloseReason::new("delivery", "failed"),
4781 UndeliveredFrame {
4782 module_id: None,
4783 sink: &probe_sink,
4784 },
4785 )
4786 .unwrap());
4787 assert!(close.try_recv().is_ok());
4788
4789 let (published, client, channel, epoch) = setup_successor(Some(true));
4790 let mut close = published.register_connection_close(client);
4791 assert!(!published
4792 .escalate_client_delivery_failure(
4793 client,
4794 channel,
4795 epoch,
4796 CloseReason::new("delivery", "stale failure"),
4797 UndeliveredFrame {
4798 module_id: None,
4799 sink: &probe_sink,
4800 },
4801 )
4802 .unwrap());
4803 assert!(close.try_recv().is_err());
4804 }
4805
4806 #[test]
4807 fn route_concentration_separates_client_count_from_routes_per_client() {
4808 let (forwarding, module_connection, _, client, sink, _client_rx) =
4812 route_fixture("concentration");
4813 assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
4814
4815 for corr in [70_u64, 71] {
4816 let pending =
4817 begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
4818 forwarding
4819 .complete_pending_relay(
4820 module_connection,
4821 pending.corr,
4822 RouteBindRelayOutcome::Accepted,
4823 )
4824 .unwrap();
4825 }
4826
4827 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4829 assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
4830 }
4831
4832 #[test]
4833 fn cleanup_and_accepted_resolution_have_one_lock_winner() {
4834 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4835 route_fixture("cleanup-race");
4836 let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
4837 forwarding
4838 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4839 .unwrap();
4840 let released = forwarding.cleanup_connection(client).unwrap();
4841 assert_eq!(released.len(), 1);
4842 let completion = forwarding
4843 .complete_pending_relay(
4844 module_connection,
4845 pending.corr,
4846 RouteBindRelayOutcome::Accepted,
4847 )
4848 .unwrap();
4849 assert!(!completion.settled);
4850 assert!(client_rx.try_recv().is_err());
4851 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4852
4853 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4854 route_fixture("accepted-race");
4855 let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
4856 forwarding
4857 .complete_pending_relay(
4858 module_connection,
4859 pending.corr,
4860 RouteBindRelayOutcome::Accepted,
4861 )
4862 .unwrap();
4863 assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
4864 let released = forwarding.cleanup_connection(client).unwrap();
4865 assert_eq!(released.len(), 1);
4866 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4867 }
4868
4869 #[test]
4870 fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
4871 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4872 route_fixture("drain-gap");
4873 let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
4874 forwarding
4875 .complete_pending_relay(
4876 module_connection,
4877 live.corr,
4878 RouteBindRelayOutcome::Accepted,
4879 )
4880 .unwrap();
4881 client_rx.try_recv().unwrap();
4882 let binding = match forwarding
4883 .lookup_data_route(client, live.client_channel, live.client_epoch)
4884 .unwrap()
4885 {
4886 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4887 other => panic!("expected live route, got {other:?}"),
4888 };
4889
4890 let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
4891 forwarding
4892 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4893 .unwrap();
4894 let control_rpc = forwarding
4895 .begin_module_control_rpc_for(
4896 "drain-gap",
4897 "health.check",
4898 Instant::now() + Duration::from_secs(1),
4899 )
4900 .unwrap();
4901 let target = forwarding
4902 .begin_module_drain("drain-gap", RouteCloseReason::Reload)
4903 .unwrap()
4904 .unwrap();
4905 assert!(matches!(
4906 control_rpc.receiver.blocking_recv().unwrap(),
4907 ModuleControlRpcOutcome::ModuleGone(_)
4908 ));
4909 assert_eq!(target.abandoned_bindings.len(), 1);
4910 assert!(binding.flow.sem.is_closed());
4911 assert!(
4912 !forwarding
4913 .complete_pending_relay(
4914 module_connection,
4915 pending.corr,
4916 RouteBindRelayOutcome::Accepted,
4917 )
4918 .unwrap()
4919 .settled
4920 );
4921 assert!(matches!(
4922 forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
4923 Err(ForwardingError::ModuleReloading { .. })
4924 ));
4925 let released = forwarding
4926 .release_module_endpoint_routes(target.endpoint)
4927 .unwrap();
4928 assert_eq!(released.len(), 1);
4929 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4930 }
4931
4932 #[test]
4940 fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
4941 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4942 route_fixture("closing-client");
4943
4944 let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
4947 forwarding
4948 .complete_pending_relay(
4949 module_connection,
4950 live.corr,
4951 RouteBindRelayOutcome::Accepted,
4952 )
4953 .unwrap();
4954 client_rx.try_recv().unwrap();
4955
4956 let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
4958 forwarding
4959 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4960 .unwrap();
4961
4962 assert!(forwarding
4964 .escalate_client_delivery_failure(
4965 client,
4966 live.client_channel,
4967 live.client_epoch,
4968 CloseReason::new(
4969 "module_to_client_delivery_failed",
4970 "client egress refused a module frame",
4971 ),
4972 UndeliveredFrame {
4973 module_id: None,
4974 sink: &sink,
4975 },
4976 )
4977 .unwrap());
4978 assert!(!sink.is_closed());
4979
4980 let completion = forwarding
4981 .complete_pending_relay(
4982 module_connection,
4983 pending.corr,
4984 RouteBindRelayOutcome::Accepted,
4985 )
4986 .expect("a closing client must not turn a module's ack into an error");
4987
4988 assert!(completion.settled);
4989 let abandoned = completion
4990 .abandoned
4991 .expect("the module must be told to drop the binding it just created");
4992 assert_eq!(abandoned.connection_id, module_connection);
4993 assert_eq!(abandoned.channel, pending.module_channel);
4994 assert_eq!(abandoned.epoch, pending.module_epoch);
4995 assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
4996 assert!(matches!(
4997 pending.receiver.blocking_recv().unwrap(),
4998 RouteBindRelayOutcome::ModuleGone(_)
4999 ));
5000 assert!(client_rx.try_recv().is_err());
5003 assert_eq!(forwarding.active_binding_count().unwrap(), 1);
5004
5005 assert!(forwarding
5008 .has_live_module_connection("closing-client")
5009 .unwrap());
5010 let cotenant = ConnectionId::new(201);
5011 let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
5012 let cotenant_route = begin_test_route(
5013 &forwarding,
5014 cotenant,
5015 FrameSink::new(cotenant_tx),
5016 62,
5017 "closing-client",
5018 );
5019 assert_eq!(cotenant_route.endpoint, endpoint);
5020 forwarding
5021 .complete_pending_relay(
5022 module_connection,
5023 cotenant_route.corr,
5024 RouteBindRelayOutcome::Accepted,
5025 )
5026 .unwrap();
5027 assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
5028 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
5029 }
5030
5031 #[test]
5032 fn pending_route_permit_is_released_on_rejection_and_abort() {
5033 let forwarding = ForwardingTable::default();
5034 let module_connection = ConnectionId::new(300);
5035 let client = ConnectionId::new(301);
5036 let (module_tx, _module_rx) = mpsc::channel(1);
5037 let endpoint = forwarding
5038 .register_module_connection(
5039 module_connection,
5040 "permit".into(),
5041 2,
5042 Concurrency::ModuleManaged,
5043 FrameSink::new(module_tx),
5044 )
5045 .unwrap();
5046 let (client_tx, mut client_rx) = mpsc::channel(1);
5047 let sink = FrameSink::new(client_tx);
5048 let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
5049 assert!(sink.try_send(test_ping(999)).is_err());
5050 forwarding
5051 .complete_pending_relay(
5052 module_connection,
5053 rejected.corr,
5054 RouteBindRelayOutcome::Rejected(ErrorBody {
5055 code: "no".into(),
5056 message: "rejected".into(),
5057 detail: None,
5058 }),
5059 )
5060 .unwrap();
5061 sink.try_send(test_ping(1000)).unwrap();
5062 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
5063
5064 let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
5065 assert!(sink.try_send(test_ping(1001)).is_err());
5066 forwarding
5067 .abort_pending_relay(
5068 endpoint,
5069 aborted.corr,
5070 RouteBindRelayOutcome::ModuleGone("abort".into()),
5071 )
5072 .unwrap();
5073 sink.try_send(test_ping(1002)).unwrap();
5074 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
5075
5076 let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
5077 forwarding
5078 .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
5079 .unwrap();
5080 drop(client_rx);
5081 let completion = forwarding
5082 .complete_pending_relay(
5083 module_connection,
5084 receiver_closed.corr,
5085 RouteBindRelayOutcome::Accepted,
5086 )
5087 .unwrap();
5088 assert!(completion.abandoned.is_some());
5089 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
5090 }
5091
5092 #[test]
5098 fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
5099 let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
5100 route_fixture("closing-set-leak");
5101
5102 const CONNECTIONS: u64 = 32;
5103 for index in 0..CONNECTIONS {
5104 let client = ConnectionId::new(1000 + index);
5105 let (client_tx, _client_rx) = mpsc::channel(8);
5106 let route = begin_test_route(
5107 &forwarding,
5108 client,
5109 FrameSink::new(client_tx),
5110 index + 1,
5111 "closing-set-leak",
5112 );
5113 forwarding
5114 .complete_pending_relay(
5115 module_connection,
5116 route.corr,
5117 RouteBindRelayOutcome::Accepted,
5118 )
5119 .unwrap();
5120 forwarding.cleanup_connection(client).unwrap();
5121 }
5122 forwarding.cleanup_connection(module_connection).unwrap();
5123
5124 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
5125 }
5126
5127 #[test]
5134 fn closing_connection_is_refused_new_work_until_cleanup_completes() {
5135 let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
5136 route_fixture("closing-gate");
5137
5138 let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
5141 forwarding
5142 .complete_pending_relay(
5143 module_connection,
5144 live.corr,
5145 RouteBindRelayOutcome::Accepted,
5146 )
5147 .unwrap();
5148 client_rx.try_recv().unwrap();
5149
5150 assert!(forwarding
5153 .escalate_client_delivery_failure(
5154 client,
5155 live.client_channel,
5156 live.client_epoch,
5157 CloseReason::new(
5158 "module_to_client_delivery_failed",
5159 "client egress refused a module frame",
5160 ),
5161 UndeliveredFrame {
5162 module_id: None,
5163 sink: &sink,
5164 },
5165 )
5166 .unwrap());
5167 assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
5168
5169 assert!(matches!(
5171 forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
5172 Err(ForwardingError::ConnectionClosing { connection_id })
5173 if connection_id == client
5174 ));
5175 let (late_tx, _late_rx) = mpsc::channel(1);
5177 assert!(matches!(
5178 forwarding.register_module_connection(
5179 client,
5180 "late-module".into(),
5181 2,
5182 Concurrency::ModuleManaged,
5183 FrameSink::new(late_tx),
5184 ),
5185 Err(ForwardingError::ConnectionClosing { connection_id })
5186 if connection_id == client
5187 ));
5188
5189 forwarding.cleanup_connection(client).unwrap();
5193 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
5194 }
5195}
5196
5197#[cfg(test)]
5200mod swap_slot_tests {
5201 use std::time::Duration;
5202
5203 use super::*;
5204 use tokio::sync::mpsc;
5205
5206 const MODULE_ID: &str = "swapped";
5207
5208 struct SwapFixture {
5209 forwarding: ForwardingTable,
5210 incumbent_connection: ConnectionId,
5211 incumbent: ModuleEndpointId,
5212 candidate_connection: ConnectionId,
5213 candidate: ModuleEndpointId,
5214 _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
5215 }
5216
5217 fn swap_fixture() -> SwapFixture {
5218 let forwarding = ForwardingTable::default();
5219 let incumbent_connection = ConnectionId::new(100);
5220 let candidate_connection = ConnectionId::new(110);
5221 let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
5222 let incumbent = forwarding
5223 .register_module_connection(
5224 incumbent_connection,
5225 MODULE_ID.to_string(),
5226 2,
5227 Concurrency::ModuleManaged,
5228 FrameSink::new(incumbent_tx),
5229 )
5230 .unwrap();
5231 let (candidate_tx, candidate_rx) = mpsc::channel(8);
5232 let candidate = forwarding
5233 .register_candidate_module_connection(
5234 candidate_connection,
5235 MODULE_ID.to_string(),
5236 2,
5237 Concurrency::ModuleManaged,
5238 FrameSink::new(candidate_tx),
5239 )
5240 .unwrap();
5241 SwapFixture {
5242 forwarding,
5243 incumbent_connection,
5244 incumbent,
5245 candidate_connection,
5246 candidate,
5247 _module_rxs: vec![incumbent_rx, candidate_rx],
5248 }
5249 }
5250
5251 fn client(
5252 raw: u64,
5253 ) -> (
5254 ConnectionId,
5255 FrameSink,
5256 mpsc::Receiver<crate::router::OutboundFrame>,
5257 ) {
5258 let (tx, rx) = mpsc::channel(8);
5259 (ConnectionId::new(raw), FrameSink::new(tx), rx)
5260 }
5261
5262 fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
5263 forwarding
5264 .read_inner()
5265 .unwrap()
5266 .client_to_module
5267 .values()
5268 .map(|route| route.module_endpoint)
5269 .collect()
5270 }
5271
5272 #[test]
5273 fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
5274 let fixture = swap_fixture();
5275 let forwarding = &fixture.forwarding;
5276 assert_ne!(fixture.incumbent, fixture.candidate);
5277
5278 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
5280 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
5281 let (client_connection, client_sink, _client_rx) = client(200);
5282 let pending = forwarding
5283 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5284 .unwrap();
5285 assert_eq!(pending.endpoint, fixture.incumbent);
5286 let rpc = forwarding
5287 .begin_module_control_rpc_for(
5288 MODULE_ID,
5289 "health.check",
5290 Instant::now() + Duration::from_secs(1),
5291 )
5292 .unwrap();
5293 assert_eq!(rpc.endpoint, fixture.incumbent);
5294 let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
5295 assert_eq!(census.len(), 1, "the census lists one endpoint per id");
5296
5297 assert_eq!(
5299 forwarding
5300 .module_endpoint_for_connection(fixture.candidate_connection)
5301 .unwrap(),
5302 Some(fixture.candidate)
5303 );
5304 assert_eq!(
5305 forwarding
5306 .module_id_for_connection(fixture.candidate_connection)
5307 .unwrap()
5308 .as_deref(),
5309 Some(MODULE_ID)
5310 );
5311
5312 let (other_tx, _other_rx) = mpsc::channel(1);
5314 assert_eq!(
5315 forwarding.register_candidate_module_connection(
5316 ConnectionId::new(120),
5317 MODULE_ID.to_string(),
5318 2,
5319 Concurrency::ModuleManaged,
5320 FrameSink::new(other_tx),
5321 ),
5322 Err(ForwardingError::CandidateSlotOccupied {
5323 module_id: MODULE_ID.to_string()
5324 })
5325 );
5326 }
5327
5328 #[test]
5332 fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
5333 let fixture = swap_fixture();
5334 let forwarding = &fixture.forwarding;
5335 let (early_client, early_sink, _early_rx) = client(200);
5336 let mut early = forwarding
5337 .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
5338 .unwrap();
5339 assert_eq!(early.endpoint, fixture.incumbent);
5340 assert_eq!(
5341 forwarding
5342 .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
5343 .unwrap(),
5344 RouteBindRelayMark::Marked
5345 );
5346
5347 let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
5348 assert_eq!(
5349 cutover,
5350 ForwardingCutover {
5351 promoted: fixture.candidate,
5352 incumbent: Some(fixture.incumbent),
5353 }
5354 );
5355
5356 let (late_client, late_sink, _late_rx) = client(201);
5358 let late = forwarding
5359 .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
5360 .unwrap();
5361 assert_eq!(
5362 late.endpoint, fixture.candidate,
5363 "a route.open after cutover was reserved on the incumbent"
5364 );
5365
5366 let completion = forwarding
5368 .complete_pending_relay(
5369 fixture.incumbent_connection,
5370 early.corr,
5371 RouteBindRelayOutcome::Accepted,
5372 )
5373 .expect("a superseded endpoint's ack is not an error on its connection");
5374 assert!(completion.settled);
5375 assert!(
5376 !committed_endpoints(forwarding).contains(&fixture.incumbent),
5377 "a relay reserved before cutover committed a route on the incumbent"
5378 );
5379 let goodbye = completion
5380 .abandoned
5381 .expect("the incumbent is told to drop the binding it just created");
5382 assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
5383 assert_eq!(goodbye.channel, early.module_channel);
5384 assert_eq!(goodbye.epoch, early.module_epoch);
5385 assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
5386 match early.receiver.try_recv() {
5387 Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
5388 other => panic!("expected a retryable module_reloading answer, got {other:?}"),
5389 }
5390 assert!(matches!(
5391 forwarding
5392 .lookup_data_route(early_client, early.client_channel, early.client_epoch)
5393 .unwrap(),
5394 DataRoute::Client(DataRouteState::Absent)
5395 ));
5396
5397 assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
5400 forwarding
5401 .complete_pending_relay(
5402 fixture.candidate_connection,
5403 late.corr,
5404 RouteBindRelayOutcome::Accepted,
5405 )
5406 .unwrap();
5407 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
5408 assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
5409 }
5410
5411 #[test]
5412 fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
5413 let fixture = swap_fixture();
5414 let forwarding = &fixture.forwarding;
5415 let (bound_client, bound_sink, _bound_rx) = client(200);
5417 let bound = forwarding
5418 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
5419 .unwrap();
5420 forwarding
5421 .complete_pending_relay(
5422 fixture.incumbent_connection,
5423 bound.corr,
5424 RouteBindRelayOutcome::Accepted,
5425 )
5426 .unwrap();
5427 let (pending_client, pending_sink, _pending_rx) = client(201);
5428 let mut in_flight = forwarding
5429 .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
5430 .unwrap();
5431 forwarding
5432 .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
5433 .unwrap();
5434
5435 let incumbent = forwarding
5436 .cutover_candidate(MODULE_ID)
5437 .unwrap()
5438 .unwrap()
5439 .incumbent
5440 .unwrap();
5441 let target = forwarding
5442 .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
5443 .unwrap()
5444 .expect("the superseded incumbent is still registered");
5445
5446 assert_eq!(target.endpoint, fixture.incumbent);
5447 assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
5448 assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
5449 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
5450 assert_eq!(target.abandoned_bindings.len(), 1);
5451 assert_eq!(
5452 target.abandoned_bindings[0].channel,
5453 in_flight.module_channel
5454 );
5455 assert!(matches!(
5456 in_flight.receiver.try_recv(),
5457 Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
5458 ));
5459 assert_eq!(
5460 forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
5461 1,
5462 "the incumbent's bound route stays until its drain finishes"
5463 );
5464
5465 let (next_client, next_sink, _next_rx) = client(202);
5466 let next = forwarding
5467 .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
5468 .expect("the promoted candidate keeps accepting routes");
5469 assert_eq!(next.endpoint, fixture.candidate);
5470 }
5471
5472 #[test]
5478 fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
5479 let forwarding = ForwardingTable::default();
5480 let first_connection = ConnectionId::new(70);
5481 let (first_tx, _first_rx) = mpsc::channel(8);
5482 forwarding
5483 .register_module_connection(
5484 first_connection,
5485 MODULE_ID.to_string(),
5486 2,
5487 Concurrency::ModuleManaged,
5488 FrameSink::new(first_tx),
5489 )
5490 .unwrap();
5491 let (client_connection, client_sink, _client_rx) = client(200);
5492 let mut pending = forwarding
5493 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5494 .unwrap();
5495 let (second_tx, _second_rx) = mpsc::channel(8);
5496 forwarding
5497 .register_module_connection(
5498 ConnectionId::new(80),
5499 MODULE_ID.to_string(),
5500 2,
5501 Concurrency::ModuleManaged,
5502 FrameSink::new(second_tx),
5503 )
5504 .unwrap();
5505
5506 assert_eq!(
5507 forwarding
5508 .complete_pending_relay(
5509 first_connection,
5510 pending.corr,
5511 RouteBindRelayOutcome::Accepted
5512 )
5513 .unwrap_err(),
5514 ForwardingError::StaleModuleEndpoint
5515 );
5516 assert!(committed_endpoints(&forwarding).is_empty());
5517 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
5518 assert!(matches!(
5519 pending.receiver.try_recv(),
5520 Err(oneshot::error::TryRecvError::Closed)
5521 ));
5522 }
5523
5524 #[test]
5525 fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
5526 let fixture = swap_fixture();
5528 let forwarding = &fixture.forwarding;
5529 assert!(forwarding
5530 .cleanup_connection(fixture.candidate_connection)
5531 .unwrap()
5532 .is_empty());
5533 assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
5534 let (client_connection, client_sink, _client_rx) = client(200);
5535 assert_eq!(
5536 forwarding
5537 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
5538 .unwrap()
5539 .endpoint,
5540 fixture.incumbent
5541 );
5542
5543 let fixture = swap_fixture();
5546 let forwarding = &fixture.forwarding;
5547 let (bound_client, bound_sink, _bound_rx) = client(200);
5548 let bound = forwarding
5549 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
5550 .unwrap();
5551 forwarding
5552 .complete_pending_relay(
5553 fixture.incumbent_connection,
5554 bound.corr,
5555 RouteBindRelayOutcome::Accepted,
5556 )
5557 .unwrap();
5558 forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
5559 let released = forwarding
5560 .cleanup_connection(fixture.incumbent_connection)
5561 .unwrap();
5562 assert_eq!(released.len(), 1);
5563 assert_eq!(released[0].connection_id, bound_client);
5564 assert!(forwarding
5565 .read_inner()
5566 .unwrap()
5567 .superseded_endpoints
5568 .is_empty());
5569 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
5570 let (next_client, next_sink, _next_rx) = client(201);
5571 assert_eq!(
5572 forwarding
5573 .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
5574 .unwrap()
5575 .endpoint,
5576 fixture.candidate
5577 );
5578 }
5579}