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