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