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 Frame, ProjectRootId,
25};
26
27const DEFAULT_MODULE_MANAGED_WINDOW: usize = 32;
29
30const STATELESS_PARALLEL_WINDOW: usize = 1024;
32
33const HEALTH_PROBE_TOMBSTONE_TTL: Duration = Duration::from_secs(5 * 60);
37
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub struct ModuleEndpointId {
44 pub connection_id: ConnectionId,
45 pub generation: u64,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
50pub(crate) struct ClientRouteKey {
51 pub connection_id: ConnectionId,
52 pub channel: u16,
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
57pub(crate) struct ModuleRouteKey {
58 pub endpoint: ModuleEndpointId,
59 pub channel: u16,
60}
61
62#[derive(Debug)]
63pub(crate) struct RouteBinding {
64 pub client_connection_id: ConnectionId,
65 pub client_sink: FrameSink,
66 pub client_negotiated_ver: u8,
67 pub client_channel: u16,
68 pub client_epoch: u32,
69 pub module_id: String,
70 pub module_endpoint: ModuleEndpointId,
71 pub module_sink: FrameSink,
72 pub module_negotiated_ver: u8,
73 pub module_channel: u16,
74 pub module_epoch: u32,
75 pub principal: Principal,
76 pub project_root: Option<ProjectRootId>,
77 pub bound_at: Instant,
78 pub flow: Arc<ChannelFlow>,
79}
80
81#[derive(Debug, Clone)]
82pub(crate) enum DataRoute {
83 Client(DataRouteState),
84 Module(DataRouteState),
85}
86
87#[derive(Debug, Clone)]
88pub(crate) enum DataRouteState {
89 Bound(Arc<RouteBinding>),
90 Reserved,
91 EpochMismatch,
92 Absent,
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub(crate) enum GoodbyeTargetKind {
125 Client,
126 Module,
127}
128
129#[derive(Debug, Clone)]
130pub(crate) struct GoodbyeTarget {
131 pub connection_id: ConnectionId,
132 pub sink: FrameSink,
133 pub negotiated_ver: u8,
134 pub channel: u16,
135 pub epoch: u32,
136 pub kind: GoodbyeTargetKind,
137 pub module_id: Option<String>,
142}
143
144#[derive(Debug, Clone, Copy)]
147pub(crate) struct UndeliveredFrame<'a> {
148 pub module_id: Option<&'a str>,
150 pub sink: &'a FrameSink,
152}
153
154fn principal_label(principal: &Principal) -> String {
156 match principal {
157 Principal::Reserved { module_id } => format!("reserved:{module_id}"),
158 Principal::Direct => "direct".to_string(),
159 other => format!("{other:?}"),
160 }
161}
162
163fn connection_principals_locked(inner: &ForwardingInner, connection_id: ConnectionId) -> String {
166 let labels = inner
167 .client_to_module
168 .iter()
169 .filter(|(key, _)| key.connection_id == connection_id)
170 .map(|(_, route)| principal_label(&route.principal))
171 .collect::<std::collections::BTreeSet<_>>();
172 if labels.is_empty() {
173 "none".to_string()
174 } else {
175 labels.into_iter().collect::<Vec<_>>().join(",")
176 }
177}
178
179impl GoodbyeTarget {
180 pub(crate) fn close_on_delivery_failure(&self) -> bool {
183 matches!(self.kind, GoodbyeTargetKind::Client)
184 }
185}
186
187pub(crate) const LATE_MODULE_GOODBYE_DEADLINE: Duration = crate::supervise::DEFAULT_DRAIN_TIMEOUT;
200
201pub(crate) fn send_module_route_goodbye(
213 counters: &DaemonCounters,
214 sink: &FrameSink,
215 frame: Frame,
216 module_id: Option<&str>,
217 context: &'static str,
218) {
219 let channel = frame.header.channel;
220 let epoch = frame.header.epoch;
221 let Err(err) = sink.try_send(frame.clone()) else {
222 return;
223 };
224 let runtime = match tokio::runtime::Handle::try_current() {
227 Ok(runtime) if !sink.is_closed() => runtime,
228 _ => {
229 counters.increment_goodbye_relay_module_dropped(module_id);
230 warn!(
231 module_id = module_id.unwrap_or("unknown"),
232 route_channel = channel,
233 route_epoch = epoch,
234 error = %err,
235 context,
236 "route GOODBYE to module dropped: module connection is closed; not closing shared module connection"
237 );
238 return;
239 }
240 };
241 debug!(
242 module_id = module_id.unwrap_or("unknown"),
243 route_channel = channel,
244 route_epoch = epoch,
245 error = %err,
246 context,
247 "module egress queue refused route GOODBYE; delivering it once the module frees room"
248 );
249 let counters = counters.clone();
250 let sink = sink.clone();
251 let module_id = module_id.map(str::to_string);
252 runtime.spawn(async move {
253 let outcome = tokio::time::timeout(LATE_MODULE_GOODBYE_DEADLINE, sink.send(frame)).await;
254 let why = match outcome {
255 Ok(Ok(())) => {
256 debug!(
257 module_id = module_id.as_deref().unwrap_or("unknown"),
258 route_channel = channel,
259 route_epoch = epoch,
260 context,
261 "late route GOODBYE delivered to module"
262 );
263 return;
264 }
265 Ok(Err(err)) => err.to_string(),
266 Err(_) => format!(
267 "module egress queue had no room within {LATE_MODULE_GOODBYE_DEADLINE:?}"
268 ),
269 };
270 counters.increment_goodbye_relay_module_dropped(module_id.as_deref());
271 warn!(
272 module_id = module_id.as_deref().unwrap_or("unknown"),
273 route_channel = channel,
274 route_epoch = epoch,
275 error = %why,
276 context,
277 "route GOODBYE to module dropped under backpressure; not closing shared module connection"
278 );
279 });
280}
281
282#[derive(Debug, Clone)]
288pub(crate) struct EndpointRoute {
289 pub goodbye_target: GoodbyeTarget,
290 pub principal: Principal,
291 pub bound_at: Instant,
292 pub draining: bool,
293 pub drain_reason: Option<RouteCloseReason>,
297}
298
299#[derive(Debug)]
300pub(crate) struct PendingRouteBindRelay {
301 pub endpoint: ModuleEndpointId,
302 pub module_sink: FrameSink,
303 pub negotiated_ver: u8,
304 pub client_channel: u16,
305 pub client_epoch: u32,
306 pub module_channel: u16,
307 pub module_epoch: u32,
308 pub corr: u64,
309 pub receiver: oneshot::Receiver<RouteBindRelayOutcome>,
310}
311
312#[derive(Debug, Clone)]
313pub(crate) struct ModuleDrainTarget {
314 pub endpoint: ModuleEndpointId,
315 pub sink: FrameSink,
316 pub negotiated_ver: u8,
317 pub abandoned_bindings: Vec<GoodbyeTarget>,
318 pub excluded_subscriptions: u32,
319}
320
321#[cfg(unix)]
324#[derive(Debug, Clone)]
325pub(crate) struct ModuleConnectionTarget {
326 pub module_id: String,
327 pub endpoint: ModuleEndpointId,
328 pub sink: FrameSink,
329 pub negotiated_ver: u8,
330}
331
332#[derive(Debug, Clone)]
333pub(crate) enum RouteBindRelayOutcome {
334 Accepted,
335 Rejected(ErrorBody),
336 ModuleGone(String),
337}
338
339#[derive(Debug, Clone, Copy, PartialEq, Eq)]
341pub(crate) struct ForwardingCutover {
342 pub promoted: ModuleEndpointId,
344 pub incumbent: Option<ModuleEndpointId>,
347}
348
349#[derive(Debug)]
351pub(crate) struct ConnectionCleanup {
352 pub released: Vec<GoodbyeTarget>,
354 pub abandoned_relays: u32,
357}
358
359#[derive(Debug, Clone)]
360pub(crate) struct PendingRelayCompletion {
361 pub settled: bool,
362 pub abandoned: Option<GoodbyeTarget>,
363}
364
365#[derive(Debug)]
366pub(crate) struct PendingModuleControlRpc {
367 pub endpoint: ModuleEndpointId,
368 pub module_sink: FrameSink,
369 pub negotiated_ver: u8,
370 pub corr: u64,
371 pub receiver: oneshot::Receiver<ModuleControlRpcOutcome>,
372}
373
374#[derive(Debug, Clone)]
375pub(crate) enum ModuleControlRpcOutcome {
376 Response(ModuleControlResponse),
377 Rejected(ErrorBody),
378 ModuleGone(String),
379 MalformedResponse(String),
380 UnexpectedOp { expected: String, actual: String },
381 DeadlineElapsed,
382}
383
384#[derive(Debug, Clone, PartialEq, Eq)]
385pub(crate) enum ModuleControlRpcCompletion {
386 Unknown,
387 Settled,
388 LateHealthAnswer {
389 module_id: String,
390 latency: Duration,
391 },
392}
393
394#[derive(Debug)]
395struct PendingModuleControlRpcEntry {
396 expected_op: String,
397 deadline: Instant,
398 health_probe_started_at: Option<Instant>,
399 sender: oneshot::Sender<ModuleControlRpcOutcome>,
400}
401
402#[derive(Debug)]
403struct HealthProbeTombstone {
404 expected_op: String,
405 module_id: String,
406 probe_started_at: Instant,
407 expires_at: Instant,
408}
409
410#[derive(Debug, Clone)]
411struct RouteReservation {
412 client_key: ClientRouteKey,
413 module_key: ModuleRouteKey,
414 client_epoch: u32,
415 module_epoch: u32,
416 project_root: Option<ProjectRootId>,
417}
418
419#[derive(Debug)]
420struct PendingRouteBindRelayEntry {
421 reservation: RouteReservation,
422 client_sink: FrameSink,
423 client_negotiated_ver: u8,
424 client_permit: crate::router::EgressPermit,
425 route_open_frame: Frame,
426 principal: Principal,
427 deadline: Instant,
428 relay_enqueued: bool,
429 sender: oneshot::Sender<RouteBindRelayOutcome>,
430}
431
432#[derive(Debug, Clone)]
433pub(crate) enum RouteRelease {
434 Removed(GoodbyeTarget),
435 Stale,
436 Absent,
437}
438
439#[derive(Debug, Clone)]
440pub(crate) enum RoutePollSnapshot {
441 Bound {
442 module_id: String,
443 status: Option<String>,
444 },
445 Absent,
446}
447
448#[derive(Debug, Clone)]
449struct ModuleConnection {
450 endpoint: ModuleEndpointId,
451 sink: FrameSink,
452 negotiated_ver: u8,
453 concurrency: Concurrency,
454}
455
456#[derive(Debug, Default)]
457struct ForwardingInner {
458 daemon_draining: bool,
459 modules_by_id: HashMap<String, ModuleConnection>,
463 candidates_by_id: HashMap<String, ModuleConnection>,
469 superseded_endpoints: HashMap<ModuleEndpointId, ModuleConnection>,
475 endpoint_by_connection: HashMap<ConnectionId, ModuleEndpointId>,
476 module_id_by_endpoint: HashMap<ModuleEndpointId, String>,
477 draining_endpoints: HashMap<ModuleEndpointId, RouteCloseReason>,
481 closing_connections: HashSet<ConnectionId>,
482 next_generation: u64,
483 reserved_client: HashMap<ClientRouteKey, ModuleRouteKey>,
484 reserved_module: HashMap<ModuleRouteKey, ClientRouteKey>,
485 next_client_channel: HashMap<ConnectionId, u16>,
486 next_module_channel: HashMap<ModuleEndpointId, u16>,
487 client_slot_epochs: HashMap<ClientRouteKey, u32>,
488 module_slot_epochs: HashMap<ModuleRouteKey, u32>,
489 last_published_epoch: HashMap<ClientRouteKey, u32>,
490 client_to_module: HashMap<ClientRouteKey, Arc<RouteBinding>>,
491 module_to_client: HashMap<ModuleRouteKey, Arc<RouteBinding>>,
492 status: HashMap<(ClientRouteKey, u32), String>,
493 pending_relays: HashMap<(ModuleEndpointId, u64), PendingRouteBindRelayEntry>,
494 next_control_corr: HashMap<ModuleEndpointId, u64>,
495 pending_control_rpcs: HashMap<(ModuleEndpointId, u64), PendingModuleControlRpcEntry>,
496 health_probe_tombstones: HashMap<(ModuleEndpointId, u64), HealthProbeTombstone>,
497}
498
499#[derive(Debug, Clone)]
500pub(crate) struct CloseReason {
501 code: &'static str,
502 message: String,
503}
504
505impl CloseReason {
506 pub(crate) fn new(code: &'static str, message: impl Into<String>) -> Self {
507 Self {
508 code,
509 message: message.into(),
510 }
511 }
512}
513
514impl fmt::Display for CloseReason {
515 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
516 write!(f, "{}: {}", self.code, self.message)
517 }
518}
519
520pub(crate) type ConnectionCloseReceiver = oneshot::Receiver<CloseReason>;
521
522#[derive(Debug, Default)]
524pub struct ForwardingTable {
525 inner: Arc<RwLock<ForwardingInner>>,
526 close_registry: Mutex<HashMap<ConnectionId, oneshot::Sender<CloseReason>>>,
527 counters: DaemonCounters,
528 route_bind_breakers: RouteBindBreakers,
533 route_bind_concurrency: RouteBindConcurrency,
536 route_outages: Arc<crate::route_outage::RouteOutageTracker>,
540}
541
542impl ForwardingTable {
543 pub(crate) fn counters(&self) -> DaemonCounters {
544 self.counters.clone()
545 }
546
547 pub(crate) fn route_bind_breakers(&self) -> RouteBindBreakers {
548 self.route_bind_breakers.clone()
549 }
550
551 pub(crate) fn route_bind_concurrency(&self) -> RouteBindConcurrency {
552 self.route_bind_concurrency.clone()
553 }
554
555 pub(crate) fn route_outages(&self) -> Arc<crate::route_outage::RouteOutageTracker> {
556 Arc::clone(&self.route_outages)
557 }
558
559 pub(crate) fn register_connection_close(
560 &self,
561 connection_id: ConnectionId,
562 ) -> ConnectionCloseReceiver {
563 let (sender, receiver) = oneshot::channel();
564 let replaced = self
565 .lock_close_registry()
566 .insert(connection_id, sender)
567 .is_some();
568 if replaced {
569 warn!(
570 connection_id = connection_id.get(),
571 "replaced existing connection close registration"
572 );
573 }
574 receiver
575 }
576
577 pub(crate) fn unregister_connection_close(&self, connection_id: ConnectionId) {
578 self.lock_close_registry().remove(&connection_id);
579 }
580
581 #[cfg(unix)]
590 pub(crate) fn close_all_connections(&self, reason: &CloseReason) -> usize {
591 let senders: Vec<_> = self.lock_close_registry().drain().collect();
592 let count = senders.len();
593 for (_, sender) in senders {
594 let _ = sender.send(reason.clone());
595 }
596 count
597 }
598
599 #[cfg(unix)]
603 pub(crate) fn module_connections(
604 &self,
605 ) -> Result<Vec<ModuleConnectionTarget>, ForwardingError> {
606 let inner = self.read_inner()?;
607 let mut seen = HashSet::new();
608 Ok(inner
609 .modules_by_id
610 .values()
611 .chain(inner.candidates_by_id.values())
612 .chain(inner.superseded_endpoints.values())
613 .filter(|module| seen.insert(module.endpoint))
614 .map(|module| ModuleConnectionTarget {
615 module_id: inner
616 .module_id_by_endpoint
617 .get(&module.endpoint)
618 .cloned()
619 .unwrap_or_default(),
620 endpoint: module.endpoint,
621 sink: module.sink.clone(),
622 negotiated_ver: module.negotiated_ver,
623 })
624 .collect())
625 }
626
627 pub(crate) fn request_connection_close(
631 &self,
632 connection_id: ConnectionId,
633 reason: CloseReason,
634 ) -> bool {
635 let sender = self.lock_close_registry().remove(&connection_id);
636 if let Some(sender) = sender {
637 debug!(
638 connection_id = connection_id.get(),
639 close_reason = %reason,
640 "requesting connection close"
641 );
642 let _ = sender.send(reason);
643 true
644 } else {
645 debug!(
646 connection_id = connection_id.get(),
647 close_reason = %reason,
648 "connection close request ignored for inactive connection"
649 );
650 false
651 }
652 }
653
654 pub fn register_module_connection(
655 &self,
656 connection_id: ConnectionId,
657 module_id: String,
658 negotiated_ver: u8,
659 concurrency: Concurrency,
660 sink: FrameSink,
661 ) -> Result<ModuleEndpointId, ForwardingError> {
662 self.register_module_connection_inner(
663 connection_id,
664 module_id,
665 negotiated_ver,
666 concurrency,
667 sink,
668 None,
669 )
670 }
671
672 pub(crate) fn register_module_connection_acked(
685 &self,
686 connection_id: ConnectionId,
687 module_id: String,
688 negotiated_ver: u8,
689 concurrency: Concurrency,
690 sink: FrameSink,
691 hello_ack: Frame,
692 ) -> Result<ModuleEndpointId, ForwardingError> {
693 self.register_module_connection_inner(
694 connection_id,
695 module_id,
696 negotiated_ver,
697 concurrency,
698 sink,
699 Some(hello_ack),
700 )
701 }
702
703 fn register_module_connection_inner(
704 &self,
705 connection_id: ConnectionId,
706 module_id: String,
707 negotiated_ver: u8,
708 concurrency: Concurrency,
709 sink: FrameSink,
710 hello_ack: Option<Frame>,
711 ) -> Result<ModuleEndpointId, ForwardingError> {
712 let mut inner = self.write_inner()?;
713 if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
714 return Err(ForwardingError::ConnectionClosing { connection_id });
715 }
716 enqueue_hello_ack_locked(&sink, connection_id, hello_ack)?;
719 if let Some(old_endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
720 let _ = remove_module_connection_locked(&mut inner, old_endpoint);
721 }
722
723 inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
724 let endpoint = ModuleEndpointId {
725 connection_id,
726 generation: inner.next_generation,
727 };
728 inner.endpoint_by_connection.insert(connection_id, endpoint);
729 inner
730 .module_id_by_endpoint
731 .insert(endpoint, module_id.clone());
732 inner.next_module_channel.insert(endpoint, 1);
733 inner.next_control_corr.insert(endpoint, 1);
734 inner.modules_by_id.insert(
735 module_id.clone(),
736 ModuleConnection {
737 endpoint,
738 sink,
739 negotiated_ver,
740 concurrency,
741 },
742 );
743 drop(inner);
744
745 if let Some(discarded) = self
758 .route_bind_breakers
759 .reset_for_new_module_connection(&module_id)
760 {
761 info!(
762 module_id = %module_id,
763 discarded_consecutive_timeouts = discarded,
764 "route.bind breaker state discarded: a new module connection replaced the process it described"
765 );
766 }
767 Ok(endpoint)
768 }
769
770 #[cfg(test)]
784 pub(crate) fn register_candidate_module_connection(
785 &self,
786 connection_id: ConnectionId,
787 module_id: String,
788 negotiated_ver: u8,
789 concurrency: Concurrency,
790 sink: FrameSink,
791 ) -> Result<ModuleEndpointId, ForwardingError> {
792 self.register_candidate_module_connection_inner(
793 connection_id,
794 module_id,
795 negotiated_ver,
796 concurrency,
797 sink,
798 None,
799 )
800 }
801
802 pub(crate) fn register_candidate_module_connection_acked(
809 &self,
810 connection_id: ConnectionId,
811 module_id: String,
812 negotiated_ver: u8,
813 concurrency: Concurrency,
814 sink: FrameSink,
815 hello_ack: Frame,
816 ) -> Result<ModuleEndpointId, ForwardingError> {
817 self.register_candidate_module_connection_inner(
818 connection_id,
819 module_id,
820 negotiated_ver,
821 concurrency,
822 sink,
823 Some(hello_ack),
824 )
825 }
826
827 fn register_candidate_module_connection_inner(
828 &self,
829 connection_id: ConnectionId,
830 module_id: String,
831 negotiated_ver: u8,
832 concurrency: Concurrency,
833 sink: FrameSink,
834 hello_ack: Option<Frame>,
835 ) -> Result<ModuleEndpointId, ForwardingError> {
836 let mut inner = self.write_inner()?;
837 if inner.daemon_draining || inner.closing_connections.contains(&connection_id) {
838 return Err(ForwardingError::ConnectionClosing { connection_id });
839 }
840 if inner.candidates_by_id.contains_key(&module_id) {
841 return Err(ForwardingError::CandidateSlotOccupied { module_id });
842 }
843 enqueue_hello_ack_locked(&sink, connection_id, hello_ack)?;
846 if let Some(old_endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
847 let _ = remove_module_connection_locked(&mut inner, old_endpoint);
848 }
849
850 inner.next_generation = inner.next_generation.checked_add(1).unwrap_or(1);
851 let endpoint = ModuleEndpointId {
852 connection_id,
853 generation: inner.next_generation,
854 };
855 inner.endpoint_by_connection.insert(connection_id, endpoint);
856 inner
857 .module_id_by_endpoint
858 .insert(endpoint, module_id.clone());
859 inner.next_module_channel.insert(endpoint, 1);
860 inner.next_control_corr.insert(endpoint, 1);
861 inner.candidates_by_id.insert(
862 module_id,
863 ModuleConnection {
864 endpoint,
865 sink,
866 negotiated_ver,
867 concurrency,
868 },
869 );
870 Ok(endpoint)
871 }
872
873 pub(crate) fn cutover_candidate(
891 &self,
892 module_id: &str,
893 ) -> Result<Option<ForwardingCutover>, ForwardingError> {
894 let mut inner = self.write_inner()?;
895 if inner.daemon_draining {
896 return Err(ForwardingError::ModuleReloading {
897 module_id: module_id.to_string(),
898 });
899 }
900 let Some(candidate) = inner.candidates_by_id.remove(module_id) else {
901 return Ok(None);
902 };
903 let promoted = candidate.endpoint;
904 let incumbent = inner.modules_by_id.insert(module_id.to_string(), candidate);
905 let incumbent = incumbent.map(|incumbent| {
906 let endpoint = incumbent.endpoint;
907 inner.superseded_endpoints.insert(endpoint, incumbent);
908 endpoint
909 });
910 drop(inner);
911
912 if let Some(discarded) = self
916 .route_bind_breakers
917 .reset_for_new_module_connection(module_id)
918 {
919 info!(
920 module_id = %module_id,
921 discarded_consecutive_timeouts = discarded,
922 "route.bind breaker state discarded: a swap candidate was promoted over the process it described"
923 );
924 }
925 Ok(Some(ForwardingCutover {
926 promoted,
927 incumbent,
928 }))
929 }
930
931 #[allow(clippy::too_many_arguments)]
932 pub(crate) async fn begin_route_bind_relay_for(
933 &self,
934 client_connection_id: ConnectionId,
935 client_sink: FrameSink,
936 client_negotiated_ver: u8,
937 client_corr: u64,
938 module_id: &str,
939 principal: Principal,
940 project_root: Option<ProjectRootId>,
941 deadline: Instant,
942 ) -> Result<PendingRouteBindRelay, ForwardingError> {
943 let client_permit =
947 client_sink
948 .reserve_owned()
949 .await
950 .map_err(|_| ForwardingError::ClientEgressClosed {
951 connection_id: client_connection_id,
952 })?;
953 self.begin_route_bind_relay_inner(
954 client_connection_id,
955 client_sink,
956 client_negotiated_ver,
957 client_corr,
958 module_id,
959 principal,
960 project_root,
961 deadline,
962 client_permit,
963 )
964 }
965
966 #[cfg(test)]
967 pub(crate) fn begin_route_bind_relay_for_test(
968 &self,
969 client_connection_id: ConnectionId,
970 client_sink: FrameSink,
971 client_corr: u64,
972 module_id: &str,
973 ) -> Result<PendingRouteBindRelay, ForwardingError> {
974 let permit =
975 client_sink
976 .try_reserve_owned()
977 .map_err(|_| ForwardingError::ClientEgressClosed {
978 connection_id: client_connection_id,
979 })?;
980 self.begin_route_bind_relay_inner(
981 client_connection_id,
982 client_sink,
983 subc_protocol::PROTOCOL_VERSION,
984 client_corr,
985 module_id,
986 Principal::Direct,
987 None,
988 Instant::now() + std::time::Duration::from_secs(60),
989 permit,
990 )
991 }
992
993 pub(crate) fn begin_module_control_rpc_for(
994 &self,
995 module_id: &str,
996 expected_op: &str,
997 deadline: Instant,
998 ) -> Result<PendingModuleControlRpc, ForwardingError> {
999 self.begin_module_control_rpc_inner(module_id, expected_op, deadline, None, false)
1000 }
1001
1002 pub(crate) fn begin_health_probe_rpc_for(
1003 &self,
1004 module_id: &str,
1005 expected_op: &str,
1006 probe_started_at: Instant,
1007 deadline: Instant,
1008 ) -> Result<PendingModuleControlRpc, ForwardingError> {
1009 self.begin_module_control_rpc_inner(
1010 module_id,
1011 expected_op,
1012 deadline,
1013 Some(probe_started_at),
1014 false,
1015 )
1016 }
1017
1018 pub(crate) fn begin_drain_health_probe_rpc_for(
1019 &self,
1020 module_id: &str,
1021 expected_op: &str,
1022 probe_started_at: Instant,
1023 deadline: Instant,
1024 ) -> Result<PendingModuleControlRpc, ForwardingError> {
1025 self.begin_module_control_rpc_inner(
1026 module_id,
1027 expected_op,
1028 deadline,
1029 Some(probe_started_at),
1030 true,
1031 )
1032 }
1033
1034 pub(crate) fn begin_endpoint_health_probe_rpc_for(
1042 &self,
1043 endpoint: ModuleEndpointId,
1044 expected_op: &str,
1045 probe_started_at: Instant,
1046 deadline: Instant,
1047 ) -> Result<PendingModuleControlRpc, ForwardingError> {
1048 let inner = self.write_inner()?;
1049 let module = module_connection_for_endpoint_locked(&inner, endpoint)
1050 .cloned()
1051 .ok_or(ForwardingError::NoModuleConnection)?;
1052 let module_id = inner
1053 .module_id_by_endpoint
1054 .get(&endpoint)
1055 .cloned()
1056 .unwrap_or_default();
1057 self.begin_control_rpc_locked(
1060 inner,
1061 &module_id,
1062 module,
1063 expected_op,
1064 deadline,
1065 Some(probe_started_at),
1066 true,
1067 )
1068 }
1069
1070 fn begin_module_control_rpc_inner(
1071 &self,
1072 module_id: &str,
1073 expected_op: &str,
1074 deadline: Instant,
1075 health_probe_started_at: Option<Instant>,
1076 allow_draining: bool,
1077 ) -> Result<PendingModuleControlRpc, ForwardingError> {
1078 let inner = self.write_inner()?;
1079 let module = inner
1080 .modules_by_id
1081 .get(module_id)
1082 .cloned()
1083 .ok_or(ForwardingError::NoModuleConnection)?;
1084 self.begin_control_rpc_locked(
1085 inner,
1086 module_id,
1087 module,
1088 expected_op,
1089 deadline,
1090 health_probe_started_at,
1091 allow_draining,
1092 )
1093 }
1094
1095 #[allow(clippy::too_many_arguments)]
1096 fn begin_control_rpc_locked(
1097 &self,
1098 mut inner: RwLockWriteGuard<'_, ForwardingInner>,
1099 module_id: &str,
1100 module: ModuleConnection,
1101 expected_op: &str,
1102 deadline: Instant,
1103 health_probe_started_at: Option<Instant>,
1104 allow_draining: bool,
1105 ) -> Result<PendingModuleControlRpc, ForwardingError> {
1106 if !allow_draining && inner.draining_endpoints.contains_key(&module.endpoint) {
1107 return Err(ForwardingError::ModuleReloading {
1108 module_id: module_id.to_string(),
1109 });
1110 }
1111 if inner
1112 .closing_connections
1113 .contains(&module.endpoint.connection_id)
1114 {
1115 return Err(ForwardingError::ConnectionClosing {
1116 connection_id: module.endpoint.connection_id,
1117 });
1118 }
1119 if health_probe_started_at.is_some() {
1120 inner
1124 .health_probe_tombstones
1125 .retain(|(endpoint, _), _| *endpoint != module.endpoint);
1126 }
1127 let corr = match inner.allocate_control_corr(module.endpoint) {
1128 Ok(corr) => corr,
1129 Err(err) => {
1130 drop(inner);
1131 self.request_connection_close(
1132 module.endpoint.connection_id,
1133 CloseReason::new(
1134 "control_correlation_exhausted",
1135 "daemon-originated channel-0 correlation space exhausted",
1136 ),
1137 );
1138 return Err(err);
1139 }
1140 };
1141 let (sender, receiver) = oneshot::channel();
1142 inner.pending_control_rpcs.insert(
1143 (module.endpoint, corr),
1144 PendingModuleControlRpcEntry {
1145 expected_op: expected_op.to_string(),
1146 deadline,
1147 health_probe_started_at,
1148 sender,
1149 },
1150 );
1151
1152 Ok(PendingModuleControlRpc {
1153 endpoint: module.endpoint,
1154 module_sink: module.sink,
1155 negotiated_ver: module.negotiated_ver,
1156 corr,
1157 receiver,
1158 })
1159 }
1160
1161 #[allow(clippy::too_many_arguments)]
1162 fn begin_route_bind_relay_inner(
1163 &self,
1164 client_connection_id: ConnectionId,
1165 client_sink: FrameSink,
1166 client_negotiated_ver: u8,
1167 client_corr: u64,
1168 expected_module_id: &str,
1169 principal: Principal,
1170 project_root: Option<ProjectRootId>,
1171 deadline: Instant,
1172 client_permit: crate::router::EgressPermit,
1173 ) -> Result<PendingRouteBindRelay, ForwardingError> {
1174 let mut inner = self.write_inner()?;
1175 if inner.closing_connections.contains(&client_connection_id) {
1176 return Err(ForwardingError::ConnectionClosing {
1177 connection_id: client_connection_id,
1178 });
1179 }
1180 let module = inner
1181 .modules_by_id
1182 .get(expected_module_id)
1183 .cloned()
1184 .ok_or(ForwardingError::NoModuleConnection)?;
1185 if inner.draining_endpoints.contains_key(&module.endpoint) {
1186 return Err(ForwardingError::ModuleReloading {
1187 module_id: expected_module_id.to_string(),
1188 });
1189 }
1190 if inner
1191 .closing_connections
1192 .contains(&module.endpoint.connection_id)
1193 {
1194 return Err(ForwardingError::ConnectionClosing {
1195 connection_id: module.endpoint.connection_id,
1196 });
1197 }
1198
1199 let corr = match inner.allocate_control_corr(module.endpoint) {
1200 Ok(corr) => corr,
1201 Err(err) => {
1202 drop(inner);
1203 self.request_connection_close(
1204 module.endpoint.connection_id,
1205 CloseReason::new(
1206 "control_correlation_exhausted",
1207 "daemon-originated channel-0 correlation space exhausted",
1208 ),
1209 );
1210 return Err(err);
1211 }
1212 };
1213 let (client_channel, client_epoch, module_channel, module_epoch) =
1214 inner.allocate_route_slots(client_connection_id, module.endpoint)?;
1215 let client_key = ClientRouteKey {
1216 connection_id: client_connection_id,
1217 channel: client_channel,
1218 };
1219 let module_key = ModuleRouteKey {
1220 endpoint: module.endpoint,
1221 channel: module_channel,
1222 };
1223 let reservation = RouteReservation {
1224 client_key,
1225 module_key,
1226 client_epoch,
1227 module_epoch,
1228 project_root,
1229 };
1230 let response_body = serde_json::to_vec(&ClientControlResponse::RouteOpen {
1231 route_channel: client_channel,
1232 route_epoch: client_epoch,
1233 })
1234 .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
1235 let route_open_frame = Frame::build_with_version(
1236 client_negotiated_ver,
1237 FrameType::Response,
1238 Flags::new(false, Priority::Passive, false),
1239 0,
1240 0,
1241 client_corr,
1242 response_body,
1243 )
1244 .map_err(|err| ForwardingError::RouteOpenBuild(err.to_string()))?;
1245 let (sender, receiver) = oneshot::channel();
1246 inner.reserved_client.insert(client_key, module_key);
1247 inner.reserved_module.insert(module_key, client_key);
1248 inner.pending_relays.insert(
1249 (module.endpoint, corr),
1250 PendingRouteBindRelayEntry {
1251 reservation,
1252 client_sink,
1253 client_negotiated_ver,
1254 client_permit,
1255 route_open_frame,
1256 principal,
1257 deadline,
1258 relay_enqueued: false,
1259 sender,
1260 },
1261 );
1262
1263 Ok(PendingRouteBindRelay {
1264 endpoint: module.endpoint,
1265 module_sink: module.sink,
1266 negotiated_ver: module.negotiated_ver,
1267 client_channel,
1268 client_epoch,
1269 module_channel,
1270 module_epoch,
1271 corr,
1272 receiver,
1273 })
1274 }
1275
1276 pub(crate) fn mark_route_bind_relay_enqueued(
1277 &self,
1278 endpoint: ModuleEndpointId,
1279 corr: u64,
1280 ) -> Result<bool, ForwardingError> {
1281 let mut inner = self.write_inner()?;
1282 let Some(pending) = inner.pending_relays.get_mut(&(endpoint, corr)) else {
1283 return Ok(false);
1284 };
1285 pending.relay_enqueued = true;
1286 Ok(true)
1287 }
1288
1289 pub(crate) fn release_client_route(
1290 &self,
1291 client_connection_id: ConnectionId,
1292 client_channel: u16,
1293 expected_epoch: u32,
1294 ) -> Result<RouteRelease, ForwardingError> {
1295 let mut inner = self.write_inner()?;
1296 let release = release_client_route_locked(
1297 &mut inner,
1298 ClientRouteKey {
1299 connection_id: client_connection_id,
1300 channel: client_channel,
1301 },
1302 expected_epoch,
1303 );
1304 self.record_route_release(&release);
1305 Ok(release)
1306 }
1307
1308 pub(crate) fn release_module_route(
1309 &self,
1310 module_connection_id: ConnectionId,
1311 module_channel: u16,
1312 expected_epoch: u32,
1313 ) -> Result<RouteRelease, ForwardingError> {
1314 let mut inner = self.write_inner()?;
1315 let Some(endpoint) = inner
1316 .endpoint_by_connection
1317 .get(&module_connection_id)
1318 .copied()
1319 else {
1320 return Ok(RouteRelease::Absent);
1321 };
1322 let release = release_module_route_locked(
1323 &mut inner,
1324 ModuleRouteKey {
1325 endpoint,
1326 channel: module_channel,
1327 },
1328 expected_epoch,
1329 );
1330 self.record_route_release(&release);
1331 Ok(release)
1332 }
1333
1334 pub(crate) fn abort_pending_relay(
1335 &self,
1336 endpoint: ModuleEndpointId,
1337 corr: u64,
1338 outcome: RouteBindRelayOutcome,
1339 ) -> Result<Option<GoodbyeTarget>, ForwardingError> {
1340 let mut inner = self.write_inner()?;
1341 let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1342 return Ok(None);
1343 };
1344 release_reserved_route_locked(
1345 &mut inner,
1346 pending.reservation.client_key,
1347 pending.reservation.module_key,
1348 );
1349 let target = pending
1350 .relay_enqueued
1351 .then(|| abandoned_route_target(&inner, &pending.reservation));
1352 let _ = pending.sender.send(outcome);
1353 Ok(target.flatten())
1354 }
1355
1356 pub(crate) fn cancel_module_control_rpc(
1357 &self,
1358 endpoint: ModuleEndpointId,
1359 corr: u64,
1360 ) -> Result<(), ForwardingError> {
1361 self.write_inner()?
1362 .pending_control_rpcs
1363 .remove(&(endpoint, corr));
1364 Ok(())
1365 }
1366
1367 pub(crate) fn tombstone_health_probe_rpc(
1368 &self,
1369 endpoint: ModuleEndpointId,
1370 corr: u64,
1371 ) -> Result<bool, ForwardingError> {
1372 let key = (endpoint, corr);
1373 let expires_at = Instant::now() + HEALTH_PROBE_TOMBSTONE_TTL;
1374 {
1375 let mut inner = self.write_inner()?;
1376 let Some(pending) = inner.pending_control_rpcs.remove(&key) else {
1377 return Ok(false);
1378 };
1379 let Some(probe_started_at) = pending.health_probe_started_at else {
1380 inner.pending_control_rpcs.insert(key, pending);
1381 return Ok(false);
1382 };
1383 let module_id = inner
1384 .module_id_by_endpoint
1385 .get(&endpoint)
1386 .cloned()
1387 .unwrap_or_else(|| "unknown".to_string());
1388 inner.health_probe_tombstones.insert(
1389 key,
1390 HealthProbeTombstone {
1391 expected_op: pending.expected_op,
1392 module_id,
1393 probe_started_at,
1394 expires_at,
1395 },
1396 );
1397 }
1398 self.schedule_health_probe_tombstone_expiration(key, expires_at);
1399 Ok(true)
1400 }
1401
1402 fn schedule_health_probe_tombstone_expiration(
1403 &self,
1404 key: (ModuleEndpointId, u64),
1405 expires_at: Instant,
1406 ) {
1407 let inner = Arc::downgrade(&self.inner);
1408 tokio::spawn(async move {
1409 tokio::time::sleep_until(expires_at).await;
1410 let Some(inner) = inner.upgrade() else {
1411 return;
1412 };
1413 let Ok(mut inner) = inner.write() else {
1414 return;
1415 };
1416 let expired = inner
1417 .health_probe_tombstones
1418 .get(&key)
1419 .is_some_and(|tombstone| tombstone.expires_at <= Instant::now());
1420 if expired {
1421 inner.health_probe_tombstones.remove(&key);
1422 }
1423 });
1424 }
1425
1426 pub(crate) fn complete_pending_relay(
1427 &self,
1428 connection_id: ConnectionId,
1429 corr: u64,
1430 outcome: RouteBindRelayOutcome,
1431 ) -> Result<PendingRelayCompletion, ForwardingError> {
1432 let mut inner = self.write_inner()?;
1433 let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1434 return Ok(PendingRelayCompletion {
1435 settled: false,
1436 abandoned: None,
1437 });
1438 };
1439 let Some(pending) = inner.pending_relays.remove(&(endpoint, corr)) else {
1440 return Ok(PendingRelayCompletion {
1441 settled: false,
1442 abandoned: None,
1443 });
1444 };
1445
1446 if Instant::now() >= pending.deadline {
1447 release_reserved_route_locked(
1448 &mut inner,
1449 pending.reservation.client_key,
1450 pending.reservation.module_key,
1451 );
1452 let abandoned = matches!(outcome, RouteBindRelayOutcome::Accepted)
1453 .then(|| abandoned_route_target(&inner, &pending.reservation))
1454 .flatten();
1455 let _ = pending
1456 .sender
1457 .send(RouteBindRelayOutcome::Rejected(ErrorBody {
1458 code: "module_timeout".to_string(),
1459 message: "route.bind response arrived after its daemon deadline".to_string(),
1460 detail: None,
1461 }));
1462 return Ok(PendingRelayCompletion {
1463 settled: true,
1464 abandoned,
1465 });
1466 }
1467
1468 match outcome {
1469 RouteBindRelayOutcome::Accepted
1488 if pending.client_sink.is_closed()
1489 || inner
1490 .closing_connections
1491 .contains(&pending.reservation.client_key.connection_id) =>
1492 {
1493 let reason = if pending.client_sink.is_closed() {
1494 "client egress closed before route publication"
1495 } else {
1496 "client connection is closing before route publication"
1497 };
1498 release_reserved_route_locked(
1499 &mut inner,
1500 pending.reservation.client_key,
1501 pending.reservation.module_key,
1502 );
1503 let abandoned = pending
1504 .relay_enqueued
1505 .then(|| abandoned_route_target(&inner, &pending.reservation))
1506 .flatten();
1507 let _ = pending
1508 .sender
1509 .send(RouteBindRelayOutcome::ModuleGone(reason.to_string()));
1510 return Ok(PendingRelayCompletion {
1511 settled: true,
1512 abandoned,
1513 });
1514 }
1515 RouteBindRelayOutcome::Accepted
1532 if inner.superseded_endpoints.contains_key(&endpoint) =>
1533 {
1534 release_reserved_route_locked(
1535 &mut inner,
1536 pending.reservation.client_key,
1537 pending.reservation.module_key,
1538 );
1539 let abandoned = abandoned_route_target(&inner, &pending.reservation);
1542 let module_id = inner
1543 .module_id_by_endpoint
1544 .get(&endpoint)
1545 .cloned()
1546 .unwrap_or_else(|| "unknown".to_string());
1547 let _ = pending
1548 .sender
1549 .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
1550 "module_reloading",
1551 format!("module_id '{module_id}' is reloading"),
1552 )));
1553 return Ok(PendingRelayCompletion {
1554 settled: true,
1555 abandoned,
1556 });
1557 }
1558 RouteBindRelayOutcome::Accepted => {
1559 let abandoned = commit_route_locked(&mut inner, pending)?;
1560 return Ok(PendingRelayCompletion {
1561 settled: true,
1562 abandoned,
1563 });
1564 }
1565 terminal => {
1566 release_reserved_route_locked(
1567 &mut inner,
1568 pending.reservation.client_key,
1569 pending.reservation.module_key,
1570 );
1571 let _ = pending.sender.send(terminal);
1572 }
1573 }
1574 Ok(PendingRelayCompletion {
1575 settled: true,
1576 abandoned: None,
1577 })
1578 }
1579
1580 pub(crate) fn pending_module_control_op(
1581 &self,
1582 connection_id: ConnectionId,
1583 corr: u64,
1584 ) -> Result<Option<String>, ForwardingError> {
1585 let inner = self.read_inner()?;
1586 let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1587 return Ok(None);
1588 };
1589 let key = (endpoint, corr);
1590 Ok(inner
1591 .pending_control_rpcs
1592 .get(&key)
1593 .map(|pending| pending.expected_op.clone())
1594 .or_else(|| {
1595 inner
1596 .health_probe_tombstones
1597 .get(&key)
1598 .filter(|tombstone| tombstone.expires_at > Instant::now())
1599 .map(|tombstone| tombstone.expected_op.clone())
1600 }))
1601 }
1602
1603 pub(crate) fn complete_module_control_rpc(
1604 &self,
1605 connection_id: ConnectionId,
1606 corr: u64,
1607 actual_op: Option<&str>,
1608 outcome: ModuleControlRpcOutcome,
1609 ) -> Result<ModuleControlRpcCompletion, ForwardingError> {
1610 let now = Instant::now();
1611 let mut inner = self.write_inner()?;
1612 let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1613 return Ok(ModuleControlRpcCompletion::Unknown);
1614 };
1615 let key = (endpoint, corr);
1616 if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
1617 if now >= pending.deadline {
1618 let late_health_answer = pending.health_probe_started_at.map(|probe_started_at| {
1619 ModuleControlRpcCompletion::LateHealthAnswer {
1620 module_id: inner
1621 .module_id_by_endpoint
1622 .get(&endpoint)
1623 .cloned()
1624 .unwrap_or_else(|| "unknown".to_string()),
1625 latency: now.saturating_duration_since(probe_started_at),
1626 }
1627 });
1628 let _ = pending
1629 .sender
1630 .send(ModuleControlRpcOutcome::DeadlineElapsed);
1631 return Ok(late_health_answer.unwrap_or(ModuleControlRpcCompletion::Settled));
1632 }
1633 let outcome = match actual_op {
1634 Some(actual) if actual != pending.expected_op => {
1635 ModuleControlRpcOutcome::UnexpectedOp {
1636 expected: pending.expected_op,
1637 actual: actual.to_string(),
1638 }
1639 }
1640 _ => outcome,
1641 };
1642 let _ = pending.sender.send(outcome);
1643 return Ok(ModuleControlRpcCompletion::Settled);
1644 }
1645
1646 let Some(tombstone) = inner.health_probe_tombstones.remove(&key) else {
1647 return Ok(ModuleControlRpcCompletion::Unknown);
1648 };
1649 if tombstone.expires_at <= now {
1650 return Ok(ModuleControlRpcCompletion::Unknown);
1651 }
1652 Ok(ModuleControlRpcCompletion::LateHealthAnswer {
1653 module_id: tombstone.module_id,
1654 latency: now.saturating_duration_since(tombstone.probe_started_at),
1655 })
1656 }
1657
1658 #[cfg(test)]
1659 pub(crate) fn health_probe_tombstone_count(&self) -> Result<usize, ForwardingError> {
1660 Ok(self.read_inner()?.health_probe_tombstones.len())
1661 }
1662
1663 #[cfg(test)]
1664 pub(crate) fn closing_connection_count(&self) -> Result<usize, ForwardingError> {
1665 Ok(self.read_inner()?.closing_connections.len())
1666 }
1667
1668 #[cfg(test)]
1671 pub(crate) fn reserved_route_count(&self) -> Result<(usize, usize), ForwardingError> {
1672 let inner = self.read_inner()?;
1673 Ok((inner.reserved_client.len(), inner.reserved_module.len()))
1674 }
1675
1676 pub(crate) fn module_endpoint_for_connection(
1677 &self,
1678 connection_id: ConnectionId,
1679 ) -> Result<Option<ModuleEndpointId>, ForwardingError> {
1680 Ok(self
1681 .read_inner()?
1682 .endpoint_by_connection
1683 .get(&connection_id)
1684 .copied())
1685 }
1686
1687 pub(crate) fn module_id_for_connection(
1690 &self,
1691 connection_id: ConnectionId,
1692 ) -> Result<Option<String>, ForwardingError> {
1693 let inner = self.read_inner()?;
1694 Ok(inner
1695 .endpoint_by_connection
1696 .get(&connection_id)
1697 .and_then(|endpoint| inner.module_id_by_endpoint.get(endpoint))
1698 .cloned())
1699 }
1700
1701 pub(crate) fn module_route_epoch_was_allocated(
1708 &self,
1709 connection_id: ConnectionId,
1710 channel: u16,
1711 epoch: u32,
1712 ) -> Result<bool, ForwardingError> {
1713 let inner = self.read_inner()?;
1714 let Some(endpoint) = inner.endpoint_by_connection.get(&connection_id).copied() else {
1715 return Ok(false);
1716 };
1717 Ok(inner
1718 .module_slot_epochs
1719 .get(&ModuleRouteKey { endpoint, channel })
1720 .is_some_and(|last| epoch != 0 && epoch <= *last))
1721 }
1722
1723 pub(crate) fn has_live_module_connection(
1724 &self,
1725 module_id: &str,
1726 ) -> Result<bool, ForwardingError> {
1727 Ok(self.read_inner()?.modules_by_id.contains_key(module_id))
1728 }
1729
1730 pub(crate) fn lookup_data_route(
1731 &self,
1732 connection_id: ConnectionId,
1733 channel: u16,
1734 epoch: u32,
1735 ) -> Result<DataRoute, ForwardingError> {
1736 let inner = self.read_inner()?;
1737 let state = if let Some(endpoint) =
1738 inner.endpoint_by_connection.get(&connection_id).copied()
1739 {
1740 let key = ModuleRouteKey { endpoint, channel };
1741 match inner.module_to_client.get(&key) {
1742 Some(route) if route.module_epoch == epoch => {
1743 DataRouteState::Bound(Arc::clone(route))
1744 }
1745 Some(_) => DataRouteState::EpochMismatch,
1746 None if inner.reserved_module.contains_key(&key)
1747 && inner.module_slot_epochs.get(&key).copied() == Some(epoch) =>
1748 {
1749 DataRouteState::Reserved
1750 }
1751 None if inner.reserved_module.contains_key(&key) => DataRouteState::EpochMismatch,
1752 None => DataRouteState::Absent,
1753 }
1754 } else {
1755 let key = ClientRouteKey {
1756 connection_id,
1757 channel,
1758 };
1759 match inner.client_to_module.get(&key) {
1760 Some(route) if route.client_epoch == epoch => {
1761 DataRouteState::Bound(Arc::clone(route))
1762 }
1763 Some(_) => DataRouteState::EpochMismatch,
1764 None if inner.reserved_client.contains_key(&key)
1765 && inner.client_slot_epochs.get(&key).copied() == Some(epoch) =>
1766 {
1767 DataRouteState::Reserved
1768 }
1769 None if inner.reserved_client.contains_key(&key) => DataRouteState::EpochMismatch,
1770 None => DataRouteState::Absent,
1771 }
1772 };
1773 Ok(
1774 if inner.endpoint_by_connection.contains_key(&connection_id) {
1775 DataRoute::Module(state)
1776 } else {
1777 DataRoute::Client(state)
1778 },
1779 )
1780 }
1781
1782 #[cfg(test)]
1783 pub(crate) fn inject_client_slot_epoch(
1784 &self,
1785 connection_id: ConnectionId,
1786 channel: u16,
1787 last_epoch: u32,
1788 ) {
1789 let mut inner = self.write_inner().expect("forwarding lock");
1790 inner.client_slot_epochs.insert(
1791 ClientRouteKey {
1792 connection_id,
1793 channel,
1794 },
1795 last_epoch,
1796 );
1797 inner.next_client_channel.insert(connection_id, channel);
1798 }
1799
1800 #[cfg(test)]
1801 pub(crate) fn inject_module_slot_epoch(
1802 &self,
1803 endpoint: ModuleEndpointId,
1804 channel: u16,
1805 last_epoch: u32,
1806 ) {
1807 let mut inner = self.write_inner().expect("forwarding lock");
1808 inner
1809 .module_slot_epochs
1810 .insert(ModuleRouteKey { endpoint, channel }, last_epoch);
1811 inner.next_module_channel.insert(endpoint, channel);
1812 }
1813
1814 #[cfg(test)]
1815 pub(crate) fn inject_control_corr(&self, endpoint: ModuleEndpointId, next_corr: u64) {
1816 self.write_inner()
1817 .expect("forwarding lock")
1818 .next_control_corr
1819 .insert(endpoint, next_corr);
1820 }
1821
1822 pub(crate) fn cache_status(
1823 &self,
1824 endpoint: ModuleEndpointId,
1825 module_channel: u16,
1826 module_epoch: u32,
1827 status: String,
1828 ) -> Result<bool, ForwardingError> {
1829 let mut inner = self.write_inner()?;
1830 if !inner.module_id_by_endpoint.contains_key(&endpoint) {
1831 return Err(ForwardingError::StaleModuleEndpoint);
1832 }
1833
1834 let module_key = ModuleRouteKey {
1835 endpoint,
1836 channel: module_channel,
1837 };
1838 let handle = if let Some(route) = inner.module_to_client.get(&module_key) {
1839 (route.module_epoch == module_epoch).then_some((
1840 ClientRouteKey {
1841 connection_id: route.client_connection_id,
1842 channel: route.client_channel,
1843 },
1844 route.client_epoch,
1845 ))
1846 } else if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
1847 (inner.module_slot_epochs.get(&module_key).copied() == Some(module_epoch)).then_some((
1848 client_key,
1849 inner
1850 .client_slot_epochs
1851 .get(&client_key)
1852 .copied()
1853 .unwrap_or(0),
1854 ))
1855 } else {
1856 None
1857 };
1858
1859 if let Some(handle) = handle {
1860 inner.status.insert(handle, status);
1861 Ok(true)
1862 } else {
1863 debug!(
1864 module_channel,
1865 module_epoch,
1866 generation = endpoint.generation,
1867 connection_id = endpoint.connection_id.get(),
1868 "dropping stale status update for module route handle"
1869 );
1870 Ok(false)
1871 }
1872 }
1873
1874 pub(crate) fn route_poll_snapshot(
1875 &self,
1876 client_connection_id: ConnectionId,
1877 client_channel: u16,
1878 client_epoch: u32,
1879 ) -> Result<RoutePollSnapshot, ForwardingError> {
1880 let inner = self.read_inner()?;
1881 let client_key = ClientRouteKey {
1882 connection_id: client_connection_id,
1883 channel: client_channel,
1884 };
1885 let Some(route) = inner.client_to_module.get(&client_key) else {
1886 return Ok(RoutePollSnapshot::Absent);
1887 };
1888 if route.client_epoch != client_epoch
1889 || !inner
1890 .module_id_by_endpoint
1891 .contains_key(&route.module_endpoint)
1892 {
1893 return Ok(RoutePollSnapshot::Absent);
1894 }
1895 Ok(RoutePollSnapshot::Bound {
1896 module_id: route.module_id.clone(),
1897 status: inner.status.get(&(client_key, client_epoch)).cloned(),
1898 })
1899 }
1900
1901 pub fn active_binding_count(&self) -> Result<usize, ForwardingError> {
1902 Ok(self.read_inner()?.client_to_module.len())
1903 }
1904
1905 pub fn client_route_concentration(&self) -> Result<(usize, usize), ForwardingError> {
1915 let inner = self.read_inner()?;
1916 let mut per_connection: HashMap<ConnectionId, usize> = HashMap::new();
1917 for key in inner.client_to_module.keys() {
1918 *per_connection.entry(key.connection_id).or_insert(0) += 1;
1919 }
1920 let max = per_connection.values().copied().max().unwrap_or(0);
1921 Ok((per_connection.len(), max))
1922 }
1923
1924 pub fn has_route_channel(&self, route_channel: u16) -> Result<bool, ForwardingError> {
1925 let inner = self.read_inner()?;
1926 Ok(inner
1927 .client_to_module
1928 .keys()
1929 .any(|key| key.channel == route_channel))
1930 }
1931
1932 pub(crate) fn is_daemon_draining(&self) -> Result<bool, ForwardingError> {
1934 Ok(self.read_inner()?.daemon_draining)
1935 }
1936
1937 #[cfg(unix)]
1940 pub(crate) fn begin_daemon_drain(&self) -> Result<Vec<String>, ForwardingError> {
1941 let mut inner = self.write_inner()?;
1942 inner.daemon_draining = true;
1943 let modules = inner
1944 .modules_by_id
1945 .iter()
1946 .map(|(id, module)| (id.clone(), module.endpoint))
1947 .collect::<Vec<_>>();
1948 for (_, endpoint) in &modules {
1949 inner
1950 .draining_endpoints
1951 .insert(*endpoint, RouteCloseReason::Restart);
1952 }
1953 let off_slot_endpoints = inner
1958 .candidates_by_id
1959 .values()
1960 .map(|module| module.endpoint)
1961 .chain(inner.superseded_endpoints.keys().copied())
1962 .collect::<Vec<_>>();
1963 for endpoint in off_slot_endpoints {
1964 inner
1965 .draining_endpoints
1966 .insert(endpoint, RouteCloseReason::Restart);
1967 }
1968 Ok(modules.into_iter().map(|(id, _)| id).collect())
1969 }
1970
1971 pub(crate) fn begin_module_drain(
1978 &self,
1979 module_id: &str,
1980 reason: RouteCloseReason,
1981 ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
1982 let mut inner = self.write_inner()?;
1983 let Some(module) = inner.modules_by_id.get(module_id).cloned() else {
1984 return Ok(None);
1985 };
1986 Ok(Some(begin_drain_locked(
1987 &mut inner, module_id, module, reason,
1988 )))
1989 }
1990
1991 pub(crate) fn begin_endpoint_drain(
1998 &self,
1999 endpoint: ModuleEndpointId,
2000 reason: RouteCloseReason,
2001 ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
2002 let mut inner = self.write_inner()?;
2003 let Some(module) = module_connection_for_endpoint_locked(&inner, endpoint).cloned() else {
2004 return Ok(None);
2005 };
2006 let module_id = inner
2007 .module_id_by_endpoint
2008 .get(&endpoint)
2009 .cloned()
2010 .expect("an endpoint resolved to a module connection has a module id");
2011 Ok(Some(begin_drain_locked(
2012 &mut inner, &module_id, module, reason,
2013 )))
2014 }
2015}
2016
2017fn begin_drain_locked(
2021 inner: &mut ForwardingInner,
2022 module_id: &str,
2023 module: ModuleConnection,
2024 reason: RouteCloseReason,
2025) -> ModuleDrainTarget {
2026 {
2027 let endpoint = module.endpoint;
2028 inner.draining_endpoints.insert(endpoint, reason);
2029
2030 let flows = inner
2031 .client_to_module
2032 .values()
2033 .filter(|route| route.module_endpoint == endpoint)
2034 .map(|route| Arc::clone(&route.flow))
2035 .collect::<Vec<_>>();
2036 let excluded_subscriptions = flows
2037 .into_iter()
2038 .map(|flow| flow.begin_drain())
2039 .fold(0u32, u32::saturating_add);
2040
2041 let pending_keys = inner
2042 .pending_relays
2043 .keys()
2044 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2045 .copied()
2046 .collect::<Vec<_>>();
2047 let mut abandoned_bindings = Vec::new();
2048 for key in pending_keys {
2049 let Some(pending) = inner.pending_relays.remove(&key) else {
2050 continue;
2051 };
2052 release_reserved_route_locked(
2053 inner,
2054 pending.reservation.client_key,
2055 pending.reservation.module_key,
2056 );
2057 if pending.relay_enqueued {
2058 if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2059 abandoned_bindings.push(target);
2060 }
2061 }
2062 let _ = pending
2063 .sender
2064 .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
2065 "module_reloading",
2066 format!("module_id '{module_id}' is reloading"),
2067 )));
2068 }
2069
2070 let pending_control_keys = inner
2071 .pending_control_rpcs
2072 .keys()
2073 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2074 .copied()
2075 .collect::<Vec<_>>();
2076 for key in pending_control_keys {
2077 if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
2078 let _ = pending
2079 .sender
2080 .send(ModuleControlRpcOutcome::ModuleGone(format!(
2081 "module '{module_id}' began draining during module-control RPC"
2082 )));
2083 }
2084 }
2085
2086 ModuleDrainTarget {
2087 endpoint,
2088 sink: module.sink,
2089 negotiated_ver: module.negotiated_ver,
2090 abandoned_bindings,
2091 excluded_subscriptions,
2092 }
2093 }
2094}
2095
2096#[derive(Debug, Default, PartialEq, Eq)]
2099pub(crate) struct DrainHoldouts {
2100 pub(crate) requests: usize,
2103 pub(crate) routes: usize,
2105 pub(crate) total_routes: usize,
2107 pub(crate) top_connections: Vec<(u64, usize)>,
2110 pub(crate) held: Vec<(u16, u64)>,
2119}
2120
2121pub(crate) const DRAIN_HELD_REQUESTS_LISTED: usize = 32;
2123
2124impl ForwardingTable {
2125 pub(crate) fn endpoint_drain_holdouts(
2127 &self,
2128 endpoint: ModuleEndpointId,
2129 ) -> Result<DrainHoldouts, ForwardingError> {
2130 let inner = self.read_inner()?;
2131 let mut holdouts = DrainHoldouts::default();
2132 let mut by_connection: HashMap<u64, usize> = HashMap::new();
2133 for (key, route) in &inner.client_to_module {
2134 if route.module_endpoint != endpoint {
2135 continue;
2136 }
2137 holdouts.total_routes += 1;
2138 let held = route.flow.drain_in_flight();
2139 if held == 0 {
2140 continue;
2141 }
2142 holdouts.requests += held;
2143 holdouts.routes += 1;
2144 *by_connection.entry(key.connection_id.get()).or_default() += held;
2145 holdouts.held.extend(
2146 route
2147 .flow
2148 .drain_held_corrs()
2149 .into_iter()
2150 .map(|corr| (route.module_channel, corr)),
2151 );
2152 }
2153 holdouts.held.sort_unstable();
2154 holdouts.held.truncate(DRAIN_HELD_REQUESTS_LISTED);
2155 let mut connections = by_connection.into_iter().collect::<Vec<_>>();
2156 connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
2157 connections.truncate(3);
2158 holdouts.top_connections = connections;
2159 Ok(holdouts)
2160 }
2161
2162 pub(crate) fn endpoint_in_flight_count(
2163 &self,
2164 endpoint: ModuleEndpointId,
2165 ) -> Result<usize, ForwardingError> {
2166 let inner = self.read_inner()?;
2167 Ok(inner
2168 .client_to_module
2169 .values()
2170 .filter(|route| route.module_endpoint == endpoint)
2171 .map(|route| route.flow.drain_in_flight())
2172 .sum())
2173 }
2174
2175 pub(crate) fn endpoint_is_draining(
2176 &self,
2177 endpoint: ModuleEndpointId,
2178 ) -> Result<bool, ForwardingError> {
2179 Ok(self
2180 .read_inner()?
2181 .draining_endpoints
2182 .contains_key(&endpoint))
2183 }
2184
2185 pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
2186 let inner = self.read_inner()?;
2187 Ok(inner
2188 .modules_by_id
2189 .get(module_id)
2190 .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
2191 }
2192
2193 pub(crate) fn release_module_endpoint_routes(
2194 &self,
2195 endpoint: ModuleEndpointId,
2196 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2197 let mut inner = self.write_inner()?;
2198 let routes = inner
2199 .module_to_client
2200 .iter()
2201 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2202 .map(|(module_key, route)| (*module_key, route.module_epoch))
2203 .collect::<Vec<_>>();
2204 let mut released = Vec::with_capacity(routes.len());
2205 for (module_key, epoch) in routes {
2206 if let RouteRelease::Removed(target) =
2207 release_module_route_locked(&mut inner, module_key, epoch)
2208 {
2209 released.push(target);
2210 }
2211 }
2212 Ok(released)
2213 }
2214
2215 pub(crate) fn endpoint_routes(
2221 &self,
2222 endpoint: ModuleEndpointId,
2223 ) -> Result<Vec<EndpointRoute>, ForwardingError> {
2224 let inner = self.read_inner()?;
2225 Ok(endpoint_routes_locked(&inner, endpoint))
2226 }
2227
2228 pub(crate) fn route_census(
2230 &self,
2231 module_id: Option<&str>,
2232 ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
2233 let inner = self.read_inner()?;
2234 let mut endpoints = inner
2235 .modules_by_id
2236 .iter()
2237 .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
2238 .map(|(id, module)| (id.clone(), module.endpoint))
2239 .collect::<Vec<_>>();
2240 endpoints.sort_by(|left, right| left.0.cmp(&right.0));
2241 Ok(endpoints
2242 .into_iter()
2243 .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
2244 .collect())
2245 }
2246
2247 pub(crate) fn live_roots(
2249 &self,
2250 module_id: &str,
2251 ) -> Result<ModuleControlResponseToModule, ForwardingError> {
2252 let inner = self.read_inner()?;
2253 let endpoint = inner
2254 .modules_by_id
2255 .get(module_id)
2256 .map(|module| module.endpoint);
2257 let mut roots = BTreeMap::new();
2258 let mut unknown_root_bindings = 0;
2259 let mut total_bindings = 0;
2260 if let Some(endpoint) = endpoint {
2261 for binding in inner
2262 .module_to_client
2263 .values()
2264 .filter(|binding| binding.module_endpoint == endpoint)
2265 {
2266 total_bindings += 1;
2267 if let Some(root) = &binding.project_root {
2268 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2269 entry.0 += 1;
2270 } else {
2271 unknown_root_bindings += 1;
2272 }
2273 }
2274 for pending in inner
2275 .pending_relays
2276 .values()
2277 .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
2278 {
2279 total_bindings += 1;
2280 if let Some(root) = &pending.reservation.project_root {
2281 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2282 entry.1 += 1;
2283 } else {
2284 unknown_root_bindings += 1;
2285 }
2286 }
2287 }
2288 Ok(ModuleControlResponseToModule::LiveRoots {
2289 roots: roots
2290 .into_iter()
2291 .map(|(project_root, (bound, pending))| LiveRoot {
2292 project_root,
2293 bound,
2294 pending,
2295 })
2296 .collect(),
2297 unknown_root_bindings,
2298 total_bindings,
2299 })
2300 }
2301
2302 pub(crate) fn connection_has_client_routes(
2308 &self,
2309 connection_id: ConnectionId,
2310 ) -> Result<bool, ForwardingError> {
2311 let inner = self.read_inner()?;
2312 let has = inner
2313 .client_to_module
2314 .keys()
2315 .any(|key| key.connection_id == connection_id)
2316 || inner
2317 .reserved_client
2318 .keys()
2319 .any(|key| key.connection_id == connection_id);
2320 Ok(has)
2321 }
2322
2323 pub(crate) fn cleanup_connection(
2324 &self,
2325 connection_id: ConnectionId,
2326 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2327 self.cleanup_connection_counted(connection_id)
2328 .map(|cleanup| cleanup.released)
2329 }
2330
2331 pub(crate) fn cleanup_connection_counted(
2336 &self,
2337 connection_id: ConnectionId,
2338 ) -> Result<ConnectionCleanup, ForwardingError> {
2339 let mut inner = self.write_inner()?;
2340 inner.closing_connections.insert(connection_id);
2341 let cleanup = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
2342 remove_module_connection_locked(&mut inner, endpoint)
2343 } else {
2344 ConnectionCleanup {
2345 released: Self::cleanup_client_connection_locked(&mut inner, connection_id),
2346 abandoned_relays: 0,
2347 }
2348 };
2349 inner.closing_connections.remove(&connection_id);
2358 Ok(cleanup)
2359 }
2360
2361 fn cleanup_client_connection_locked(
2362 inner: &mut ForwardingInner,
2363 connection_id: ConnectionId,
2364 ) -> Vec<GoodbyeTarget> {
2365 let routes = inner
2366 .client_to_module
2367 .iter()
2368 .filter(|(key, _)| key.connection_id == connection_id)
2369 .map(|(key, route)| (*key, route.client_epoch))
2370 .collect::<Vec<_>>();
2371 let mut released = Vec::with_capacity(routes.len());
2372 for (client_key, epoch) in routes {
2373 if let RouteRelease::Removed(target) =
2374 release_client_route_locked(inner, client_key, epoch)
2375 {
2376 released.push(target);
2377 }
2378 }
2379
2380 let pending_keys = inner
2381 .pending_relays
2382 .iter()
2383 .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2384 .map(|(key, _)| *key)
2385 .collect::<Vec<_>>();
2386 for key in pending_keys {
2387 let Some(pending) = inner.pending_relays.remove(&key) else {
2388 continue;
2389 };
2390 release_reserved_route_locked(
2391 inner,
2392 pending.reservation.client_key,
2393 pending.reservation.module_key,
2394 );
2395 if pending.relay_enqueued {
2396 if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2397 released.push(target);
2398 }
2399 }
2400 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2401 "client connection closed during route.bind relay".to_string(),
2402 ));
2403 }
2404
2405 let orphaned = inner
2406 .reserved_client
2407 .iter()
2408 .filter(|(key, _)| key.connection_id == connection_id)
2409 .map(|(client, module)| (*client, *module))
2410 .collect::<Vec<_>>();
2411 for (client_key, module_key) in orphaned {
2412 release_reserved_route_locked(inner, client_key, module_key);
2413 }
2414 inner.next_client_channel.remove(&connection_id);
2415 inner
2416 .client_slot_epochs
2417 .retain(|key, _| key.connection_id != connection_id);
2418 inner
2419 .last_published_epoch
2420 .retain(|key, _| key.connection_id != connection_id);
2421 inner
2422 .status
2423 .retain(|(key, _), _| key.connection_id != connection_id);
2424
2425 released
2426 }
2427
2428 pub(crate) fn escalate_client_delivery_failure(
2437 &self,
2438 connection_id: ConnectionId,
2439 channel: u16,
2440 expected_epoch: u32,
2441 reason: CloseReason,
2442 undelivered: UndeliveredFrame<'_>,
2443 ) -> Result<bool, ForwardingError> {
2444 let principals = {
2445 let mut inner = self.write_inner()?;
2446 let key = ClientRouteKey {
2447 connection_id,
2448 channel,
2449 };
2450 if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2451 None
2452 } else {
2453 inner.closing_connections.insert(connection_id);
2454 Some(connection_principals_locked(&inner, connection_id))
2455 }
2456 };
2457 let Some(principals) = principals else {
2458 return Ok(false);
2459 };
2460 let backlog = undelivered.sink.backlog();
2461 let close_reason = reason.to_string();
2462 if self.request_connection_close(connection_id, reason) {
2463 warn!(
2464 connection_id = connection_id.get(),
2465 principals = %principals,
2466 module_id = undelivered.module_id.unwrap_or("unknown"),
2467 client_channel = channel,
2468 queued_bytes = backlog.queued_bytes,
2469 queued_frames = backlog.queued_frames,
2470 oldest_queued_ms = backlog
2471 .oldest_age
2472 .map(|age| age.as_millis() as u64)
2473 .unwrap_or(0),
2474 close_reason = %close_reason,
2475 "closing client connection: its egress queue could not take a frame"
2476 );
2477 }
2478 Ok(true)
2479 }
2480
2481 fn record_route_release(&self, release: &RouteRelease) {
2482 match release {
2483 RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2484 RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2485 RouteRelease::Absent => {}
2486 }
2487 }
2488
2489 fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2490 self.inner.read().map_err(|_| ForwardingError::Poisoned)
2491 }
2492
2493 fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2494 self.inner.write().map_err(|_| ForwardingError::Poisoned)
2495 }
2496
2497 fn lock_close_registry(
2498 &self,
2499 ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2500 self.close_registry
2501 .lock()
2502 .unwrap_or_else(|poisoned| poisoned.into_inner())
2503 }
2504}
2505
2506impl ForwardingInner {
2507 fn allocate_route_slots(
2508 &mut self,
2509 connection_id: ConnectionId,
2510 endpoint: ModuleEndpointId,
2511 ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2512 let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2513 let mut client_channel = client_start;
2514 let client_channel = loop {
2515 let key = ClientRouteKey {
2516 connection_id,
2517 channel: client_channel,
2518 };
2519 let eligible = !self.client_to_module.contains_key(&key)
2520 && !self.reserved_client.contains_key(&key)
2521 && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2522 if eligible {
2523 break client_channel;
2524 }
2525 client_channel = next_channel(client_channel);
2526 if client_channel == client_start {
2527 return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2528 }
2529 };
2530
2531 let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2532 let mut module_channel = module_start;
2533 let module_channel = loop {
2534 let key = ModuleRouteKey {
2535 endpoint,
2536 channel: module_channel,
2537 };
2538 let eligible = !self.module_to_client.contains_key(&key)
2539 && !self.reserved_module.contains_key(&key)
2540 && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2541 if eligible {
2542 break module_channel;
2543 }
2544 module_channel = next_channel(module_channel);
2545 if module_channel == module_start {
2546 return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2547 }
2548 };
2549
2550 let client_key = ClientRouteKey {
2551 connection_id,
2552 channel: client_channel,
2553 };
2554 let module_key = ModuleRouteKey {
2555 endpoint,
2556 channel: module_channel,
2557 };
2558 let client_epoch = self
2559 .client_slot_epochs
2560 .get(&client_key)
2561 .copied()
2562 .unwrap_or(0)
2563 + 1;
2564 let module_epoch = self
2565 .module_slot_epochs
2566 .get(&module_key)
2567 .copied()
2568 .unwrap_or(0)
2569 + 1;
2570 self.client_slot_epochs.insert(client_key, client_epoch);
2571 self.module_slot_epochs.insert(module_key, module_epoch);
2572 self.next_client_channel
2573 .insert(connection_id, next_channel(client_channel));
2574 self.next_module_channel
2575 .insert(endpoint, next_channel(module_channel));
2576 Ok((client_channel, client_epoch, module_channel, module_epoch))
2577 }
2578
2579 fn allocate_control_corr(
2580 &mut self,
2581 endpoint: ModuleEndpointId,
2582 ) -> Result<u64, ForwardingError> {
2583 let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2584 if candidate == 0 {
2585 self.closing_connections.insert(endpoint.connection_id);
2586 return Err(ForwardingError::RelayCorrelationExhausted);
2587 }
2588 self.next_control_corr.insert(
2589 endpoint,
2590 if candidate == u64::MAX {
2591 0
2592 } else {
2593 candidate + 1
2594 },
2595 );
2596 Ok(candidate)
2597 }
2598}
2599
2600fn next_channel(channel: u16) -> u16 {
2601 let next = channel.wrapping_add(1);
2602 if next == 0 {
2603 1
2604 } else {
2605 next
2606 }
2607}
2608
2609fn endpoint_routes_locked(
2610 inner: &ForwardingInner,
2611 endpoint: ModuleEndpointId,
2612) -> Vec<EndpointRoute> {
2613 let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2614 let draining = drain_reason.is_some();
2615 let mut routes = inner
2616 .module_to_client
2617 .iter()
2618 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2619 .map(|(_, route)| EndpointRoute {
2620 goodbye_target: GoodbyeTarget {
2621 connection_id: route.client_connection_id,
2622 sink: route.client_sink.clone(),
2623 negotiated_ver: route.client_negotiated_ver,
2624 channel: route.client_channel,
2625 epoch: route.client_epoch,
2626 kind: GoodbyeTargetKind::Client,
2627 module_id: Some(route.module_id.clone()),
2628 },
2629 principal: route.principal.clone(),
2630 bound_at: route.bound_at,
2631 draining,
2632 drain_reason,
2633 })
2634 .collect::<Vec<_>>();
2635 routes.sort_by_key(|route| {
2636 (
2637 route.goodbye_target.connection_id.get(),
2638 route.goodbye_target.channel,
2639 route.goodbye_target.epoch,
2640 )
2641 });
2642 routes
2643}
2644
2645fn release_reserved_route_locked(
2646 inner: &mut ForwardingInner,
2647 client_key: ClientRouteKey,
2648 module_key: ModuleRouteKey,
2649) {
2650 if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
2651 inner.reserved_client.remove(&client_key);
2652 }
2653 if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
2654 inner.reserved_module.remove(&module_key);
2655 }
2656 inner.status.retain(|(key, _), _| *key != client_key);
2657}
2658
2659fn release_client_route_locked(
2660 inner: &mut ForwardingInner,
2661 client_key: ClientRouteKey,
2662 expected_epoch: u32,
2663) -> RouteRelease {
2664 let Some(route) = inner.client_to_module.get(&client_key) else {
2665 return RouteRelease::Absent;
2666 };
2667 if route.client_epoch != expected_epoch {
2668 return RouteRelease::Stale;
2669 }
2670 let route = inner
2671 .client_to_module
2672 .remove(&client_key)
2673 .expect("route checked under the same forwarding lock");
2674 route.flow.close();
2675 inner.module_to_client.remove(&ModuleRouteKey {
2676 endpoint: route.module_endpoint,
2677 channel: route.module_channel,
2678 });
2679 inner.status.remove(&(client_key, expected_epoch));
2680 RouteRelease::Removed(GoodbyeTarget {
2681 connection_id: route.module_endpoint.connection_id,
2682 sink: route.module_sink.clone(),
2683 negotiated_ver: route.module_negotiated_ver,
2684 channel: route.module_channel,
2685 epoch: route.module_epoch,
2686 kind: GoodbyeTargetKind::Module,
2687 module_id: Some(route.module_id.clone()),
2688 })
2689}
2690
2691fn release_module_route_locked(
2692 inner: &mut ForwardingInner,
2693 module_key: ModuleRouteKey,
2694 expected_epoch: u32,
2695) -> RouteRelease {
2696 let Some(route) = inner.module_to_client.get(&module_key) else {
2697 return RouteRelease::Absent;
2698 };
2699 if route.module_epoch != expected_epoch {
2700 return RouteRelease::Stale;
2701 }
2702 let route = inner
2703 .module_to_client
2704 .remove(&module_key)
2705 .expect("route checked under the same forwarding lock");
2706 route.flow.close();
2707 let client_key = ClientRouteKey {
2708 connection_id: route.client_connection_id,
2709 channel: route.client_channel,
2710 };
2711 inner.client_to_module.remove(&client_key);
2712 inner.status.remove(&(client_key, route.client_epoch));
2713 RouteRelease::Removed(GoodbyeTarget {
2714 connection_id: route.client_connection_id,
2715 sink: route.client_sink.clone(),
2716 negotiated_ver: route.client_negotiated_ver,
2717 channel: route.client_channel,
2718 epoch: route.client_epoch,
2719 kind: GoodbyeTargetKind::Client,
2720 module_id: Some(route.module_id.clone()),
2721 })
2722}
2723
2724fn commit_route_locked(
2725 inner: &mut ForwardingInner,
2726 pending: PendingRouteBindRelayEntry,
2727) -> Result<Option<GoodbyeTarget>, ForwardingError> {
2728 let reservation = pending.reservation;
2729 if inner
2730 .closing_connections
2731 .contains(&reservation.client_key.connection_id)
2732 {
2733 return Err(ForwardingError::ConnectionClosing {
2734 connection_id: reservation.client_key.connection_id,
2735 });
2736 }
2737 let module_id = inner
2738 .module_id_by_endpoint
2739 .get(&reservation.module_key.endpoint)
2740 .cloned()
2741 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2742 if inner
2743 .draining_endpoints
2744 .contains_key(&reservation.module_key.endpoint)
2745 {
2746 return Err(ForwardingError::ModuleReloading { module_id });
2747 }
2748 if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
2749 || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
2750 {
2751 return Err(ForwardingError::UnknownReservation {
2752 client_channel: reservation.client_key.channel,
2753 module_channel: reservation.module_key.channel,
2754 });
2755 }
2756 let module = inner
2757 .modules_by_id
2758 .get(&module_id)
2759 .filter(|module| module.endpoint == reservation.module_key.endpoint)
2760 .cloned()
2761 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2762 let binding = Arc::new(RouteBinding {
2763 client_connection_id: reservation.client_key.connection_id,
2764 client_sink: pending.client_sink,
2765 client_negotiated_ver: pending.client_negotiated_ver,
2766 client_channel: reservation.client_key.channel,
2767 client_epoch: reservation.client_epoch,
2768 module_id,
2769 module_endpoint: reservation.module_key.endpoint,
2770 module_sink: module.sink,
2771 module_negotiated_ver: module.negotiated_ver,
2772 module_channel: reservation.module_key.channel,
2773 module_epoch: reservation.module_epoch,
2774 principal: pending.principal,
2775 project_root: reservation.project_root.clone(),
2776 bound_at: Instant::now(),
2777 flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
2778 });
2779 inner
2780 .client_to_module
2781 .insert(reservation.client_key, Arc::clone(&binding));
2782 inner
2783 .module_to_client
2784 .insert(reservation.module_key, binding);
2785 let previous_published = inner
2786 .last_published_epoch
2787 .insert(reservation.client_key, reservation.client_epoch);
2788
2789 let client_writer_closed = pending.client_permit.send(pending.route_open_frame);
2794 if client_writer_closed {
2795 let abandoned = pending
2796 .relay_enqueued
2797 .then(|| abandoned_route_target(inner, &reservation))
2798 .flatten();
2799 if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
2800 route.flow.close();
2801 }
2802 inner.module_to_client.remove(&reservation.module_key);
2803 inner
2804 .status
2805 .remove(&(reservation.client_key, reservation.client_epoch));
2806 match previous_published {
2807 Some(epoch) => {
2808 inner
2809 .last_published_epoch
2810 .insert(reservation.client_key, epoch);
2811 }
2812 None => {
2813 inner.last_published_epoch.remove(&reservation.client_key);
2814 }
2815 }
2816 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2817 "client egress closed during route publication".to_string(),
2818 ));
2819 return Ok(abandoned);
2820 }
2821
2822 let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
2823 Ok(None)
2824}
2825
2826fn module_connection_for_endpoint_locked(
2834 inner: &ForwardingInner,
2835 endpoint: ModuleEndpointId,
2836) -> Option<&ModuleConnection> {
2837 let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
2838 inner
2839 .modules_by_id
2840 .get(module_id)
2841 .filter(|module| module.endpoint == endpoint)
2842 .or_else(|| {
2843 inner
2844 .candidates_by_id
2845 .get(module_id)
2846 .filter(|module| module.endpoint == endpoint)
2847 })
2848 .or_else(|| inner.superseded_endpoints.get(&endpoint))
2849}
2850
2851fn abandoned_route_target(
2852 inner: &ForwardingInner,
2853 reservation: &RouteReservation,
2854) -> Option<GoodbyeTarget> {
2855 let module_id = inner
2856 .module_id_by_endpoint
2857 .get(&reservation.module_key.endpoint)?;
2858 let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
2859 (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
2860 connection_id: module.endpoint.connection_id,
2861 sink: module.sink.clone(),
2862 negotiated_ver: module.negotiated_ver,
2863 channel: reservation.module_key.channel,
2864 epoch: reservation.module_epoch,
2865 kind: GoodbyeTargetKind::Module,
2866 module_id: Some(module_id.clone()),
2867 })
2868}
2869
2870fn enqueue_hello_ack_locked(
2875 sink: &FrameSink,
2876 connection_id: ConnectionId,
2877 hello_ack: Option<Frame>,
2878) -> Result<(), ForwardingError> {
2879 let Some(hello_ack) = hello_ack else {
2880 return Ok(());
2881 };
2882 sink.try_send(hello_ack)
2883 .map_err(|_| ForwardingError::ModuleEgressUnavailable { connection_id })
2884}
2885
2886fn remove_module_connection_locked(
2887 inner: &mut ForwardingInner,
2888 endpoint: ModuleEndpointId,
2889) -> ConnectionCleanup {
2890 inner.draining_endpoints.remove(&endpoint);
2891 let module_id = inner.module_id_by_endpoint.remove(&endpoint);
2892 if let Some(module_id) = module_id.as_ref() {
2893 if inner
2894 .modules_by_id
2895 .get(module_id)
2896 .is_some_and(|module| module.endpoint == endpoint)
2897 {
2898 inner.modules_by_id.remove(module_id);
2899 }
2900 if inner
2901 .candidates_by_id
2902 .get(module_id)
2903 .is_some_and(|module| module.endpoint == endpoint)
2904 {
2905 inner.candidates_by_id.remove(module_id);
2906 }
2907 }
2908 inner.superseded_endpoints.remove(&endpoint);
2909 inner.endpoint_by_connection.remove(&endpoint.connection_id);
2910 inner.next_module_channel.remove(&endpoint);
2911 inner.next_control_corr.remove(&endpoint);
2912 inner
2913 .health_probe_tombstones
2914 .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
2915 inner
2916 .module_slot_epochs
2917 .retain(|key, _| key.endpoint != endpoint);
2918 let reserved_module_keys: Vec<ModuleRouteKey> = inner
2919 .reserved_module
2920 .keys()
2921 .filter(|module_key| module_key.endpoint == endpoint)
2922 .copied()
2923 .collect();
2924 for module_key in reserved_module_keys {
2925 if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2926 release_reserved_route_locked(inner, client_key, module_key);
2927 }
2928 }
2929
2930 let pending_keys: Vec<_> = inner
2931 .pending_relays
2932 .keys()
2933 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2934 .copied()
2935 .collect();
2936 let pending: Vec<_> = pending_keys
2937 .into_iter()
2938 .filter_map(|key| inner.pending_relays.remove(&key))
2939 .collect();
2940 let abandoned_relays = u32::try_from(pending.len()).unwrap_or(u32::MAX);
2941 for pending in pending {
2942 let module_label = module_id.as_deref().unwrap_or("unknown");
2943 let _ = pending
2944 .sender
2945 .send(RouteBindRelayOutcome::ModuleGone(format!(
2946 "module '{module_label}' connection closed during route.bind relay"
2947 )));
2948 }
2949
2950 let pending_control_keys: Vec<_> = inner
2951 .pending_control_rpcs
2952 .keys()
2953 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2954 .copied()
2955 .collect();
2956 let pending_control: Vec<_> = pending_control_keys
2957 .into_iter()
2958 .filter_map(|key| inner.pending_control_rpcs.remove(&key))
2959 .collect();
2960 for pending in pending_control {
2961 let module_label = module_id.as_deref().unwrap_or("unknown");
2962 let _ = pending
2963 .sender
2964 .send(ModuleControlRpcOutcome::ModuleGone(format!(
2965 "module '{module_label}' connection closed during module-control RPC"
2966 )));
2967 }
2968
2969 let module_routes = inner
2970 .module_to_client
2971 .iter()
2972 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2973 .map(|(module_key, route)| (*module_key, route.module_epoch))
2974 .collect::<Vec<_>>();
2975 let mut released = Vec::with_capacity(module_routes.len());
2976 for (module_key, epoch) in module_routes {
2977 if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
2978 {
2979 released.push(target);
2980 }
2981 }
2982 ConnectionCleanup {
2983 released,
2984 abandoned_relays,
2985 }
2986}
2987
2988#[derive(Debug, Clone, Copy)]
2989struct RequestCredit {
2990 subscription: bool,
2991 excluded_from_drain: bool,
2992}
2993
2994#[derive(Debug, Default)]
2995struct CreditLedger {
2996 by_corr: HashMap<u64, Vec<RequestCredit>>,
2997}
2998
2999impl CreditLedger {
3000 fn acquire(&mut self, corr: u64, subscription: bool) {
3001 self.by_corr.entry(corr).or_default().push(RequestCredit {
3002 subscription,
3003 excluded_from_drain: false,
3004 });
3005 }
3006
3007 fn release(&mut self, corr: u64) -> bool {
3008 let Some(credits) = self.by_corr.get_mut(&corr) else {
3009 return false;
3010 };
3011 let released = credits.pop().is_some();
3012 if credits.is_empty() {
3013 self.by_corr.remove(&corr);
3014 }
3015 released
3016 }
3017
3018 fn capture_subscription_exclusions(&mut self) -> u32 {
3019 let mut excluded = 0u32;
3020 for credit in self.by_corr.values_mut().flatten() {
3021 if credit.subscription && !credit.excluded_from_drain {
3022 credit.excluded_from_drain = true;
3023 excluded = excluded.saturating_add(1);
3024 }
3025 }
3026 excluded
3027 }
3028
3029 #[cfg(test)]
3030 fn in_flight(&self) -> usize {
3031 self.by_corr.values().map(Vec::len).sum()
3032 }
3033
3034 fn drain_in_flight(&self) -> usize {
3035 self.by_corr
3036 .values()
3037 .flatten()
3038 .filter(|credit| !credit.excluded_from_drain)
3039 .count()
3040 }
3041
3042 fn drain_held_corrs(&self) -> Vec<u64> {
3045 let mut corrs = self
3046 .by_corr
3047 .iter()
3048 .flat_map(|(corr, credits)| {
3049 credits
3050 .iter()
3051 .filter(|credit| !credit.excluded_from_drain)
3052 .map(move |_| *corr)
3053 })
3054 .collect::<Vec<_>>();
3055 corrs.sort_unstable();
3056 corrs
3057 }
3058}
3059
3060#[derive(Debug, Default)]
3061struct ChannelFlowState {
3062 closed: bool,
3063 credits: CreditLedger,
3064}
3065
3066#[derive(Debug)]
3068pub(crate) struct ChannelFlow {
3069 sem: Semaphore,
3070 window: usize,
3071 state: Mutex<ChannelFlowState>,
3072}
3073
3074impl ChannelFlow {
3075 pub(crate) fn new(window: usize) -> Self {
3076 debug_assert!(window > 0, "flow-control window must be non-zero");
3077 Self {
3078 sem: Semaphore::new(window),
3079 window,
3080 state: Mutex::new(ChannelFlowState::default()),
3081 }
3082 }
3083
3084 #[cfg(test)]
3085 pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
3086 self.acquire_tagged(0, false).await
3087 }
3088
3089 pub(crate) async fn acquire_tagged(
3090 &self,
3091 corr: u64,
3092 subscription: bool,
3093 ) -> Result<(), ChannelFlowClosed> {
3094 let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
3095 let mut state = self
3096 .state
3097 .lock()
3098 .unwrap_or_else(|poisoned| poisoned.into_inner());
3099 if state.closed {
3100 return Err(ChannelFlowClosed);
3101 }
3102 state.credits.acquire(corr, subscription);
3103 permit.forget();
3104 Ok(())
3105 }
3106
3107 #[cfg(test)]
3108 pub(crate) fn release(&self) {
3109 self.release_corr(0);
3110 }
3111
3112 pub(crate) fn release_corr(&self, corr: u64) {
3113 let released = self
3114 .state
3115 .lock()
3116 .unwrap_or_else(|poisoned| poisoned.into_inner())
3117 .credits
3118 .release(corr);
3119 if !released {
3120 warn!(
3124 window = self.window,
3125 available = self.sem.available_permits(),
3126 "flow-control over-release ignored"
3127 );
3128 return;
3129 }
3130 if !self.sem.is_closed() {
3131 self.sem.add_permits(1);
3132 }
3133 }
3134
3135 #[cfg(test)]
3136 pub(crate) fn in_flight(&self) -> usize {
3137 self.state
3138 .lock()
3139 .unwrap_or_else(|poisoned| poisoned.into_inner())
3140 .credits
3141 .in_flight()
3142 }
3143
3144 pub(crate) fn drain_in_flight(&self) -> usize {
3145 self.state
3146 .lock()
3147 .unwrap_or_else(|poisoned| poisoned.into_inner())
3148 .credits
3149 .drain_in_flight()
3150 }
3151
3152 pub(crate) fn drain_held_corrs(&self) -> Vec<u64> {
3153 self.state
3154 .lock()
3155 .unwrap_or_else(|poisoned| poisoned.into_inner())
3156 .credits
3157 .drain_held_corrs()
3158 }
3159
3160 #[cfg(test)]
3161 pub(crate) fn available_permits(&self) -> usize {
3162 self.sem.available_permits()
3163 }
3164
3165 pub(crate) fn begin_drain(&self) -> u32 {
3166 let mut state = self
3167 .state
3168 .lock()
3169 .unwrap_or_else(|poisoned| poisoned.into_inner());
3170 state.closed = true;
3171 self.sem.close();
3172 state.credits.capture_subscription_exclusions()
3173 }
3174
3175 pub(crate) fn close(&self) {
3176 self.state
3177 .lock()
3178 .unwrap_or_else(|poisoned| poisoned.into_inner())
3179 .closed = true;
3180 self.sem.close();
3181 }
3182}
3183
3184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3185pub(crate) struct ChannelFlowClosed;
3186
3187impl fmt::Display for ChannelFlowClosed {
3188 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3189 write!(f, "flow-control window closed")
3190 }
3191}
3192
3193impl Error for ChannelFlowClosed {}
3194
3195fn window_for(concurrency: &Concurrency) -> usize {
3196 match concurrency {
3197 Concurrency::Serial => 1,
3198 Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
3199 Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
3200 }
3201}
3202
3203#[derive(Debug, Clone, PartialEq, Eq)]
3204pub enum ForwardingError {
3205 NoModuleConnection,
3206 ModuleReloading {
3207 module_id: String,
3208 },
3209 StaleModuleEndpoint,
3210 UnknownReservation {
3211 client_channel: u16,
3212 module_channel: u16,
3213 },
3214 ClientRouteChannelExhausted {
3215 connection_id: ConnectionId,
3216 },
3217 ModuleRouteChannelExhausted {
3218 endpoint: ModuleEndpointId,
3219 },
3220 RelayCorrelationExhausted,
3221 ConnectionClosing {
3222 connection_id: ConnectionId,
3223 },
3224 ClientEgressClosed {
3225 connection_id: ConnectionId,
3226 },
3227 RouteOpenBuild(String),
3228 CandidateSlotOccupied {
3230 module_id: String,
3231 },
3232 ModuleEgressUnavailable {
3235 connection_id: ConnectionId,
3236 },
3237 Poisoned,
3238}
3239
3240impl fmt::Display for ForwardingError {
3241 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3242 match self {
3243 Self::NoModuleConnection => write!(f, "no module connection is registered"),
3244 Self::ModuleReloading { module_id } => {
3245 write!(f, "module_id '{module_id}' is reloading")
3246 }
3247 Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
3248 Self::UnknownReservation {
3249 client_channel,
3250 module_channel,
3251 } => write!(
3252 f,
3253 "route reservation client channel {client_channel} / module channel {module_channel} was not found"
3254 ),
3255 Self::ClientRouteChannelExhausted { connection_id } => write!(
3256 f,
3257 "no client route channels are available for connection {}",
3258 connection_id.get()
3259 ),
3260 Self::ModuleRouteChannelExhausted { endpoint } => write!(
3261 f,
3262 "no module route channels are available for endpoint generation {} on connection {}",
3263 endpoint.generation,
3264 endpoint.connection_id.get()
3265 ),
3266 Self::RelayCorrelationExhausted => {
3267 write!(f, "module control correlation ids are exhausted")
3268 }
3269 Self::ConnectionClosing { connection_id } => write!(
3270 f,
3271 "connection {} is closing and cannot accept route allocation",
3272 connection_id.get()
3273 ),
3274 Self::ClientEgressClosed { connection_id } => write!(
3275 f,
3276 "client connection {} egress is closed",
3277 connection_id.get()
3278 ),
3279 Self::RouteOpenBuild(message) => {
3280 write!(f, "failed to prebuild route.open response: {message}")
3281 }
3282 Self::CandidateSlotOccupied { module_id } => write!(
3283 f,
3284 "module_id '{module_id}' already has a swap candidate registered"
3285 ),
3286 Self::ModuleEgressUnavailable { connection_id } => write!(
3287 f,
3288 "module connection {} egress is unavailable; HELLO_ACK could not be queued",
3289 connection_id.get()
3290 ),
3291 Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
3292 }
3293 }
3294}
3295
3296impl Error for ForwardingError {}
3297
3298#[cfg(test)]
3299mod tests {
3300 use std::time::Duration;
3301
3302 use super::*;
3303 use tokio::sync::mpsc;
3304
3305 #[test]
3306 fn ordinary_long_running_request_is_not_excluded_from_drain() {
3307 let mut ledger = CreditLedger::default();
3308 ledger.acquire(1, false);
3309
3310 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3311 assert_eq!(ledger.drain_in_flight(), 1);
3312 }
3313
3314 #[test]
3315 fn bit_set_subscription_is_excluded_and_counted() {
3316 let mut ledger = CreditLedger::default();
3317 ledger.acquire(1, true);
3318
3319 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3320 assert_eq!(ledger.drain_in_flight(), 0);
3321 }
3322
3323 #[test]
3324 fn subscription_opened_after_drain_snapshot_is_not_excluded() {
3325 let mut ledger = CreditLedger::default();
3326 ledger.acquire(1, true);
3327 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3328
3329 ledger.acquire(2, true);
3330
3331 assert_eq!(ledger.drain_in_flight(), 1);
3332 }
3333
3334 #[test]
3335 fn drain_with_no_subscriptions_reports_zero_excluded() {
3336 let mut ledger = CreditLedger::default();
3337 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3338 }
3339
3340 fn test_hello_ack(corr: u64) -> Frame {
3341 Frame::build(
3342 FrameType::HelloAck,
3343 Flags::new(false, Priority::Passive, false),
3344 0,
3345 0,
3346 corr,
3347 Vec::new(),
3348 )
3349 .unwrap()
3350 }
3351
3352 #[test]
3356 fn acked_registration_that_cannot_queue_its_hello_ack_inserts_nothing() {
3357 let forwarding = ForwardingTable::default();
3358
3359 let (closed_tx, closed_rx) = mpsc::channel(8);
3360 drop(closed_rx);
3361 let closed = ConnectionId::new(1);
3362 assert_eq!(
3363 forwarding.register_module_connection_acked(
3364 closed,
3365 "closed".to_string(),
3366 2,
3367 Concurrency::ModuleManaged,
3368 FrameSink::new(closed_tx),
3369 test_hello_ack(1),
3370 ),
3371 Err(ForwardingError::ModuleEgressUnavailable {
3372 connection_id: closed
3373 })
3374 );
3375
3376 let (full_tx, _full_rx) = mpsc::channel(1);
3377 let full_sink = FrameSink::new(full_tx);
3378 full_sink.try_send(test_hello_ack(99)).unwrap();
3379 let full = ConnectionId::new(2);
3380 assert_eq!(
3381 forwarding.register_module_connection_acked(
3382 full,
3383 "full".to_string(),
3384 2,
3385 Concurrency::ModuleManaged,
3386 full_sink.clone(),
3387 test_hello_ack(2),
3388 ),
3389 Err(ForwardingError::ModuleEgressUnavailable {
3390 connection_id: full
3391 })
3392 );
3393 assert_eq!(
3394 forwarding.register_candidate_module_connection_acked(
3395 full,
3396 "full".to_string(),
3397 2,
3398 Concurrency::ModuleManaged,
3399 full_sink,
3400 test_hello_ack(3),
3401 ),
3402 Err(ForwardingError::ModuleEgressUnavailable {
3403 connection_id: full
3404 })
3405 );
3406
3407 for (connection, module_id) in [(closed, "closed"), (full, "full")] {
3408 assert_eq!(
3409 forwarding
3410 .module_endpoint_for_connection(connection)
3411 .unwrap(),
3412 None
3413 );
3414 let (client_tx, _client_rx) = mpsc::channel(8);
3415 assert_eq!(
3416 forwarding
3417 .begin_route_bind_relay_for_test(
3418 ConnectionId::new(50),
3419 FrameSink::new(client_tx),
3420 1,
3421 module_id,
3422 )
3423 .err(),
3424 Some(ForwardingError::NoModuleConnection)
3425 );
3426 }
3427 assert!(forwarding.read_inner().unwrap().candidates_by_id.is_empty());
3428 }
3429
3430 #[test]
3433 fn acked_registration_queues_the_hello_ack_first() {
3434 let forwarding = ForwardingTable::default();
3435 let (active_tx, mut active_rx) = mpsc::channel(8);
3436 forwarding
3437 .register_module_connection_acked(
3438 ConnectionId::new(1),
3439 "acked".to_string(),
3440 2,
3441 Concurrency::ModuleManaged,
3442 FrameSink::new(active_tx),
3443 test_hello_ack(11),
3444 )
3445 .unwrap();
3446 let (candidate_tx, mut candidate_rx) = mpsc::channel(8);
3447 forwarding
3448 .register_candidate_module_connection_acked(
3449 ConnectionId::new(2),
3450 "acked".to_string(),
3451 2,
3452 Concurrency::ModuleManaged,
3453 FrameSink::new(candidate_tx),
3454 test_hello_ack(12),
3455 )
3456 .unwrap();
3457
3458 let active_first = active_rx.try_recv().unwrap().frame;
3459 assert_eq!(active_first.header.ty, FrameType::HelloAck);
3460 assert_eq!(active_first.header.corr, 11);
3461 let candidate_first = candidate_rx.try_recv().unwrap().frame;
3462 assert_eq!(candidate_first.header.ty, FrameType::HelloAck);
3463 assert_eq!(candidate_first.header.corr, 12);
3464 }
3465
3466 #[test]
3467 fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
3468 let forwarding = ForwardingTable::default();
3469 let module_connection = ConnectionId::new(10);
3470 let exhausted_client = ConnectionId::new(20);
3471 let second_client = ConnectionId::new(30);
3472 let (module_tx, _module_rx) = mpsc::channel(1);
3473 let endpoint = forwarding
3474 .register_module_connection(
3475 module_connection,
3476 "route-limit-provider".to_string(),
3477 1,
3478 Concurrency::ModuleManaged,
3479 FrameSink::new(module_tx),
3480 )
3481 .unwrap();
3482
3483 {
3484 let mut inner = forwarding.inner.write().unwrap();
3485 for channel in 1..=u16::MAX {
3486 inner.reserved_client.insert(
3487 ClientRouteKey {
3488 connection_id: exhausted_client,
3489 channel,
3490 },
3491 ModuleRouteKey {
3492 endpoint,
3493 channel: 1,
3494 },
3495 );
3496 }
3497 }
3498
3499 let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
3500 let err = forwarding
3501 .begin_route_bind_relay_for_test(
3502 exhausted_client,
3503 FrameSink::new(exhausted_tx),
3504 1,
3505 "route-limit-provider",
3506 )
3507 .unwrap_err();
3508 assert!(matches!(
3509 err,
3510 ForwardingError::ClientRouteChannelExhausted { connection_id }
3511 if connection_id == exhausted_client
3512 ));
3513
3514 let (second_tx, _second_rx) = mpsc::channel(1);
3515 let pending = forwarding
3516 .begin_route_bind_relay_for_test(
3517 second_client,
3518 FrameSink::new(second_tx),
3519 2,
3520 "route-limit-provider",
3521 )
3522 .unwrap();
3523 assert_eq!(pending.client_channel, 1);
3524 }
3525
3526 #[test]
3527 fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
3528 let forwarding = ForwardingTable::default();
3529 let module_connection = ConnectionId::new(40);
3530 let client = ConnectionId::new(50);
3531 let (module_tx, _module_rx) = mpsc::channel(1);
3532 forwarding
3533 .register_module_connection(
3534 module_connection,
3535 "slot-reuse-provider".to_string(),
3536 1,
3537 Concurrency::ModuleManaged,
3538 FrameSink::new(module_tx),
3539 )
3540 .unwrap();
3541
3542 let (client_tx, _client_rx) = mpsc::channel(1);
3543 let client_sink = FrameSink::new(client_tx);
3544 let mut wrapped_channel = None;
3545 for index in 0..=usize::from(u16::MAX) {
3546 let pending = forwarding
3547 .begin_route_bind_relay_for_test(
3548 client,
3549 client_sink.clone(),
3550 index as u64 + 1,
3551 "slot-reuse-provider",
3552 )
3553 .unwrap();
3554 if index == usize::from(u16::MAX) {
3555 wrapped_channel = Some(pending.module_channel);
3556 }
3557 forwarding
3558 .abort_pending_relay(
3559 pending.endpoint,
3560 pending.corr,
3561 RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
3562 )
3563 .unwrap();
3564 }
3565
3566 assert_eq!(wrapped_channel, Some(1));
3567 }
3568
3569 #[test]
3570 fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
3571 let forwarding = ForwardingTable::default();
3572 let client = ConnectionId::new(60);
3573 forwarding
3574 .inner
3575 .write()
3576 .unwrap()
3577 .next_client_channel
3578 .insert(client, 41);
3579
3580 let released = forwarding.cleanup_connection(client).unwrap();
3581
3582 assert!(released.is_empty());
3583 assert!(!forwarding
3584 .inner
3585 .read()
3586 .unwrap()
3587 .next_client_channel
3588 .contains_key(&client));
3589 }
3590
3591 #[test]
3592 fn stale_module_cleanup_preserves_fast_reconnect_successor() {
3593 let forwarding = ForwardingTable::default();
3594 let module_id = "fast-reconnect-provider";
3595 let first_connection = ConnectionId::new(70);
3596 let second_connection = ConnectionId::new(80);
3597 let (first_tx, _first_rx) = mpsc::channel(1);
3598 let first_endpoint = forwarding
3599 .register_module_connection(
3600 first_connection,
3601 module_id.to_string(),
3602 1,
3603 Concurrency::ModuleManaged,
3604 FrameSink::new(first_tx),
3605 )
3606 .unwrap();
3607 let (second_tx, _second_rx) = mpsc::channel(1);
3608 let second_endpoint = forwarding
3609 .register_module_connection(
3610 second_connection,
3611 module_id.to_string(),
3612 1,
3613 Concurrency::ModuleManaged,
3614 FrameSink::new(second_tx),
3615 )
3616 .unwrap();
3617 assert_ne!(first_endpoint, second_endpoint);
3618
3619 let released = forwarding.cleanup_connection(first_connection).unwrap();
3620
3621 assert!(released.is_empty());
3622 assert_eq!(
3623 forwarding
3624 .inner
3625 .read()
3626 .unwrap()
3627 .modules_by_id
3628 .get(module_id)
3629 .map(|module| module.endpoint),
3630 Some(second_endpoint)
3631 );
3632 assert!(forwarding.has_live_module_connection(module_id).unwrap());
3633 let control_rpc = forwarding
3634 .begin_module_control_rpc_for(
3635 module_id,
3636 "health.check",
3637 Instant::now() + Duration::from_secs(1),
3638 )
3639 .unwrap();
3640 assert_eq!(control_rpc.endpoint, second_endpoint);
3641 }
3642
3643 fn route_fixture(
3644 module_id: &str,
3645 ) -> (
3646 ForwardingTable,
3647 ConnectionId,
3648 ModuleEndpointId,
3649 ConnectionId,
3650 FrameSink,
3651 mpsc::Receiver<crate::router::OutboundFrame>,
3652 ) {
3653 let forwarding = ForwardingTable::default();
3654 let module_connection = ConnectionId::new(100);
3655 let client_connection = ConnectionId::new(200);
3656 let (module_tx, _module_rx) = mpsc::channel(8);
3657 let endpoint = forwarding
3658 .register_module_connection(
3659 module_connection,
3660 module_id.to_string(),
3661 2,
3662 Concurrency::ModuleManaged,
3663 FrameSink::new(module_tx),
3664 )
3665 .unwrap();
3666 let (client_tx, client_rx) = mpsc::channel(8);
3667 (
3668 forwarding,
3669 module_connection,
3670 endpoint,
3671 client_connection,
3672 FrameSink::new(client_tx),
3673 client_rx,
3674 )
3675 }
3676
3677 #[test]
3678 #[cfg(unix)]
3679 fn daemon_drain_gates_current_and_racing_provider_registrations() {
3680 let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
3681 assert_eq!(forwarding.begin_daemon_drain().unwrap(), ["provider"]);
3682 assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
3683 assert!(matches!(
3684 forwarding.register_module_connection(
3685 ConnectionId::new(300),
3686 "late-provider".into(),
3687 2,
3688 Concurrency::ModuleManaged,
3689 sink,
3690 ),
3691 Err(ForwardingError::ConnectionClosing { .. })
3692 ));
3693 }
3694
3695 fn test_ping(corr: u64) -> Frame {
3696 Frame::build(
3697 FrameType::Ping,
3698 Flags::new(false, Priority::Passive, false),
3699 0,
3700 0,
3701 corr,
3702 Vec::new(),
3703 )
3704 .unwrap()
3705 }
3706
3707 fn begin_test_route(
3708 forwarding: &ForwardingTable,
3709 client_connection: ConnectionId,
3710 client_sink: FrameSink,
3711 corr: u64,
3712 module_id: &str,
3713 ) -> PendingRouteBindRelay {
3714 forwarding
3715 .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
3716 .unwrap()
3717 }
3718
3719 #[tokio::test]
3725 async fn pending_route_open_completes_behind_queued_data_frames() {
3726 assert_eq!(
3727 crate::server::MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
3728 8,
3729 "the per-connection pending route.open limit is its own constant"
3730 );
3731 let (forwarding, module_connection, _endpoint, client, _unused_sink, _unused_rx) =
3732 route_fixture("open-behind-data");
3733 let (sink, mut client_rx) = crate::server::connection_egress();
3734 const DATA_FRAMES: usize = 1_000;
3735 let data = |corr: u64| {
3736 Frame::build(
3737 FrameType::StreamData,
3738 Flags::new(false, Priority::Interactive, false),
3739 9,
3740 1,
3741 corr,
3742 vec![b'x'; 200],
3743 )
3744 .unwrap()
3745 };
3746 for corr in 0..DATA_FRAMES as u64 {
3747 sink.try_send(data(corr)).unwrap();
3748 }
3749 let data_bytes = DATA_FRAMES * (subc_protocol::HEADER_LEN + 200);
3750 assert_eq!(sink.backlog().queued_bytes, data_bytes);
3751
3752 let pending = tokio::time::timeout(
3753 Duration::from_secs(5),
3754 forwarding.begin_route_bind_relay_for(
3755 client,
3756 sink.clone(),
3757 subc_protocol::PROTOCOL_VERSION,
3758 4_242,
3759 "open-behind-data",
3760 Principal::Direct,
3761 None,
3762 Instant::now() + Duration::from_secs(60),
3763 ),
3764 )
3765 .await
3766 .expect("reserving the route.open slot must not wait behind data frames")
3767 .unwrap();
3768 forwarding
3769 .complete_pending_relay(
3770 module_connection,
3771 pending.corr,
3772 RouteBindRelayOutcome::Accepted,
3773 )
3774 .unwrap();
3775
3776 let backlog = sink.backlog();
3777 assert_eq!(backlog.queued_frames, DATA_FRAMES + 1);
3778 assert!(
3779 backlog.queued_bytes > data_bytes,
3780 "the route.open response must be counted in queued bytes: {backlog:?}"
3781 );
3782 for corr in 0..DATA_FRAMES as u64 {
3783 assert_eq!(client_rx.recv().await.unwrap().header.corr, corr);
3784 }
3785 let open = client_rx.recv().await.unwrap();
3786 assert_eq!(open.header.corr, 4_242);
3787 assert_eq!(open.header.ty, FrameType::Response);
3788 drop(open);
3789 assert_eq!(sink.backlog().queued_bytes, 0);
3790 assert_eq!(sink.backlog().queued_frames, 0);
3791 }
3792
3793 #[tokio::test]
3798 async fn drain_holdouts_count_held_requests_and_name_the_connection() {
3799 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3800 route_fixture("holdouts");
3801 let mut bound = |corr| {
3802 let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
3803 forwarding
3804 .complete_pending_relay(
3805 module_connection,
3806 route.corr,
3807 RouteBindRelayOutcome::Accepted,
3808 )
3809 .unwrap();
3810 client_rx.try_recv().unwrap();
3811 match forwarding
3812 .lookup_data_route(client, route.client_channel, route.client_epoch)
3813 .unwrap()
3814 {
3815 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3816 other => panic!("expected live route, got {other:?}"),
3817 }
3818 };
3819 let holding = bound(61);
3820 let _idle = bound(62);
3821 holding.flow.acquire_tagged(7, false).await.unwrap();
3822 holding.flow.acquire_tagged(2, false).await.unwrap();
3823 holding.flow.acquire_tagged(3, true).await.unwrap();
3824 forwarding
3825 .begin_module_drain("holdouts", RouteCloseReason::Restart)
3826 .unwrap();
3827
3828 let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
3829 assert_eq!(
3830 holdouts,
3831 DrainHoldouts {
3832 requests: 2,
3833 routes: 1,
3834 total_routes: 2,
3835 top_connections: vec![(client.get(), 2)],
3836 held: vec![(holding.module_channel, 2), (holding.module_channel, 7)],
3839 }
3840 );
3841 }
3842
3843 #[test]
3844 fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
3845 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3846 route_fixture("census");
3847 let pending = begin_test_route(&forwarding, client, sink, 1, "census");
3848 forwarding
3849 .complete_pending_relay(
3850 module_connection,
3851 pending.corr,
3852 RouteBindRelayOutcome::Accepted,
3853 )
3854 .unwrap();
3855
3856 let routes = forwarding.endpoint_routes(endpoint).unwrap();
3857 assert_eq!(routes.len(), 1);
3858 assert!(matches!(routes[0].principal, Principal::Direct));
3859 assert_eq!(routes[0].goodbye_target.connection_id, client);
3860 assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
3861 assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
3862 assert!(!routes[0].draining);
3863
3864 forwarding
3865 .begin_module_drain("census", RouteCloseReason::Restart)
3866 .unwrap();
3867 let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
3868 assert_eq!(draining_routes.len(), 1);
3869 assert!(draining_routes[0].draining);
3870 }
3871
3872 #[test]
3873 fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
3874 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
3875 let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
3876 assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
3877 forwarding
3878 .abort_pending_relay(
3879 first.endpoint,
3880 first.corr,
3881 RouteBindRelayOutcome::ModuleGone("abort".into()),
3882 )
3883 .unwrap();
3884 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3885 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3886
3887 let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
3888 assert_eq!(second.client_channel, first.client_channel);
3889 assert_eq!(second.module_channel, first.module_channel);
3890 assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
3891 }
3892
3893 #[test]
3894 fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
3895 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3896 route_fixture("epoch-release");
3897 let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
3898 forwarding
3899 .complete_pending_relay(
3900 module_connection,
3901 first.corr,
3902 RouteBindRelayOutcome::Accepted,
3903 )
3904 .unwrap();
3905 assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
3906 assert!(matches!(
3907 forwarding
3908 .release_client_route(client, first.client_channel, first.client_epoch)
3909 .unwrap(),
3910 RouteRelease::Removed(_)
3911 ));
3912 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3913 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3914
3915 let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
3916 forwarding
3917 .complete_pending_relay(
3918 module_connection,
3919 second.corr,
3920 RouteBindRelayOutcome::Accepted,
3921 )
3922 .unwrap();
3923 assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
3924 assert!(matches!(
3925 forwarding
3926 .release_client_route(client, second.client_channel, first.client_epoch)
3927 .unwrap(),
3928 RouteRelease::Stale
3929 ));
3930 assert!(!forwarding
3931 .cache_status(
3932 endpoint,
3933 second.module_channel,
3934 first.module_epoch,
3935 "stale".into(),
3936 )
3937 .unwrap());
3938 assert!(forwarding
3939 .cache_status(
3940 endpoint,
3941 second.module_channel,
3942 second.module_epoch,
3943 "current".into(),
3944 )
3945 .unwrap());
3946 match forwarding
3947 .route_poll_snapshot(client, second.client_channel, second.client_epoch)
3948 .unwrap()
3949 {
3950 RoutePollSnapshot::Bound { status, .. } => {
3951 assert_eq!(status.as_deref(), Some("current"));
3952 }
3953 RoutePollSnapshot::Absent => panic!("successor binding was removed"),
3954 }
3955 let counters = forwarding.counters().snapshot();
3956 assert_eq!(counters["route_released_epoch_fenced"], 1);
3957 assert_eq!(counters["route_release_stale_skipped"], 1);
3958 }
3959
3960 #[test]
3961 fn max_epoch_reservation_retires_only_that_slot() {
3962 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
3963 forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
3964 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
3965 let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
3966 assert_eq!(
3967 (final_use.client_channel, final_use.client_epoch),
3968 (7, u32::MAX)
3969 );
3970 assert_eq!(
3971 (final_use.module_channel, final_use.module_epoch),
3972 (9, u32::MAX)
3973 );
3974 forwarding
3975 .abort_pending_relay(
3976 endpoint,
3977 final_use.corr,
3978 RouteBindRelayOutcome::ModuleGone("abort".into()),
3979 )
3980 .unwrap();
3981 forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
3982 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
3983 let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
3984 assert_ne!(next.client_channel, 7);
3985 assert_ne!(next.module_channel, 9);
3986 assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
3987 }
3988
3989 #[test]
3990 fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
3991 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3992 route_fixture("corr-shared");
3993 let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
3994 assert_eq!(bind.corr, 1);
3995 forwarding
3996 .abort_pending_relay(
3997 endpoint,
3998 bind.corr,
3999 RouteBindRelayOutcome::ModuleGone("abort".into()),
4000 )
4001 .unwrap();
4002 let rpc = forwarding
4003 .begin_module_control_rpc_for(
4004 "corr-shared",
4005 "health.check",
4006 Instant::now() - Duration::from_millis(1),
4007 )
4008 .unwrap();
4009 assert_eq!(rpc.corr, 2);
4010 assert_eq!(
4011 forwarding
4012 .complete_module_control_rpc(
4013 module_connection,
4014 rpc.corr,
4015 Some("health.check"),
4016 ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
4017 status: subc_protocol::session::HealthStatus::Ok,
4018 detail: None,
4019 metrics: None,
4020 }),
4021 )
4022 .unwrap(),
4023 ModuleControlRpcCompletion::Settled
4024 );
4025 assert!(matches!(
4026 rpc.receiver.blocking_recv().unwrap(),
4027 ModuleControlRpcOutcome::DeadlineElapsed
4028 ));
4029 }
4030
4031 #[tokio::test(start_paused = true)]
4032 async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
4033 let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
4034 let probe_started_at = Instant::now();
4035 let rpc = forwarding
4036 .begin_health_probe_rpc_for(
4037 "tombstone-ttl",
4038 "health.check",
4039 probe_started_at,
4040 probe_started_at + Duration::from_secs(5),
4041 )
4042 .unwrap();
4043 assert!(forwarding
4044 .tombstone_health_probe_rpc(endpoint, rpc.corr)
4045 .unwrap());
4046 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
4047
4048 tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
4049 tokio::task::yield_now().await;
4050
4051 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
4052 }
4053
4054 #[test]
4055 fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
4056 let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
4057 let mut close = forwarding.register_connection_close(endpoint.connection_id);
4058 forwarding.inject_control_corr(endpoint, u64::MAX);
4059 let final_rpc = forwarding
4060 .begin_module_control_rpc_for(
4061 "corr-max",
4062 "health.check",
4063 Instant::now() + Duration::from_secs(1),
4064 )
4065 .unwrap();
4066 assert_eq!(final_rpc.corr, u64::MAX);
4067 forwarding
4068 .cancel_module_control_rpc(endpoint, final_rpc.corr)
4069 .unwrap();
4070 assert!(matches!(
4071 forwarding.begin_module_control_rpc_for(
4072 "corr-max",
4073 "health.check",
4074 Instant::now() + Duration::from_secs(1),
4075 ),
4076 Err(ForwardingError::RelayCorrelationExhausted)
4077 ));
4078 assert!(close.try_recv().is_ok());
4079 }
4080
4081 #[test]
4082 fn publication_epoch_controls_delivery_failure_escalation() {
4083 fn setup_successor(
4084 commit_successor: Option<bool>,
4085 ) -> (ForwardingTable, ConnectionId, u16, u32) {
4086 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4087 route_fixture("escalation");
4088 let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
4089 forwarding
4090 .complete_pending_relay(
4091 module_connection,
4092 first.corr,
4093 RouteBindRelayOutcome::Accepted,
4094 )
4095 .unwrap();
4096 client_rx.try_recv().unwrap();
4097 assert!(matches!(
4098 forwarding
4099 .release_client_route(client, first.client_channel, first.client_epoch)
4100 .unwrap(),
4101 RouteRelease::Removed(_)
4102 ));
4103 if let Some(commit_successor) = commit_successor {
4104 forwarding.inject_client_slot_epoch(
4105 client,
4106 first.client_channel,
4107 first.client_epoch,
4108 );
4109 forwarding.inject_module_slot_epoch(
4110 endpoint,
4111 first.module_channel,
4112 first.module_epoch,
4113 );
4114 let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
4115 if commit_successor {
4116 forwarding
4117 .complete_pending_relay(
4118 module_connection,
4119 successor.corr,
4120 RouteBindRelayOutcome::Accepted,
4121 )
4122 .unwrap();
4123 client_rx.try_recv().unwrap();
4124 } else {
4125 forwarding
4126 .abort_pending_relay(
4127 endpoint,
4128 successor.corr,
4129 RouteBindRelayOutcome::ModuleGone("abort".into()),
4130 )
4131 .unwrap();
4132 }
4133 }
4134 (forwarding, client, first.client_channel, first.client_epoch)
4135 }
4136
4137 let probe_sink = FrameSink::new(mpsc::channel(1).0);
4138 let (no_successor, client, channel, epoch) = setup_successor(None);
4139 let mut close = no_successor.register_connection_close(client);
4140 assert!(no_successor
4141 .escalate_client_delivery_failure(
4142 client,
4143 channel,
4144 epoch,
4145 CloseReason::new("delivery", "failed"),
4146 UndeliveredFrame {
4147 module_id: None,
4148 sink: &probe_sink,
4149 },
4150 )
4151 .unwrap());
4152 assert!(close.try_recv().is_ok());
4153
4154 let (aborted, client, channel, epoch) = setup_successor(Some(false));
4155 let mut close = aborted.register_connection_close(client);
4156 assert!(aborted
4157 .escalate_client_delivery_failure(
4158 client,
4159 channel,
4160 epoch,
4161 CloseReason::new("delivery", "failed"),
4162 UndeliveredFrame {
4163 module_id: None,
4164 sink: &probe_sink,
4165 },
4166 )
4167 .unwrap());
4168 assert!(close.try_recv().is_ok());
4169
4170 let (published, client, channel, epoch) = setup_successor(Some(true));
4171 let mut close = published.register_connection_close(client);
4172 assert!(!published
4173 .escalate_client_delivery_failure(
4174 client,
4175 channel,
4176 epoch,
4177 CloseReason::new("delivery", "stale failure"),
4178 UndeliveredFrame {
4179 module_id: None,
4180 sink: &probe_sink,
4181 },
4182 )
4183 .unwrap());
4184 assert!(close.try_recv().is_err());
4185 }
4186
4187 #[test]
4188 fn route_concentration_separates_client_count_from_routes_per_client() {
4189 let (forwarding, module_connection, _, client, sink, _client_rx) =
4193 route_fixture("concentration");
4194 assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
4195
4196 for corr in [70_u64, 71] {
4197 let pending =
4198 begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
4199 forwarding
4200 .complete_pending_relay(
4201 module_connection,
4202 pending.corr,
4203 RouteBindRelayOutcome::Accepted,
4204 )
4205 .unwrap();
4206 }
4207
4208 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4210 assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
4211 }
4212
4213 #[test]
4214 fn cleanup_and_accepted_resolution_have_one_lock_winner() {
4215 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4216 route_fixture("cleanup-race");
4217 let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
4218 forwarding
4219 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4220 .unwrap();
4221 let released = forwarding.cleanup_connection(client).unwrap();
4222 assert_eq!(released.len(), 1);
4223 let completion = forwarding
4224 .complete_pending_relay(
4225 module_connection,
4226 pending.corr,
4227 RouteBindRelayOutcome::Accepted,
4228 )
4229 .unwrap();
4230 assert!(!completion.settled);
4231 assert!(client_rx.try_recv().is_err());
4232 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4233
4234 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4235 route_fixture("accepted-race");
4236 let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
4237 forwarding
4238 .complete_pending_relay(
4239 module_connection,
4240 pending.corr,
4241 RouteBindRelayOutcome::Accepted,
4242 )
4243 .unwrap();
4244 assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
4245 let released = forwarding.cleanup_connection(client).unwrap();
4246 assert_eq!(released.len(), 1);
4247 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4248 }
4249
4250 #[test]
4251 fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
4252 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4253 route_fixture("drain-gap");
4254 let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
4255 forwarding
4256 .complete_pending_relay(
4257 module_connection,
4258 live.corr,
4259 RouteBindRelayOutcome::Accepted,
4260 )
4261 .unwrap();
4262 client_rx.try_recv().unwrap();
4263 let binding = match forwarding
4264 .lookup_data_route(client, live.client_channel, live.client_epoch)
4265 .unwrap()
4266 {
4267 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4268 other => panic!("expected live route, got {other:?}"),
4269 };
4270
4271 let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
4272 forwarding
4273 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4274 .unwrap();
4275 let control_rpc = forwarding
4276 .begin_module_control_rpc_for(
4277 "drain-gap",
4278 "health.check",
4279 Instant::now() + Duration::from_secs(1),
4280 )
4281 .unwrap();
4282 let target = forwarding
4283 .begin_module_drain("drain-gap", RouteCloseReason::Reload)
4284 .unwrap()
4285 .unwrap();
4286 assert!(matches!(
4287 control_rpc.receiver.blocking_recv().unwrap(),
4288 ModuleControlRpcOutcome::ModuleGone(_)
4289 ));
4290 assert_eq!(target.abandoned_bindings.len(), 1);
4291 assert!(binding.flow.sem.is_closed());
4292 assert!(
4293 !forwarding
4294 .complete_pending_relay(
4295 module_connection,
4296 pending.corr,
4297 RouteBindRelayOutcome::Accepted,
4298 )
4299 .unwrap()
4300 .settled
4301 );
4302 assert!(matches!(
4303 forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
4304 Err(ForwardingError::ModuleReloading { .. })
4305 ));
4306 let released = forwarding
4307 .release_module_endpoint_routes(target.endpoint)
4308 .unwrap();
4309 assert_eq!(released.len(), 1);
4310 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4311 }
4312
4313 #[test]
4321 fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
4322 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4323 route_fixture("closing-client");
4324
4325 let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
4328 forwarding
4329 .complete_pending_relay(
4330 module_connection,
4331 live.corr,
4332 RouteBindRelayOutcome::Accepted,
4333 )
4334 .unwrap();
4335 client_rx.try_recv().unwrap();
4336
4337 let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
4339 forwarding
4340 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4341 .unwrap();
4342
4343 assert!(forwarding
4345 .escalate_client_delivery_failure(
4346 client,
4347 live.client_channel,
4348 live.client_epoch,
4349 CloseReason::new(
4350 "module_to_client_delivery_failed",
4351 "client egress refused a module frame",
4352 ),
4353 UndeliveredFrame {
4354 module_id: None,
4355 sink: &sink,
4356 },
4357 )
4358 .unwrap());
4359 assert!(!sink.is_closed());
4360
4361 let completion = forwarding
4362 .complete_pending_relay(
4363 module_connection,
4364 pending.corr,
4365 RouteBindRelayOutcome::Accepted,
4366 )
4367 .expect("a closing client must not turn a module's ack into an error");
4368
4369 assert!(completion.settled);
4370 let abandoned = completion
4371 .abandoned
4372 .expect("the module must be told to drop the binding it just created");
4373 assert_eq!(abandoned.connection_id, module_connection);
4374 assert_eq!(abandoned.channel, pending.module_channel);
4375 assert_eq!(abandoned.epoch, pending.module_epoch);
4376 assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
4377 assert!(matches!(
4378 pending.receiver.blocking_recv().unwrap(),
4379 RouteBindRelayOutcome::ModuleGone(_)
4380 ));
4381 assert!(client_rx.try_recv().is_err());
4384 assert_eq!(forwarding.active_binding_count().unwrap(), 1);
4385
4386 assert!(forwarding
4389 .has_live_module_connection("closing-client")
4390 .unwrap());
4391 let cotenant = ConnectionId::new(201);
4392 let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
4393 let cotenant_route = begin_test_route(
4394 &forwarding,
4395 cotenant,
4396 FrameSink::new(cotenant_tx),
4397 62,
4398 "closing-client",
4399 );
4400 assert_eq!(cotenant_route.endpoint, endpoint);
4401 forwarding
4402 .complete_pending_relay(
4403 module_connection,
4404 cotenant_route.corr,
4405 RouteBindRelayOutcome::Accepted,
4406 )
4407 .unwrap();
4408 assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
4409 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4410 }
4411
4412 #[test]
4413 fn pending_route_permit_is_released_on_rejection_and_abort() {
4414 let forwarding = ForwardingTable::default();
4415 let module_connection = ConnectionId::new(300);
4416 let client = ConnectionId::new(301);
4417 let (module_tx, _module_rx) = mpsc::channel(1);
4418 let endpoint = forwarding
4419 .register_module_connection(
4420 module_connection,
4421 "permit".into(),
4422 2,
4423 Concurrency::ModuleManaged,
4424 FrameSink::new(module_tx),
4425 )
4426 .unwrap();
4427 let (client_tx, mut client_rx) = mpsc::channel(1);
4428 let sink = FrameSink::new(client_tx);
4429 let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
4430 assert!(sink.try_send(test_ping(999)).is_err());
4431 forwarding
4432 .complete_pending_relay(
4433 module_connection,
4434 rejected.corr,
4435 RouteBindRelayOutcome::Rejected(ErrorBody {
4436 code: "no".into(),
4437 message: "rejected".into(),
4438 detail: None,
4439 }),
4440 )
4441 .unwrap();
4442 sink.try_send(test_ping(1000)).unwrap();
4443 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
4444
4445 let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
4446 assert!(sink.try_send(test_ping(1001)).is_err());
4447 forwarding
4448 .abort_pending_relay(
4449 endpoint,
4450 aborted.corr,
4451 RouteBindRelayOutcome::ModuleGone("abort".into()),
4452 )
4453 .unwrap();
4454 sink.try_send(test_ping(1002)).unwrap();
4455 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
4456
4457 let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
4458 forwarding
4459 .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
4460 .unwrap();
4461 drop(client_rx);
4462 let completion = forwarding
4463 .complete_pending_relay(
4464 module_connection,
4465 receiver_closed.corr,
4466 RouteBindRelayOutcome::Accepted,
4467 )
4468 .unwrap();
4469 assert!(completion.abandoned.is_some());
4470 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4471 }
4472
4473 #[test]
4479 fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
4480 let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
4481 route_fixture("closing-set-leak");
4482
4483 const CONNECTIONS: u64 = 32;
4484 for index in 0..CONNECTIONS {
4485 let client = ConnectionId::new(1000 + index);
4486 let (client_tx, _client_rx) = mpsc::channel(8);
4487 let route = begin_test_route(
4488 &forwarding,
4489 client,
4490 FrameSink::new(client_tx),
4491 index + 1,
4492 "closing-set-leak",
4493 );
4494 forwarding
4495 .complete_pending_relay(
4496 module_connection,
4497 route.corr,
4498 RouteBindRelayOutcome::Accepted,
4499 )
4500 .unwrap();
4501 forwarding.cleanup_connection(client).unwrap();
4502 }
4503 forwarding.cleanup_connection(module_connection).unwrap();
4504
4505 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
4506 }
4507
4508 #[test]
4515 fn closing_connection_is_refused_new_work_until_cleanup_completes() {
4516 let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
4517 route_fixture("closing-gate");
4518
4519 let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
4522 forwarding
4523 .complete_pending_relay(
4524 module_connection,
4525 live.corr,
4526 RouteBindRelayOutcome::Accepted,
4527 )
4528 .unwrap();
4529 client_rx.try_recv().unwrap();
4530
4531 assert!(forwarding
4534 .escalate_client_delivery_failure(
4535 client,
4536 live.client_channel,
4537 live.client_epoch,
4538 CloseReason::new(
4539 "module_to_client_delivery_failed",
4540 "client egress refused a module frame",
4541 ),
4542 UndeliveredFrame {
4543 module_id: None,
4544 sink: &sink,
4545 },
4546 )
4547 .unwrap());
4548 assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
4549
4550 assert!(matches!(
4552 forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
4553 Err(ForwardingError::ConnectionClosing { connection_id })
4554 if connection_id == client
4555 ));
4556 let (late_tx, _late_rx) = mpsc::channel(1);
4558 assert!(matches!(
4559 forwarding.register_module_connection(
4560 client,
4561 "late-module".into(),
4562 2,
4563 Concurrency::ModuleManaged,
4564 FrameSink::new(late_tx),
4565 ),
4566 Err(ForwardingError::ConnectionClosing { connection_id })
4567 if connection_id == client
4568 ));
4569
4570 forwarding.cleanup_connection(client).unwrap();
4574 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
4575 }
4576}
4577
4578#[cfg(test)]
4581mod swap_slot_tests {
4582 use std::time::Duration;
4583
4584 use super::*;
4585 use tokio::sync::mpsc;
4586
4587 const MODULE_ID: &str = "swapped";
4588
4589 struct SwapFixture {
4590 forwarding: ForwardingTable,
4591 incumbent_connection: ConnectionId,
4592 incumbent: ModuleEndpointId,
4593 candidate_connection: ConnectionId,
4594 candidate: ModuleEndpointId,
4595 _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
4596 }
4597
4598 fn swap_fixture() -> SwapFixture {
4599 let forwarding = ForwardingTable::default();
4600 let incumbent_connection = ConnectionId::new(100);
4601 let candidate_connection = ConnectionId::new(110);
4602 let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
4603 let incumbent = forwarding
4604 .register_module_connection(
4605 incumbent_connection,
4606 MODULE_ID.to_string(),
4607 2,
4608 Concurrency::ModuleManaged,
4609 FrameSink::new(incumbent_tx),
4610 )
4611 .unwrap();
4612 let (candidate_tx, candidate_rx) = mpsc::channel(8);
4613 let candidate = forwarding
4614 .register_candidate_module_connection(
4615 candidate_connection,
4616 MODULE_ID.to_string(),
4617 2,
4618 Concurrency::ModuleManaged,
4619 FrameSink::new(candidate_tx),
4620 )
4621 .unwrap();
4622 SwapFixture {
4623 forwarding,
4624 incumbent_connection,
4625 incumbent,
4626 candidate_connection,
4627 candidate,
4628 _module_rxs: vec![incumbent_rx, candidate_rx],
4629 }
4630 }
4631
4632 fn client(
4633 raw: u64,
4634 ) -> (
4635 ConnectionId,
4636 FrameSink,
4637 mpsc::Receiver<crate::router::OutboundFrame>,
4638 ) {
4639 let (tx, rx) = mpsc::channel(8);
4640 (ConnectionId::new(raw), FrameSink::new(tx), rx)
4641 }
4642
4643 fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
4644 forwarding
4645 .read_inner()
4646 .unwrap()
4647 .client_to_module
4648 .values()
4649 .map(|route| route.module_endpoint)
4650 .collect()
4651 }
4652
4653 #[test]
4654 fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
4655 let fixture = swap_fixture();
4656 let forwarding = &fixture.forwarding;
4657 assert_ne!(fixture.incumbent, fixture.candidate);
4658
4659 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4661 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4662 let (client_connection, client_sink, _client_rx) = client(200);
4663 let pending = forwarding
4664 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4665 .unwrap();
4666 assert_eq!(pending.endpoint, fixture.incumbent);
4667 let rpc = forwarding
4668 .begin_module_control_rpc_for(
4669 MODULE_ID,
4670 "health.check",
4671 Instant::now() + Duration::from_secs(1),
4672 )
4673 .unwrap();
4674 assert_eq!(rpc.endpoint, fixture.incumbent);
4675 let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
4676 assert_eq!(census.len(), 1, "the census lists one endpoint per id");
4677
4678 assert_eq!(
4680 forwarding
4681 .module_endpoint_for_connection(fixture.candidate_connection)
4682 .unwrap(),
4683 Some(fixture.candidate)
4684 );
4685 assert_eq!(
4686 forwarding
4687 .module_id_for_connection(fixture.candidate_connection)
4688 .unwrap()
4689 .as_deref(),
4690 Some(MODULE_ID)
4691 );
4692
4693 let (other_tx, _other_rx) = mpsc::channel(1);
4695 assert_eq!(
4696 forwarding.register_candidate_module_connection(
4697 ConnectionId::new(120),
4698 MODULE_ID.to_string(),
4699 2,
4700 Concurrency::ModuleManaged,
4701 FrameSink::new(other_tx),
4702 ),
4703 Err(ForwardingError::CandidateSlotOccupied {
4704 module_id: MODULE_ID.to_string()
4705 })
4706 );
4707 }
4708
4709 #[test]
4713 fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
4714 let fixture = swap_fixture();
4715 let forwarding = &fixture.forwarding;
4716 let (early_client, early_sink, _early_rx) = client(200);
4717 let mut early = forwarding
4718 .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
4719 .unwrap();
4720 assert_eq!(early.endpoint, fixture.incumbent);
4721 assert!(forwarding
4722 .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
4723 .unwrap());
4724
4725 let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4726 assert_eq!(
4727 cutover,
4728 ForwardingCutover {
4729 promoted: fixture.candidate,
4730 incumbent: Some(fixture.incumbent),
4731 }
4732 );
4733
4734 let (late_client, late_sink, _late_rx) = client(201);
4736 let late = forwarding
4737 .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
4738 .unwrap();
4739 assert_eq!(
4740 late.endpoint, fixture.candidate,
4741 "a route.open after cutover was reserved on the incumbent"
4742 );
4743
4744 let completion = forwarding
4746 .complete_pending_relay(
4747 fixture.incumbent_connection,
4748 early.corr,
4749 RouteBindRelayOutcome::Accepted,
4750 )
4751 .expect("a superseded endpoint's ack is not an error on its connection");
4752 assert!(completion.settled);
4753 assert!(
4754 !committed_endpoints(forwarding).contains(&fixture.incumbent),
4755 "a relay reserved before cutover committed a route on the incumbent"
4756 );
4757 let goodbye = completion
4758 .abandoned
4759 .expect("the incumbent is told to drop the binding it just created");
4760 assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
4761 assert_eq!(goodbye.channel, early.module_channel);
4762 assert_eq!(goodbye.epoch, early.module_epoch);
4763 assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
4764 match early.receiver.try_recv() {
4765 Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
4766 other => panic!("expected a retryable module_reloading answer, got {other:?}"),
4767 }
4768 assert!(matches!(
4769 forwarding
4770 .lookup_data_route(early_client, early.client_channel, early.client_epoch)
4771 .unwrap(),
4772 DataRoute::Client(DataRouteState::Absent)
4773 ));
4774
4775 assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4778 forwarding
4779 .complete_pending_relay(
4780 fixture.candidate_connection,
4781 late.corr,
4782 RouteBindRelayOutcome::Accepted,
4783 )
4784 .unwrap();
4785 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4786 assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
4787 }
4788
4789 #[test]
4790 fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
4791 let fixture = swap_fixture();
4792 let forwarding = &fixture.forwarding;
4793 let (bound_client, bound_sink, _bound_rx) = client(200);
4795 let bound = forwarding
4796 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4797 .unwrap();
4798 forwarding
4799 .complete_pending_relay(
4800 fixture.incumbent_connection,
4801 bound.corr,
4802 RouteBindRelayOutcome::Accepted,
4803 )
4804 .unwrap();
4805 let (pending_client, pending_sink, _pending_rx) = client(201);
4806 let mut in_flight = forwarding
4807 .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
4808 .unwrap();
4809 forwarding
4810 .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
4811 .unwrap();
4812
4813 let incumbent = forwarding
4814 .cutover_candidate(MODULE_ID)
4815 .unwrap()
4816 .unwrap()
4817 .incumbent
4818 .unwrap();
4819 let target = forwarding
4820 .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
4821 .unwrap()
4822 .expect("the superseded incumbent is still registered");
4823
4824 assert_eq!(target.endpoint, fixture.incumbent);
4825 assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
4826 assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
4827 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4828 assert_eq!(target.abandoned_bindings.len(), 1);
4829 assert_eq!(
4830 target.abandoned_bindings[0].channel,
4831 in_flight.module_channel
4832 );
4833 assert!(matches!(
4834 in_flight.receiver.try_recv(),
4835 Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
4836 ));
4837 assert_eq!(
4838 forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
4839 1,
4840 "the incumbent's bound route stays until its drain finishes"
4841 );
4842
4843 let (next_client, next_sink, _next_rx) = client(202);
4844 let next = forwarding
4845 .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
4846 .expect("the promoted candidate keeps accepting routes");
4847 assert_eq!(next.endpoint, fixture.candidate);
4848 }
4849
4850 #[test]
4856 fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
4857 let forwarding = ForwardingTable::default();
4858 let first_connection = ConnectionId::new(70);
4859 let (first_tx, _first_rx) = mpsc::channel(8);
4860 forwarding
4861 .register_module_connection(
4862 first_connection,
4863 MODULE_ID.to_string(),
4864 2,
4865 Concurrency::ModuleManaged,
4866 FrameSink::new(first_tx),
4867 )
4868 .unwrap();
4869 let (client_connection, client_sink, _client_rx) = client(200);
4870 let mut pending = forwarding
4871 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4872 .unwrap();
4873 let (second_tx, _second_rx) = mpsc::channel(8);
4874 forwarding
4875 .register_module_connection(
4876 ConnectionId::new(80),
4877 MODULE_ID.to_string(),
4878 2,
4879 Concurrency::ModuleManaged,
4880 FrameSink::new(second_tx),
4881 )
4882 .unwrap();
4883
4884 assert_eq!(
4885 forwarding
4886 .complete_pending_relay(
4887 first_connection,
4888 pending.corr,
4889 RouteBindRelayOutcome::Accepted
4890 )
4891 .unwrap_err(),
4892 ForwardingError::StaleModuleEndpoint
4893 );
4894 assert!(committed_endpoints(&forwarding).is_empty());
4895 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4896 assert!(matches!(
4897 pending.receiver.try_recv(),
4898 Err(oneshot::error::TryRecvError::Closed)
4899 ));
4900 }
4901
4902 #[test]
4903 fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
4904 let fixture = swap_fixture();
4906 let forwarding = &fixture.forwarding;
4907 assert!(forwarding
4908 .cleanup_connection(fixture.candidate_connection)
4909 .unwrap()
4910 .is_empty());
4911 assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
4912 let (client_connection, client_sink, _client_rx) = client(200);
4913 assert_eq!(
4914 forwarding
4915 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4916 .unwrap()
4917 .endpoint,
4918 fixture.incumbent
4919 );
4920
4921 let fixture = swap_fixture();
4924 let forwarding = &fixture.forwarding;
4925 let (bound_client, bound_sink, _bound_rx) = client(200);
4926 let bound = forwarding
4927 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4928 .unwrap();
4929 forwarding
4930 .complete_pending_relay(
4931 fixture.incumbent_connection,
4932 bound.corr,
4933 RouteBindRelayOutcome::Accepted,
4934 )
4935 .unwrap();
4936 forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4937 let released = forwarding
4938 .cleanup_connection(fixture.incumbent_connection)
4939 .unwrap();
4940 assert_eq!(released.len(), 1);
4941 assert_eq!(released[0].connection_id, bound_client);
4942 assert!(forwarding
4943 .read_inner()
4944 .unwrap()
4945 .superseded_endpoints
4946 .is_empty());
4947 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4948 let (next_client, next_sink, _next_rx) = client(201);
4949 assert_eq!(
4950 forwarding
4951 .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
4952 .unwrap()
4953 .endpoint,
4954 fixture.candidate
4955 );
4956 }
4957}