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