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