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