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