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 #[cfg(unix)]
1875 pub(crate) fn begin_daemon_drain(&self) -> Result<Vec<String>, ForwardingError> {
1876 let mut inner = self.write_inner()?;
1877 inner.daemon_draining = true;
1878 let modules = inner
1879 .modules_by_id
1880 .iter()
1881 .map(|(id, module)| (id.clone(), module.endpoint))
1882 .collect::<Vec<_>>();
1883 for (_, endpoint) in &modules {
1884 inner
1885 .draining_endpoints
1886 .insert(*endpoint, RouteCloseReason::Restart);
1887 }
1888 let off_slot_endpoints = inner
1893 .candidates_by_id
1894 .values()
1895 .map(|module| module.endpoint)
1896 .chain(inner.superseded_endpoints.keys().copied())
1897 .collect::<Vec<_>>();
1898 for endpoint in off_slot_endpoints {
1899 inner
1900 .draining_endpoints
1901 .insert(endpoint, RouteCloseReason::Restart);
1902 }
1903 Ok(modules.into_iter().map(|(id, _)| id).collect())
1904 }
1905
1906 pub(crate) fn begin_module_drain(
1913 &self,
1914 module_id: &str,
1915 reason: RouteCloseReason,
1916 ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
1917 let mut inner = self.write_inner()?;
1918 let Some(module) = inner.modules_by_id.get(module_id).cloned() else {
1919 return Ok(None);
1920 };
1921 Ok(Some(begin_drain_locked(
1922 &mut inner, module_id, module, reason,
1923 )))
1924 }
1925
1926 pub(crate) fn begin_endpoint_drain(
1933 &self,
1934 endpoint: ModuleEndpointId,
1935 reason: RouteCloseReason,
1936 ) -> Result<Option<ModuleDrainTarget>, ForwardingError> {
1937 let mut inner = self.write_inner()?;
1938 let Some(module) = module_connection_for_endpoint_locked(&inner, endpoint).cloned() else {
1939 return Ok(None);
1940 };
1941 let module_id = inner
1942 .module_id_by_endpoint
1943 .get(&endpoint)
1944 .cloned()
1945 .expect("an endpoint resolved to a module connection has a module id");
1946 Ok(Some(begin_drain_locked(
1947 &mut inner, &module_id, module, reason,
1948 )))
1949 }
1950}
1951
1952fn begin_drain_locked(
1956 inner: &mut ForwardingInner,
1957 module_id: &str,
1958 module: ModuleConnection,
1959 reason: RouteCloseReason,
1960) -> ModuleDrainTarget {
1961 {
1962 let endpoint = module.endpoint;
1963 inner.draining_endpoints.insert(endpoint, reason);
1964
1965 let flows = inner
1966 .client_to_module
1967 .values()
1968 .filter(|route| route.module_endpoint == endpoint)
1969 .map(|route| Arc::clone(&route.flow))
1970 .collect::<Vec<_>>();
1971 let excluded_subscriptions = flows
1972 .into_iter()
1973 .map(|flow| flow.begin_drain())
1974 .fold(0u32, u32::saturating_add);
1975
1976 let pending_keys = inner
1977 .pending_relays
1978 .keys()
1979 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
1980 .copied()
1981 .collect::<Vec<_>>();
1982 let mut abandoned_bindings = Vec::new();
1983 for key in pending_keys {
1984 let Some(pending) = inner.pending_relays.remove(&key) else {
1985 continue;
1986 };
1987 release_reserved_route_locked(
1988 inner,
1989 pending.reservation.client_key,
1990 pending.reservation.module_key,
1991 );
1992 if pending.relay_enqueued {
1993 if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
1994 abandoned_bindings.push(target);
1995 }
1996 }
1997 let _ = pending
1998 .sender
1999 .send(RouteBindRelayOutcome::Rejected(ErrorBody::new(
2000 "module_reloading",
2001 format!("module_id '{module_id}' is reloading"),
2002 )));
2003 }
2004
2005 let pending_control_keys = inner
2006 .pending_control_rpcs
2007 .keys()
2008 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2009 .copied()
2010 .collect::<Vec<_>>();
2011 for key in pending_control_keys {
2012 if let Some(pending) = inner.pending_control_rpcs.remove(&key) {
2013 let _ = pending
2014 .sender
2015 .send(ModuleControlRpcOutcome::ModuleGone(format!(
2016 "module '{module_id}' began draining during module-control RPC"
2017 )));
2018 }
2019 }
2020
2021 ModuleDrainTarget {
2022 endpoint,
2023 sink: module.sink,
2024 negotiated_ver: module.negotiated_ver,
2025 abandoned_bindings,
2026 excluded_subscriptions,
2027 }
2028 }
2029}
2030
2031#[derive(Debug, Default, PartialEq, Eq)]
2034pub(crate) struct DrainHoldouts {
2035 pub(crate) requests: usize,
2038 pub(crate) routes: usize,
2040 pub(crate) total_routes: usize,
2042 pub(crate) top_connections: Vec<(u64, usize)>,
2045 pub(crate) held: Vec<(u16, u64)>,
2054}
2055
2056pub(crate) const DRAIN_HELD_REQUESTS_LISTED: usize = 32;
2058
2059impl ForwardingTable {
2060 pub(crate) fn endpoint_drain_holdouts(
2062 &self,
2063 endpoint: ModuleEndpointId,
2064 ) -> Result<DrainHoldouts, ForwardingError> {
2065 let inner = self.read_inner()?;
2066 let mut holdouts = DrainHoldouts::default();
2067 let mut by_connection: HashMap<u64, usize> = HashMap::new();
2068 for (key, route) in &inner.client_to_module {
2069 if route.module_endpoint != endpoint {
2070 continue;
2071 }
2072 holdouts.total_routes += 1;
2073 let held = route.flow.drain_in_flight();
2074 if held == 0 {
2075 continue;
2076 }
2077 holdouts.requests += held;
2078 holdouts.routes += 1;
2079 *by_connection.entry(key.connection_id.get()).or_default() += held;
2080 holdouts.held.extend(
2081 route
2082 .flow
2083 .drain_held_corrs()
2084 .into_iter()
2085 .map(|corr| (route.module_channel, corr)),
2086 );
2087 }
2088 holdouts.held.sort_unstable();
2089 holdouts.held.truncate(DRAIN_HELD_REQUESTS_LISTED);
2090 let mut connections = by_connection.into_iter().collect::<Vec<_>>();
2091 connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
2092 connections.truncate(3);
2093 holdouts.top_connections = connections;
2094 Ok(holdouts)
2095 }
2096
2097 pub(crate) fn endpoint_in_flight_count(
2098 &self,
2099 endpoint: ModuleEndpointId,
2100 ) -> Result<usize, ForwardingError> {
2101 let inner = self.read_inner()?;
2102 Ok(inner
2103 .client_to_module
2104 .values()
2105 .filter(|route| route.module_endpoint == endpoint)
2106 .map(|route| route.flow.drain_in_flight())
2107 .sum())
2108 }
2109
2110 pub(crate) fn endpoint_is_draining(
2111 &self,
2112 endpoint: ModuleEndpointId,
2113 ) -> Result<bool, ForwardingError> {
2114 Ok(self
2115 .read_inner()?
2116 .draining_endpoints
2117 .contains_key(&endpoint))
2118 }
2119
2120 pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
2121 let inner = self.read_inner()?;
2122 Ok(inner
2123 .modules_by_id
2124 .get(module_id)
2125 .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
2126 }
2127
2128 pub(crate) fn release_module_endpoint_routes(
2129 &self,
2130 endpoint: ModuleEndpointId,
2131 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2132 let mut inner = self.write_inner()?;
2133 let routes = inner
2134 .module_to_client
2135 .iter()
2136 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2137 .map(|(module_key, route)| (*module_key, route.module_epoch))
2138 .collect::<Vec<_>>();
2139 let mut released = Vec::with_capacity(routes.len());
2140 for (module_key, epoch) in routes {
2141 if let RouteRelease::Removed(target) =
2142 release_module_route_locked(&mut inner, module_key, epoch)
2143 {
2144 released.push(target);
2145 }
2146 }
2147 Ok(released)
2148 }
2149
2150 pub(crate) fn endpoint_routes(
2156 &self,
2157 endpoint: ModuleEndpointId,
2158 ) -> Result<Vec<EndpointRoute>, ForwardingError> {
2159 let inner = self.read_inner()?;
2160 Ok(endpoint_routes_locked(&inner, endpoint))
2161 }
2162
2163 pub(crate) fn route_census(
2165 &self,
2166 module_id: Option<&str>,
2167 ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
2168 let inner = self.read_inner()?;
2169 let mut endpoints = inner
2170 .modules_by_id
2171 .iter()
2172 .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
2173 .map(|(id, module)| (id.clone(), module.endpoint))
2174 .collect::<Vec<_>>();
2175 endpoints.sort_by(|left, right| left.0.cmp(&right.0));
2176 Ok(endpoints
2177 .into_iter()
2178 .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
2179 .collect())
2180 }
2181
2182 pub(crate) fn live_roots(
2184 &self,
2185 module_id: &str,
2186 ) -> Result<ModuleControlResponseToModule, ForwardingError> {
2187 let inner = self.read_inner()?;
2188 let endpoint = inner
2189 .modules_by_id
2190 .get(module_id)
2191 .map(|module| module.endpoint);
2192 let mut roots = BTreeMap::new();
2193 let mut unknown_root_bindings = 0;
2194 let mut total_bindings = 0;
2195 if let Some(endpoint) = endpoint {
2196 for binding in inner
2197 .module_to_client
2198 .values()
2199 .filter(|binding| binding.module_endpoint == endpoint)
2200 {
2201 total_bindings += 1;
2202 if let Some(root) = &binding.project_root {
2203 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2204 entry.0 += 1;
2205 } else {
2206 unknown_root_bindings += 1;
2207 }
2208 }
2209 for pending in inner
2210 .pending_relays
2211 .values()
2212 .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
2213 {
2214 total_bindings += 1;
2215 if let Some(root) = &pending.reservation.project_root {
2216 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
2217 entry.1 += 1;
2218 } else {
2219 unknown_root_bindings += 1;
2220 }
2221 }
2222 }
2223 Ok(ModuleControlResponseToModule::LiveRoots {
2224 roots: roots
2225 .into_iter()
2226 .map(|(project_root, (bound, pending))| LiveRoot {
2227 project_root,
2228 bound,
2229 pending,
2230 })
2231 .collect(),
2232 unknown_root_bindings,
2233 total_bindings,
2234 })
2235 }
2236
2237 pub(crate) fn connection_has_client_routes(
2243 &self,
2244 connection_id: ConnectionId,
2245 ) -> Result<bool, ForwardingError> {
2246 let inner = self.read_inner()?;
2247 let has = inner
2248 .client_to_module
2249 .keys()
2250 .any(|key| key.connection_id == connection_id)
2251 || inner
2252 .reserved_client
2253 .keys()
2254 .any(|key| key.connection_id == connection_id);
2255 Ok(has)
2256 }
2257
2258 pub(crate) fn cleanup_connection(
2259 &self,
2260 connection_id: ConnectionId,
2261 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
2262 let mut inner = self.write_inner()?;
2263 inner.closing_connections.insert(connection_id);
2264 let released = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
2265 remove_module_connection_locked(&mut inner, endpoint)
2266 } else {
2267 Self::cleanup_client_connection_locked(&mut inner, connection_id)
2268 };
2269 inner.closing_connections.remove(&connection_id);
2278 Ok(released)
2279 }
2280
2281 fn cleanup_client_connection_locked(
2282 inner: &mut ForwardingInner,
2283 connection_id: ConnectionId,
2284 ) -> Vec<GoodbyeTarget> {
2285 let routes = inner
2286 .client_to_module
2287 .iter()
2288 .filter(|(key, _)| key.connection_id == connection_id)
2289 .map(|(key, route)| (*key, route.client_epoch))
2290 .collect::<Vec<_>>();
2291 let mut released = Vec::with_capacity(routes.len());
2292 for (client_key, epoch) in routes {
2293 if let RouteRelease::Removed(target) =
2294 release_client_route_locked(inner, client_key, epoch)
2295 {
2296 released.push(target);
2297 }
2298 }
2299
2300 let pending_keys = inner
2301 .pending_relays
2302 .iter()
2303 .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2304 .map(|(key, _)| *key)
2305 .collect::<Vec<_>>();
2306 for key in pending_keys {
2307 let Some(pending) = inner.pending_relays.remove(&key) else {
2308 continue;
2309 };
2310 release_reserved_route_locked(
2311 inner,
2312 pending.reservation.client_key,
2313 pending.reservation.module_key,
2314 );
2315 if pending.relay_enqueued {
2316 if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2317 released.push(target);
2318 }
2319 }
2320 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2321 "client connection closed during route.bind relay".to_string(),
2322 ));
2323 }
2324
2325 let orphaned = inner
2326 .reserved_client
2327 .iter()
2328 .filter(|(key, _)| key.connection_id == connection_id)
2329 .map(|(client, module)| (*client, *module))
2330 .collect::<Vec<_>>();
2331 for (client_key, module_key) in orphaned {
2332 release_reserved_route_locked(inner, client_key, module_key);
2333 }
2334 inner.next_client_channel.remove(&connection_id);
2335 inner
2336 .client_slot_epochs
2337 .retain(|key, _| key.connection_id != connection_id);
2338 inner
2339 .last_published_epoch
2340 .retain(|key, _| key.connection_id != connection_id);
2341 inner
2342 .status
2343 .retain(|(key, _), _| key.connection_id != connection_id);
2344
2345 released
2346 }
2347
2348 pub(crate) fn escalate_client_delivery_failure(
2357 &self,
2358 connection_id: ConnectionId,
2359 channel: u16,
2360 expected_epoch: u32,
2361 reason: CloseReason,
2362 undelivered: UndeliveredFrame<'_>,
2363 ) -> Result<bool, ForwardingError> {
2364 let principals = {
2365 let mut inner = self.write_inner()?;
2366 let key = ClientRouteKey {
2367 connection_id,
2368 channel,
2369 };
2370 if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2371 None
2372 } else {
2373 inner.closing_connections.insert(connection_id);
2374 Some(connection_principals_locked(&inner, connection_id))
2375 }
2376 };
2377 let Some(principals) = principals else {
2378 return Ok(false);
2379 };
2380 let backlog = undelivered.sink.backlog();
2381 let close_reason = reason.to_string();
2382 if self.request_connection_close(connection_id, reason) {
2383 warn!(
2384 connection_id = connection_id.get(),
2385 principals = %principals,
2386 module_id = undelivered.module_id.unwrap_or("unknown"),
2387 client_channel = channel,
2388 queued_bytes = backlog.queued_bytes,
2389 queued_frames = backlog.queued_frames,
2390 oldest_queued_ms = backlog
2391 .oldest_age
2392 .map(|age| age.as_millis() as u64)
2393 .unwrap_or(0),
2394 close_reason = %close_reason,
2395 "closing client connection: its egress queue could not take a frame"
2396 );
2397 }
2398 Ok(true)
2399 }
2400
2401 fn record_route_release(&self, release: &RouteRelease) {
2402 match release {
2403 RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2404 RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2405 RouteRelease::Absent => {}
2406 }
2407 }
2408
2409 fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2410 self.inner.read().map_err(|_| ForwardingError::Poisoned)
2411 }
2412
2413 fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2414 self.inner.write().map_err(|_| ForwardingError::Poisoned)
2415 }
2416
2417 fn lock_close_registry(
2418 &self,
2419 ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2420 self.close_registry
2421 .lock()
2422 .unwrap_or_else(|poisoned| poisoned.into_inner())
2423 }
2424}
2425
2426impl ForwardingInner {
2427 fn allocate_route_slots(
2428 &mut self,
2429 connection_id: ConnectionId,
2430 endpoint: ModuleEndpointId,
2431 ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2432 let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2433 let mut client_channel = client_start;
2434 let client_channel = loop {
2435 let key = ClientRouteKey {
2436 connection_id,
2437 channel: client_channel,
2438 };
2439 let eligible = !self.client_to_module.contains_key(&key)
2440 && !self.reserved_client.contains_key(&key)
2441 && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2442 if eligible {
2443 break client_channel;
2444 }
2445 client_channel = next_channel(client_channel);
2446 if client_channel == client_start {
2447 return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2448 }
2449 };
2450
2451 let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2452 let mut module_channel = module_start;
2453 let module_channel = loop {
2454 let key = ModuleRouteKey {
2455 endpoint,
2456 channel: module_channel,
2457 };
2458 let eligible = !self.module_to_client.contains_key(&key)
2459 && !self.reserved_module.contains_key(&key)
2460 && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2461 if eligible {
2462 break module_channel;
2463 }
2464 module_channel = next_channel(module_channel);
2465 if module_channel == module_start {
2466 return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2467 }
2468 };
2469
2470 let client_key = ClientRouteKey {
2471 connection_id,
2472 channel: client_channel,
2473 };
2474 let module_key = ModuleRouteKey {
2475 endpoint,
2476 channel: module_channel,
2477 };
2478 let client_epoch = self
2479 .client_slot_epochs
2480 .get(&client_key)
2481 .copied()
2482 .unwrap_or(0)
2483 + 1;
2484 let module_epoch = self
2485 .module_slot_epochs
2486 .get(&module_key)
2487 .copied()
2488 .unwrap_or(0)
2489 + 1;
2490 self.client_slot_epochs.insert(client_key, client_epoch);
2491 self.module_slot_epochs.insert(module_key, module_epoch);
2492 self.next_client_channel
2493 .insert(connection_id, next_channel(client_channel));
2494 self.next_module_channel
2495 .insert(endpoint, next_channel(module_channel));
2496 Ok((client_channel, client_epoch, module_channel, module_epoch))
2497 }
2498
2499 fn allocate_control_corr(
2500 &mut self,
2501 endpoint: ModuleEndpointId,
2502 ) -> Result<u64, ForwardingError> {
2503 let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2504 if candidate == 0 {
2505 self.closing_connections.insert(endpoint.connection_id);
2506 return Err(ForwardingError::RelayCorrelationExhausted);
2507 }
2508 self.next_control_corr.insert(
2509 endpoint,
2510 if candidate == u64::MAX {
2511 0
2512 } else {
2513 candidate + 1
2514 },
2515 );
2516 Ok(candidate)
2517 }
2518}
2519
2520fn next_channel(channel: u16) -> u16 {
2521 let next = channel.wrapping_add(1);
2522 if next == 0 {
2523 1
2524 } else {
2525 next
2526 }
2527}
2528
2529fn endpoint_routes_locked(
2530 inner: &ForwardingInner,
2531 endpoint: ModuleEndpointId,
2532) -> Vec<EndpointRoute> {
2533 let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2534 let draining = drain_reason.is_some();
2535 let mut routes = inner
2536 .module_to_client
2537 .iter()
2538 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2539 .map(|(_, route)| EndpointRoute {
2540 goodbye_target: GoodbyeTarget {
2541 connection_id: route.client_connection_id,
2542 sink: route.client_sink.clone(),
2543 negotiated_ver: route.client_negotiated_ver,
2544 channel: route.client_channel,
2545 epoch: route.client_epoch,
2546 kind: GoodbyeTargetKind::Client,
2547 module_id: Some(route.module_id.clone()),
2548 },
2549 principal: route.principal.clone(),
2550 bound_at: route.bound_at,
2551 draining,
2552 drain_reason,
2553 })
2554 .collect::<Vec<_>>();
2555 routes.sort_by_key(|route| {
2556 (
2557 route.goodbye_target.connection_id.get(),
2558 route.goodbye_target.channel,
2559 route.goodbye_target.epoch,
2560 )
2561 });
2562 routes
2563}
2564
2565fn release_reserved_route_locked(
2566 inner: &mut ForwardingInner,
2567 client_key: ClientRouteKey,
2568 module_key: ModuleRouteKey,
2569) {
2570 if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
2571 inner.reserved_client.remove(&client_key);
2572 }
2573 if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
2574 inner.reserved_module.remove(&module_key);
2575 }
2576 inner.status.retain(|(key, _), _| *key != client_key);
2577}
2578
2579fn release_client_route_locked(
2580 inner: &mut ForwardingInner,
2581 client_key: ClientRouteKey,
2582 expected_epoch: u32,
2583) -> RouteRelease {
2584 let Some(route) = inner.client_to_module.get(&client_key) else {
2585 return RouteRelease::Absent;
2586 };
2587 if route.client_epoch != expected_epoch {
2588 return RouteRelease::Stale;
2589 }
2590 let route = inner
2591 .client_to_module
2592 .remove(&client_key)
2593 .expect("route checked under the same forwarding lock");
2594 route.flow.close();
2595 inner.module_to_client.remove(&ModuleRouteKey {
2596 endpoint: route.module_endpoint,
2597 channel: route.module_channel,
2598 });
2599 inner.status.remove(&(client_key, expected_epoch));
2600 RouteRelease::Removed(GoodbyeTarget {
2601 connection_id: route.module_endpoint.connection_id,
2602 sink: route.module_sink.clone(),
2603 negotiated_ver: route.module_negotiated_ver,
2604 channel: route.module_channel,
2605 epoch: route.module_epoch,
2606 kind: GoodbyeTargetKind::Module,
2607 module_id: Some(route.module_id.clone()),
2608 })
2609}
2610
2611fn release_module_route_locked(
2612 inner: &mut ForwardingInner,
2613 module_key: ModuleRouteKey,
2614 expected_epoch: u32,
2615) -> RouteRelease {
2616 let Some(route) = inner.module_to_client.get(&module_key) else {
2617 return RouteRelease::Absent;
2618 };
2619 if route.module_epoch != expected_epoch {
2620 return RouteRelease::Stale;
2621 }
2622 let route = inner
2623 .module_to_client
2624 .remove(&module_key)
2625 .expect("route checked under the same forwarding lock");
2626 route.flow.close();
2627 let client_key = ClientRouteKey {
2628 connection_id: route.client_connection_id,
2629 channel: route.client_channel,
2630 };
2631 inner.client_to_module.remove(&client_key);
2632 inner.status.remove(&(client_key, route.client_epoch));
2633 RouteRelease::Removed(GoodbyeTarget {
2634 connection_id: route.client_connection_id,
2635 sink: route.client_sink.clone(),
2636 negotiated_ver: route.client_negotiated_ver,
2637 channel: route.client_channel,
2638 epoch: route.client_epoch,
2639 kind: GoodbyeTargetKind::Client,
2640 module_id: Some(route.module_id.clone()),
2641 })
2642}
2643
2644fn commit_route_locked(
2645 inner: &mut ForwardingInner,
2646 pending: PendingRouteBindRelayEntry,
2647) -> Result<Option<GoodbyeTarget>, ForwardingError> {
2648 let reservation = pending.reservation;
2649 if inner
2650 .closing_connections
2651 .contains(&reservation.client_key.connection_id)
2652 {
2653 return Err(ForwardingError::ConnectionClosing {
2654 connection_id: reservation.client_key.connection_id,
2655 });
2656 }
2657 let module_id = inner
2658 .module_id_by_endpoint
2659 .get(&reservation.module_key.endpoint)
2660 .cloned()
2661 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2662 if inner
2663 .draining_endpoints
2664 .contains_key(&reservation.module_key.endpoint)
2665 {
2666 return Err(ForwardingError::ModuleReloading { module_id });
2667 }
2668 if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
2669 || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
2670 {
2671 return Err(ForwardingError::UnknownReservation {
2672 client_channel: reservation.client_key.channel,
2673 module_channel: reservation.module_key.channel,
2674 });
2675 }
2676 let module = inner
2677 .modules_by_id
2678 .get(&module_id)
2679 .filter(|module| module.endpoint == reservation.module_key.endpoint)
2680 .cloned()
2681 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2682 let binding = Arc::new(RouteBinding {
2683 client_connection_id: reservation.client_key.connection_id,
2684 client_sink: pending.client_sink,
2685 client_negotiated_ver: pending.client_negotiated_ver,
2686 client_channel: reservation.client_key.channel,
2687 client_epoch: reservation.client_epoch,
2688 module_id,
2689 module_endpoint: reservation.module_key.endpoint,
2690 module_sink: module.sink,
2691 module_negotiated_ver: module.negotiated_ver,
2692 module_channel: reservation.module_key.channel,
2693 module_epoch: reservation.module_epoch,
2694 principal: pending.principal,
2695 project_root: reservation.project_root.clone(),
2696 bound_at: Instant::now(),
2697 flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
2698 });
2699 inner
2700 .client_to_module
2701 .insert(reservation.client_key, Arc::clone(&binding));
2702 inner
2703 .module_to_client
2704 .insert(reservation.module_key, binding);
2705 let previous_published = inner
2706 .last_published_epoch
2707 .insert(reservation.client_key, reservation.client_epoch);
2708
2709 let client_writer_closed = pending.client_permit.send(pending.route_open_frame);
2714 if client_writer_closed {
2715 let abandoned = pending
2716 .relay_enqueued
2717 .then(|| abandoned_route_target(inner, &reservation))
2718 .flatten();
2719 if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
2720 route.flow.close();
2721 }
2722 inner.module_to_client.remove(&reservation.module_key);
2723 inner
2724 .status
2725 .remove(&(reservation.client_key, reservation.client_epoch));
2726 match previous_published {
2727 Some(epoch) => {
2728 inner
2729 .last_published_epoch
2730 .insert(reservation.client_key, epoch);
2731 }
2732 None => {
2733 inner.last_published_epoch.remove(&reservation.client_key);
2734 }
2735 }
2736 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2737 "client egress closed during route publication".to_string(),
2738 ));
2739 return Ok(abandoned);
2740 }
2741
2742 let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
2743 Ok(None)
2744}
2745
2746fn module_connection_for_endpoint_locked(
2754 inner: &ForwardingInner,
2755 endpoint: ModuleEndpointId,
2756) -> Option<&ModuleConnection> {
2757 let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
2758 inner
2759 .modules_by_id
2760 .get(module_id)
2761 .filter(|module| module.endpoint == endpoint)
2762 .or_else(|| {
2763 inner
2764 .candidates_by_id
2765 .get(module_id)
2766 .filter(|module| module.endpoint == endpoint)
2767 })
2768 .or_else(|| inner.superseded_endpoints.get(&endpoint))
2769}
2770
2771fn abandoned_route_target(
2772 inner: &ForwardingInner,
2773 reservation: &RouteReservation,
2774) -> Option<GoodbyeTarget> {
2775 let module_id = inner
2776 .module_id_by_endpoint
2777 .get(&reservation.module_key.endpoint)?;
2778 let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
2779 (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
2780 connection_id: module.endpoint.connection_id,
2781 sink: module.sink.clone(),
2782 negotiated_ver: module.negotiated_ver,
2783 channel: reservation.module_key.channel,
2784 epoch: reservation.module_epoch,
2785 kind: GoodbyeTargetKind::Module,
2786 module_id: Some(module_id.clone()),
2787 })
2788}
2789
2790fn enqueue_hello_ack_locked(
2795 sink: &FrameSink,
2796 connection_id: ConnectionId,
2797 hello_ack: Option<Frame>,
2798) -> Result<(), ForwardingError> {
2799 let Some(hello_ack) = hello_ack else {
2800 return Ok(());
2801 };
2802 sink.try_send(hello_ack)
2803 .map_err(|_| ForwardingError::ModuleEgressUnavailable { connection_id })
2804}
2805
2806fn remove_module_connection_locked(
2807 inner: &mut ForwardingInner,
2808 endpoint: ModuleEndpointId,
2809) -> Vec<GoodbyeTarget> {
2810 inner.draining_endpoints.remove(&endpoint);
2811 let module_id = inner.module_id_by_endpoint.remove(&endpoint);
2812 if let Some(module_id) = module_id.as_ref() {
2813 if inner
2814 .modules_by_id
2815 .get(module_id)
2816 .is_some_and(|module| module.endpoint == endpoint)
2817 {
2818 inner.modules_by_id.remove(module_id);
2819 }
2820 if inner
2821 .candidates_by_id
2822 .get(module_id)
2823 .is_some_and(|module| module.endpoint == endpoint)
2824 {
2825 inner.candidates_by_id.remove(module_id);
2826 }
2827 }
2828 inner.superseded_endpoints.remove(&endpoint);
2829 inner.endpoint_by_connection.remove(&endpoint.connection_id);
2830 inner.next_module_channel.remove(&endpoint);
2831 inner.next_control_corr.remove(&endpoint);
2832 inner
2833 .health_probe_tombstones
2834 .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
2835 inner
2836 .module_slot_epochs
2837 .retain(|key, _| key.endpoint != endpoint);
2838 let reserved_module_keys: Vec<ModuleRouteKey> = inner
2839 .reserved_module
2840 .keys()
2841 .filter(|module_key| module_key.endpoint == endpoint)
2842 .copied()
2843 .collect();
2844 for module_key in reserved_module_keys {
2845 if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2846 release_reserved_route_locked(inner, client_key, module_key);
2847 }
2848 }
2849
2850 let pending_keys: Vec<_> = inner
2851 .pending_relays
2852 .keys()
2853 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2854 .copied()
2855 .collect();
2856 let pending: Vec<_> = pending_keys
2857 .into_iter()
2858 .filter_map(|key| inner.pending_relays.remove(&key))
2859 .collect();
2860 for pending in pending {
2861 let module_label = module_id.as_deref().unwrap_or("unknown");
2862 let _ = pending
2863 .sender
2864 .send(RouteBindRelayOutcome::ModuleGone(format!(
2865 "module '{module_label}' connection closed during route.bind relay"
2866 )));
2867 }
2868
2869 let pending_control_keys: Vec<_> = inner
2870 .pending_control_rpcs
2871 .keys()
2872 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2873 .copied()
2874 .collect();
2875 let pending_control: Vec<_> = pending_control_keys
2876 .into_iter()
2877 .filter_map(|key| inner.pending_control_rpcs.remove(&key))
2878 .collect();
2879 for pending in pending_control {
2880 let module_label = module_id.as_deref().unwrap_or("unknown");
2881 let _ = pending
2882 .sender
2883 .send(ModuleControlRpcOutcome::ModuleGone(format!(
2884 "module '{module_label}' connection closed during module-control RPC"
2885 )));
2886 }
2887
2888 let module_routes = inner
2889 .module_to_client
2890 .iter()
2891 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2892 .map(|(module_key, route)| (*module_key, route.module_epoch))
2893 .collect::<Vec<_>>();
2894 let mut released = Vec::with_capacity(module_routes.len());
2895 for (module_key, epoch) in module_routes {
2896 if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
2897 {
2898 released.push(target);
2899 }
2900 }
2901 released
2902}
2903
2904#[derive(Debug, Clone, Copy)]
2905struct RequestCredit {
2906 subscription: bool,
2907 excluded_from_drain: bool,
2908}
2909
2910#[derive(Debug, Default)]
2911struct CreditLedger {
2912 by_corr: HashMap<u64, Vec<RequestCredit>>,
2913}
2914
2915impl CreditLedger {
2916 fn acquire(&mut self, corr: u64, subscription: bool) {
2917 self.by_corr.entry(corr).or_default().push(RequestCredit {
2918 subscription,
2919 excluded_from_drain: false,
2920 });
2921 }
2922
2923 fn release(&mut self, corr: u64) -> bool {
2924 let Some(credits) = self.by_corr.get_mut(&corr) else {
2925 return false;
2926 };
2927 let released = credits.pop().is_some();
2928 if credits.is_empty() {
2929 self.by_corr.remove(&corr);
2930 }
2931 released
2932 }
2933
2934 fn capture_subscription_exclusions(&mut self) -> u32 {
2935 let mut excluded = 0u32;
2936 for credit in self.by_corr.values_mut().flatten() {
2937 if credit.subscription && !credit.excluded_from_drain {
2938 credit.excluded_from_drain = true;
2939 excluded = excluded.saturating_add(1);
2940 }
2941 }
2942 excluded
2943 }
2944
2945 #[cfg(test)]
2946 fn in_flight(&self) -> usize {
2947 self.by_corr.values().map(Vec::len).sum()
2948 }
2949
2950 fn drain_in_flight(&self) -> usize {
2951 self.by_corr
2952 .values()
2953 .flatten()
2954 .filter(|credit| !credit.excluded_from_drain)
2955 .count()
2956 }
2957
2958 fn drain_held_corrs(&self) -> Vec<u64> {
2961 let mut corrs = self
2962 .by_corr
2963 .iter()
2964 .flat_map(|(corr, credits)| {
2965 credits
2966 .iter()
2967 .filter(|credit| !credit.excluded_from_drain)
2968 .map(move |_| *corr)
2969 })
2970 .collect::<Vec<_>>();
2971 corrs.sort_unstable();
2972 corrs
2973 }
2974}
2975
2976#[derive(Debug, Default)]
2977struct ChannelFlowState {
2978 closed: bool,
2979 credits: CreditLedger,
2980}
2981
2982#[derive(Debug)]
2984pub(crate) struct ChannelFlow {
2985 sem: Semaphore,
2986 window: usize,
2987 state: Mutex<ChannelFlowState>,
2988}
2989
2990impl ChannelFlow {
2991 pub(crate) fn new(window: usize) -> Self {
2992 debug_assert!(window > 0, "flow-control window must be non-zero");
2993 Self {
2994 sem: Semaphore::new(window),
2995 window,
2996 state: Mutex::new(ChannelFlowState::default()),
2997 }
2998 }
2999
3000 #[cfg(test)]
3001 pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
3002 self.acquire_tagged(0, false).await
3003 }
3004
3005 pub(crate) async fn acquire_tagged(
3006 &self,
3007 corr: u64,
3008 subscription: bool,
3009 ) -> Result<(), ChannelFlowClosed> {
3010 let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
3011 let mut state = self
3012 .state
3013 .lock()
3014 .unwrap_or_else(|poisoned| poisoned.into_inner());
3015 if state.closed {
3016 return Err(ChannelFlowClosed);
3017 }
3018 state.credits.acquire(corr, subscription);
3019 permit.forget();
3020 Ok(())
3021 }
3022
3023 #[cfg(test)]
3024 pub(crate) fn release(&self) {
3025 self.release_corr(0);
3026 }
3027
3028 pub(crate) fn release_corr(&self, corr: u64) {
3029 let released = self
3030 .state
3031 .lock()
3032 .unwrap_or_else(|poisoned| poisoned.into_inner())
3033 .credits
3034 .release(corr);
3035 if !released {
3036 warn!(
3040 window = self.window,
3041 available = self.sem.available_permits(),
3042 "flow-control over-release ignored"
3043 );
3044 return;
3045 }
3046 if !self.sem.is_closed() {
3047 self.sem.add_permits(1);
3048 }
3049 }
3050
3051 #[cfg(test)]
3052 pub(crate) fn in_flight(&self) -> usize {
3053 self.state
3054 .lock()
3055 .unwrap_or_else(|poisoned| poisoned.into_inner())
3056 .credits
3057 .in_flight()
3058 }
3059
3060 pub(crate) fn drain_in_flight(&self) -> usize {
3061 self.state
3062 .lock()
3063 .unwrap_or_else(|poisoned| poisoned.into_inner())
3064 .credits
3065 .drain_in_flight()
3066 }
3067
3068 pub(crate) fn drain_held_corrs(&self) -> Vec<u64> {
3069 self.state
3070 .lock()
3071 .unwrap_or_else(|poisoned| poisoned.into_inner())
3072 .credits
3073 .drain_held_corrs()
3074 }
3075
3076 #[cfg(test)]
3077 pub(crate) fn available_permits(&self) -> usize {
3078 self.sem.available_permits()
3079 }
3080
3081 pub(crate) fn begin_drain(&self) -> u32 {
3082 let mut state = self
3083 .state
3084 .lock()
3085 .unwrap_or_else(|poisoned| poisoned.into_inner());
3086 state.closed = true;
3087 self.sem.close();
3088 state.credits.capture_subscription_exclusions()
3089 }
3090
3091 pub(crate) fn close(&self) {
3092 self.state
3093 .lock()
3094 .unwrap_or_else(|poisoned| poisoned.into_inner())
3095 .closed = true;
3096 self.sem.close();
3097 }
3098}
3099
3100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3101pub(crate) struct ChannelFlowClosed;
3102
3103impl fmt::Display for ChannelFlowClosed {
3104 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3105 write!(f, "flow-control window closed")
3106 }
3107}
3108
3109impl Error for ChannelFlowClosed {}
3110
3111fn window_for(concurrency: &Concurrency) -> usize {
3112 match concurrency {
3113 Concurrency::Serial => 1,
3114 Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
3115 Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
3116 }
3117}
3118
3119#[derive(Debug, Clone, PartialEq, Eq)]
3120pub enum ForwardingError {
3121 NoModuleConnection,
3122 ModuleReloading {
3123 module_id: String,
3124 },
3125 StaleModuleEndpoint,
3126 UnknownReservation {
3127 client_channel: u16,
3128 module_channel: u16,
3129 },
3130 ClientRouteChannelExhausted {
3131 connection_id: ConnectionId,
3132 },
3133 ModuleRouteChannelExhausted {
3134 endpoint: ModuleEndpointId,
3135 },
3136 RelayCorrelationExhausted,
3137 ConnectionClosing {
3138 connection_id: ConnectionId,
3139 },
3140 ClientEgressClosed {
3141 connection_id: ConnectionId,
3142 },
3143 RouteOpenBuild(String),
3144 CandidateSlotOccupied {
3146 module_id: String,
3147 },
3148 ModuleEgressUnavailable {
3151 connection_id: ConnectionId,
3152 },
3153 Poisoned,
3154}
3155
3156impl fmt::Display for ForwardingError {
3157 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3158 match self {
3159 Self::NoModuleConnection => write!(f, "no module connection is registered"),
3160 Self::ModuleReloading { module_id } => {
3161 write!(f, "module_id '{module_id}' is reloading")
3162 }
3163 Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
3164 Self::UnknownReservation {
3165 client_channel,
3166 module_channel,
3167 } => write!(
3168 f,
3169 "route reservation client channel {client_channel} / module channel {module_channel} was not found"
3170 ),
3171 Self::ClientRouteChannelExhausted { connection_id } => write!(
3172 f,
3173 "no client route channels are available for connection {}",
3174 connection_id.get()
3175 ),
3176 Self::ModuleRouteChannelExhausted { endpoint } => write!(
3177 f,
3178 "no module route channels are available for endpoint generation {} on connection {}",
3179 endpoint.generation,
3180 endpoint.connection_id.get()
3181 ),
3182 Self::RelayCorrelationExhausted => {
3183 write!(f, "module control correlation ids are exhausted")
3184 }
3185 Self::ConnectionClosing { connection_id } => write!(
3186 f,
3187 "connection {} is closing and cannot accept route allocation",
3188 connection_id.get()
3189 ),
3190 Self::ClientEgressClosed { connection_id } => write!(
3191 f,
3192 "client connection {} egress is closed",
3193 connection_id.get()
3194 ),
3195 Self::RouteOpenBuild(message) => {
3196 write!(f, "failed to prebuild route.open response: {message}")
3197 }
3198 Self::CandidateSlotOccupied { module_id } => write!(
3199 f,
3200 "module_id '{module_id}' already has a swap candidate registered"
3201 ),
3202 Self::ModuleEgressUnavailable { connection_id } => write!(
3203 f,
3204 "module connection {} egress is unavailable; HELLO_ACK could not be queued",
3205 connection_id.get()
3206 ),
3207 Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
3208 }
3209 }
3210}
3211
3212impl Error for ForwardingError {}
3213
3214#[cfg(test)]
3215mod tests {
3216 use std::time::Duration;
3217
3218 use super::*;
3219 use tokio::sync::mpsc;
3220
3221 #[test]
3222 fn ordinary_long_running_request_is_not_excluded_from_drain() {
3223 let mut ledger = CreditLedger::default();
3224 ledger.acquire(1, false);
3225
3226 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3227 assert_eq!(ledger.drain_in_flight(), 1);
3228 }
3229
3230 #[test]
3231 fn bit_set_subscription_is_excluded_and_counted() {
3232 let mut ledger = CreditLedger::default();
3233 ledger.acquire(1, true);
3234
3235 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3236 assert_eq!(ledger.drain_in_flight(), 0);
3237 }
3238
3239 #[test]
3240 fn subscription_opened_after_drain_snapshot_is_not_excluded() {
3241 let mut ledger = CreditLedger::default();
3242 ledger.acquire(1, true);
3243 assert_eq!(ledger.capture_subscription_exclusions(), 1);
3244
3245 ledger.acquire(2, true);
3246
3247 assert_eq!(ledger.drain_in_flight(), 1);
3248 }
3249
3250 #[test]
3251 fn drain_with_no_subscriptions_reports_zero_excluded() {
3252 let mut ledger = CreditLedger::default();
3253 assert_eq!(ledger.capture_subscription_exclusions(), 0);
3254 }
3255
3256 fn test_hello_ack(corr: u64) -> Frame {
3257 Frame::build(
3258 FrameType::HelloAck,
3259 Flags::new(false, Priority::Passive, false),
3260 0,
3261 0,
3262 corr,
3263 Vec::new(),
3264 )
3265 .unwrap()
3266 }
3267
3268 #[test]
3272 fn acked_registration_that_cannot_queue_its_hello_ack_inserts_nothing() {
3273 let forwarding = ForwardingTable::default();
3274
3275 let (closed_tx, closed_rx) = mpsc::channel(8);
3276 drop(closed_rx);
3277 let closed = ConnectionId::new(1);
3278 assert_eq!(
3279 forwarding.register_module_connection_acked(
3280 closed,
3281 "closed".to_string(),
3282 2,
3283 Concurrency::ModuleManaged,
3284 FrameSink::new(closed_tx),
3285 test_hello_ack(1),
3286 ),
3287 Err(ForwardingError::ModuleEgressUnavailable {
3288 connection_id: closed
3289 })
3290 );
3291
3292 let (full_tx, _full_rx) = mpsc::channel(1);
3293 let full_sink = FrameSink::new(full_tx);
3294 full_sink.try_send(test_hello_ack(99)).unwrap();
3295 let full = ConnectionId::new(2);
3296 assert_eq!(
3297 forwarding.register_module_connection_acked(
3298 full,
3299 "full".to_string(),
3300 2,
3301 Concurrency::ModuleManaged,
3302 full_sink.clone(),
3303 test_hello_ack(2),
3304 ),
3305 Err(ForwardingError::ModuleEgressUnavailable {
3306 connection_id: full
3307 })
3308 );
3309 assert_eq!(
3310 forwarding.register_candidate_module_connection_acked(
3311 full,
3312 "full".to_string(),
3313 2,
3314 Concurrency::ModuleManaged,
3315 full_sink,
3316 test_hello_ack(3),
3317 ),
3318 Err(ForwardingError::ModuleEgressUnavailable {
3319 connection_id: full
3320 })
3321 );
3322
3323 for (connection, module_id) in [(closed, "closed"), (full, "full")] {
3324 assert_eq!(
3325 forwarding
3326 .module_endpoint_for_connection(connection)
3327 .unwrap(),
3328 None
3329 );
3330 let (client_tx, _client_rx) = mpsc::channel(8);
3331 assert_eq!(
3332 forwarding
3333 .begin_route_bind_relay_for_test(
3334 ConnectionId::new(50),
3335 FrameSink::new(client_tx),
3336 1,
3337 module_id,
3338 )
3339 .err(),
3340 Some(ForwardingError::NoModuleConnection)
3341 );
3342 }
3343 assert!(forwarding.read_inner().unwrap().candidates_by_id.is_empty());
3344 }
3345
3346 #[test]
3349 fn acked_registration_queues_the_hello_ack_first() {
3350 let forwarding = ForwardingTable::default();
3351 let (active_tx, mut active_rx) = mpsc::channel(8);
3352 forwarding
3353 .register_module_connection_acked(
3354 ConnectionId::new(1),
3355 "acked".to_string(),
3356 2,
3357 Concurrency::ModuleManaged,
3358 FrameSink::new(active_tx),
3359 test_hello_ack(11),
3360 )
3361 .unwrap();
3362 let (candidate_tx, mut candidate_rx) = mpsc::channel(8);
3363 forwarding
3364 .register_candidate_module_connection_acked(
3365 ConnectionId::new(2),
3366 "acked".to_string(),
3367 2,
3368 Concurrency::ModuleManaged,
3369 FrameSink::new(candidate_tx),
3370 test_hello_ack(12),
3371 )
3372 .unwrap();
3373
3374 let active_first = active_rx.try_recv().unwrap().frame;
3375 assert_eq!(active_first.header.ty, FrameType::HelloAck);
3376 assert_eq!(active_first.header.corr, 11);
3377 let candidate_first = candidate_rx.try_recv().unwrap().frame;
3378 assert_eq!(candidate_first.header.ty, FrameType::HelloAck);
3379 assert_eq!(candidate_first.header.corr, 12);
3380 }
3381
3382 #[test]
3383 fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
3384 let forwarding = ForwardingTable::default();
3385 let module_connection = ConnectionId::new(10);
3386 let exhausted_client = ConnectionId::new(20);
3387 let second_client = ConnectionId::new(30);
3388 let (module_tx, _module_rx) = mpsc::channel(1);
3389 let endpoint = forwarding
3390 .register_module_connection(
3391 module_connection,
3392 "route-limit-provider".to_string(),
3393 1,
3394 Concurrency::ModuleManaged,
3395 FrameSink::new(module_tx),
3396 )
3397 .unwrap();
3398
3399 {
3400 let mut inner = forwarding.inner.write().unwrap();
3401 for channel in 1..=u16::MAX {
3402 inner.reserved_client.insert(
3403 ClientRouteKey {
3404 connection_id: exhausted_client,
3405 channel,
3406 },
3407 ModuleRouteKey {
3408 endpoint,
3409 channel: 1,
3410 },
3411 );
3412 }
3413 }
3414
3415 let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
3416 let err = forwarding
3417 .begin_route_bind_relay_for_test(
3418 exhausted_client,
3419 FrameSink::new(exhausted_tx),
3420 1,
3421 "route-limit-provider",
3422 )
3423 .unwrap_err();
3424 assert!(matches!(
3425 err,
3426 ForwardingError::ClientRouteChannelExhausted { connection_id }
3427 if connection_id == exhausted_client
3428 ));
3429
3430 let (second_tx, _second_rx) = mpsc::channel(1);
3431 let pending = forwarding
3432 .begin_route_bind_relay_for_test(
3433 second_client,
3434 FrameSink::new(second_tx),
3435 2,
3436 "route-limit-provider",
3437 )
3438 .unwrap();
3439 assert_eq!(pending.client_channel, 1);
3440 }
3441
3442 #[test]
3443 fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
3444 let forwarding = ForwardingTable::default();
3445 let module_connection = ConnectionId::new(40);
3446 let client = ConnectionId::new(50);
3447 let (module_tx, _module_rx) = mpsc::channel(1);
3448 forwarding
3449 .register_module_connection(
3450 module_connection,
3451 "slot-reuse-provider".to_string(),
3452 1,
3453 Concurrency::ModuleManaged,
3454 FrameSink::new(module_tx),
3455 )
3456 .unwrap();
3457
3458 let (client_tx, _client_rx) = mpsc::channel(1);
3459 let client_sink = FrameSink::new(client_tx);
3460 let mut wrapped_channel = None;
3461 for index in 0..=usize::from(u16::MAX) {
3462 let pending = forwarding
3463 .begin_route_bind_relay_for_test(
3464 client,
3465 client_sink.clone(),
3466 index as u64 + 1,
3467 "slot-reuse-provider",
3468 )
3469 .unwrap();
3470 if index == usize::from(u16::MAX) {
3471 wrapped_channel = Some(pending.module_channel);
3472 }
3473 forwarding
3474 .abort_pending_relay(
3475 pending.endpoint,
3476 pending.corr,
3477 RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
3478 )
3479 .unwrap();
3480 }
3481
3482 assert_eq!(wrapped_channel, Some(1));
3483 }
3484
3485 #[test]
3486 fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
3487 let forwarding = ForwardingTable::default();
3488 let client = ConnectionId::new(60);
3489 forwarding
3490 .inner
3491 .write()
3492 .unwrap()
3493 .next_client_channel
3494 .insert(client, 41);
3495
3496 let released = forwarding.cleanup_connection(client).unwrap();
3497
3498 assert!(released.is_empty());
3499 assert!(!forwarding
3500 .inner
3501 .read()
3502 .unwrap()
3503 .next_client_channel
3504 .contains_key(&client));
3505 }
3506
3507 #[test]
3508 fn stale_module_cleanup_preserves_fast_reconnect_successor() {
3509 let forwarding = ForwardingTable::default();
3510 let module_id = "fast-reconnect-provider";
3511 let first_connection = ConnectionId::new(70);
3512 let second_connection = ConnectionId::new(80);
3513 let (first_tx, _first_rx) = mpsc::channel(1);
3514 let first_endpoint = forwarding
3515 .register_module_connection(
3516 first_connection,
3517 module_id.to_string(),
3518 1,
3519 Concurrency::ModuleManaged,
3520 FrameSink::new(first_tx),
3521 )
3522 .unwrap();
3523 let (second_tx, _second_rx) = mpsc::channel(1);
3524 let second_endpoint = forwarding
3525 .register_module_connection(
3526 second_connection,
3527 module_id.to_string(),
3528 1,
3529 Concurrency::ModuleManaged,
3530 FrameSink::new(second_tx),
3531 )
3532 .unwrap();
3533 assert_ne!(first_endpoint, second_endpoint);
3534
3535 let released = forwarding.cleanup_connection(first_connection).unwrap();
3536
3537 assert!(released.is_empty());
3538 assert_eq!(
3539 forwarding
3540 .inner
3541 .read()
3542 .unwrap()
3543 .modules_by_id
3544 .get(module_id)
3545 .map(|module| module.endpoint),
3546 Some(second_endpoint)
3547 );
3548 assert!(forwarding.has_live_module_connection(module_id).unwrap());
3549 let control_rpc = forwarding
3550 .begin_module_control_rpc_for(
3551 module_id,
3552 "health.check",
3553 Instant::now() + Duration::from_secs(1),
3554 )
3555 .unwrap();
3556 assert_eq!(control_rpc.endpoint, second_endpoint);
3557 }
3558
3559 fn route_fixture(
3560 module_id: &str,
3561 ) -> (
3562 ForwardingTable,
3563 ConnectionId,
3564 ModuleEndpointId,
3565 ConnectionId,
3566 FrameSink,
3567 mpsc::Receiver<crate::router::OutboundFrame>,
3568 ) {
3569 let forwarding = ForwardingTable::default();
3570 let module_connection = ConnectionId::new(100);
3571 let client_connection = ConnectionId::new(200);
3572 let (module_tx, _module_rx) = mpsc::channel(8);
3573 let endpoint = forwarding
3574 .register_module_connection(
3575 module_connection,
3576 module_id.to_string(),
3577 2,
3578 Concurrency::ModuleManaged,
3579 FrameSink::new(module_tx),
3580 )
3581 .unwrap();
3582 let (client_tx, client_rx) = mpsc::channel(8);
3583 (
3584 forwarding,
3585 module_connection,
3586 endpoint,
3587 client_connection,
3588 FrameSink::new(client_tx),
3589 client_rx,
3590 )
3591 }
3592
3593 #[test]
3594 #[cfg(unix)]
3595 fn daemon_drain_gates_current_and_racing_provider_registrations() {
3596 let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
3597 assert_eq!(forwarding.begin_daemon_drain().unwrap(), ["provider"]);
3598 assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
3599 assert!(matches!(
3600 forwarding.register_module_connection(
3601 ConnectionId::new(300),
3602 "late-provider".into(),
3603 2,
3604 Concurrency::ModuleManaged,
3605 sink,
3606 ),
3607 Err(ForwardingError::ConnectionClosing { .. })
3608 ));
3609 }
3610
3611 fn test_ping(corr: u64) -> Frame {
3612 Frame::build(
3613 FrameType::Ping,
3614 Flags::new(false, Priority::Passive, false),
3615 0,
3616 0,
3617 corr,
3618 Vec::new(),
3619 )
3620 .unwrap()
3621 }
3622
3623 fn begin_test_route(
3624 forwarding: &ForwardingTable,
3625 client_connection: ConnectionId,
3626 client_sink: FrameSink,
3627 corr: u64,
3628 module_id: &str,
3629 ) -> PendingRouteBindRelay {
3630 forwarding
3631 .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
3632 .unwrap()
3633 }
3634
3635 #[tokio::test]
3641 async fn pending_route_open_completes_behind_queued_data_frames() {
3642 assert_eq!(
3643 crate::server::MAX_PENDING_ROUTE_OPENS_PER_CONNECTION,
3644 8,
3645 "the per-connection pending route.open limit is its own constant"
3646 );
3647 let (forwarding, module_connection, _endpoint, client, _unused_sink, _unused_rx) =
3648 route_fixture("open-behind-data");
3649 let (sink, mut client_rx) = crate::server::connection_egress();
3650 const DATA_FRAMES: usize = 1_000;
3651 let data = |corr: u64| {
3652 Frame::build(
3653 FrameType::StreamData,
3654 Flags::new(false, Priority::Interactive, false),
3655 9,
3656 1,
3657 corr,
3658 vec![b'x'; 200],
3659 )
3660 .unwrap()
3661 };
3662 for corr in 0..DATA_FRAMES as u64 {
3663 sink.try_send(data(corr)).unwrap();
3664 }
3665 let data_bytes = DATA_FRAMES * (subc_protocol::HEADER_LEN + 200);
3666 assert_eq!(sink.backlog().queued_bytes, data_bytes);
3667
3668 let pending = tokio::time::timeout(
3669 Duration::from_secs(5),
3670 forwarding.begin_route_bind_relay_for(
3671 client,
3672 sink.clone(),
3673 subc_protocol::PROTOCOL_VERSION,
3674 4_242,
3675 "open-behind-data",
3676 Principal::Direct,
3677 None,
3678 Instant::now() + Duration::from_secs(60),
3679 ),
3680 )
3681 .await
3682 .expect("reserving the route.open slot must not wait behind data frames")
3683 .unwrap();
3684 forwarding
3685 .complete_pending_relay(
3686 module_connection,
3687 pending.corr,
3688 RouteBindRelayOutcome::Accepted,
3689 )
3690 .unwrap();
3691
3692 let backlog = sink.backlog();
3693 assert_eq!(backlog.queued_frames, DATA_FRAMES + 1);
3694 assert!(
3695 backlog.queued_bytes > data_bytes,
3696 "the route.open response must be counted in queued bytes: {backlog:?}"
3697 );
3698 for corr in 0..DATA_FRAMES as u64 {
3699 assert_eq!(client_rx.recv().await.unwrap().header.corr, corr);
3700 }
3701 let open = client_rx.recv().await.unwrap();
3702 assert_eq!(open.header.corr, 4_242);
3703 assert_eq!(open.header.ty, FrameType::Response);
3704 drop(open);
3705 assert_eq!(sink.backlog().queued_bytes, 0);
3706 assert_eq!(sink.backlog().queued_frames, 0);
3707 }
3708
3709 #[tokio::test]
3714 async fn drain_holdouts_count_held_requests_and_name_the_connection() {
3715 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3716 route_fixture("holdouts");
3717 let mut bound = |corr| {
3718 let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
3719 forwarding
3720 .complete_pending_relay(
3721 module_connection,
3722 route.corr,
3723 RouteBindRelayOutcome::Accepted,
3724 )
3725 .unwrap();
3726 client_rx.try_recv().unwrap();
3727 match forwarding
3728 .lookup_data_route(client, route.client_channel, route.client_epoch)
3729 .unwrap()
3730 {
3731 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3732 other => panic!("expected live route, got {other:?}"),
3733 }
3734 };
3735 let holding = bound(61);
3736 let _idle = bound(62);
3737 holding.flow.acquire_tagged(7, false).await.unwrap();
3738 holding.flow.acquire_tagged(2, false).await.unwrap();
3739 holding.flow.acquire_tagged(3, true).await.unwrap();
3740 forwarding
3741 .begin_module_drain("holdouts", RouteCloseReason::Restart)
3742 .unwrap();
3743
3744 let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
3745 assert_eq!(
3746 holdouts,
3747 DrainHoldouts {
3748 requests: 2,
3749 routes: 1,
3750 total_routes: 2,
3751 top_connections: vec![(client.get(), 2)],
3752 held: vec![(holding.module_channel, 2), (holding.module_channel, 7)],
3755 }
3756 );
3757 }
3758
3759 #[test]
3760 fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
3761 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3762 route_fixture("census");
3763 let pending = begin_test_route(&forwarding, client, sink, 1, "census");
3764 forwarding
3765 .complete_pending_relay(
3766 module_connection,
3767 pending.corr,
3768 RouteBindRelayOutcome::Accepted,
3769 )
3770 .unwrap();
3771
3772 let routes = forwarding.endpoint_routes(endpoint).unwrap();
3773 assert_eq!(routes.len(), 1);
3774 assert!(matches!(routes[0].principal, Principal::Direct));
3775 assert_eq!(routes[0].goodbye_target.connection_id, client);
3776 assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
3777 assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
3778 assert!(!routes[0].draining);
3779
3780 forwarding
3781 .begin_module_drain("census", RouteCloseReason::Restart)
3782 .unwrap();
3783 let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
3784 assert_eq!(draining_routes.len(), 1);
3785 assert!(draining_routes[0].draining);
3786 }
3787
3788 #[test]
3789 fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
3790 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
3791 let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
3792 assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
3793 forwarding
3794 .abort_pending_relay(
3795 first.endpoint,
3796 first.corr,
3797 RouteBindRelayOutcome::ModuleGone("abort".into()),
3798 )
3799 .unwrap();
3800 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3801 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3802
3803 let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
3804 assert_eq!(second.client_channel, first.client_channel);
3805 assert_eq!(second.module_channel, first.module_channel);
3806 assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
3807 }
3808
3809 #[test]
3810 fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
3811 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3812 route_fixture("epoch-release");
3813 let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
3814 forwarding
3815 .complete_pending_relay(
3816 module_connection,
3817 first.corr,
3818 RouteBindRelayOutcome::Accepted,
3819 )
3820 .unwrap();
3821 assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
3822 assert!(matches!(
3823 forwarding
3824 .release_client_route(client, first.client_channel, first.client_epoch)
3825 .unwrap(),
3826 RouteRelease::Removed(_)
3827 ));
3828 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3829 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3830
3831 let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
3832 forwarding
3833 .complete_pending_relay(
3834 module_connection,
3835 second.corr,
3836 RouteBindRelayOutcome::Accepted,
3837 )
3838 .unwrap();
3839 assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
3840 assert!(matches!(
3841 forwarding
3842 .release_client_route(client, second.client_channel, first.client_epoch)
3843 .unwrap(),
3844 RouteRelease::Stale
3845 ));
3846 assert!(!forwarding
3847 .cache_status(
3848 endpoint,
3849 second.module_channel,
3850 first.module_epoch,
3851 "stale".into(),
3852 )
3853 .unwrap());
3854 assert!(forwarding
3855 .cache_status(
3856 endpoint,
3857 second.module_channel,
3858 second.module_epoch,
3859 "current".into(),
3860 )
3861 .unwrap());
3862 match forwarding
3863 .route_poll_snapshot(client, second.client_channel, second.client_epoch)
3864 .unwrap()
3865 {
3866 RoutePollSnapshot::Bound { status, .. } => {
3867 assert_eq!(status.as_deref(), Some("current"));
3868 }
3869 RoutePollSnapshot::Absent => panic!("successor binding was removed"),
3870 }
3871 let counters = forwarding.counters().snapshot();
3872 assert_eq!(counters["route_released_epoch_fenced"], 1);
3873 assert_eq!(counters["route_release_stale_skipped"], 1);
3874 }
3875
3876 #[test]
3877 fn max_epoch_reservation_retires_only_that_slot() {
3878 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
3879 forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
3880 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
3881 let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
3882 assert_eq!(
3883 (final_use.client_channel, final_use.client_epoch),
3884 (7, u32::MAX)
3885 );
3886 assert_eq!(
3887 (final_use.module_channel, final_use.module_epoch),
3888 (9, u32::MAX)
3889 );
3890 forwarding
3891 .abort_pending_relay(
3892 endpoint,
3893 final_use.corr,
3894 RouteBindRelayOutcome::ModuleGone("abort".into()),
3895 )
3896 .unwrap();
3897 forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
3898 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
3899 let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
3900 assert_ne!(next.client_channel, 7);
3901 assert_ne!(next.module_channel, 9);
3902 assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
3903 }
3904
3905 #[test]
3906 fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
3907 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3908 route_fixture("corr-shared");
3909 let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
3910 assert_eq!(bind.corr, 1);
3911 forwarding
3912 .abort_pending_relay(
3913 endpoint,
3914 bind.corr,
3915 RouteBindRelayOutcome::ModuleGone("abort".into()),
3916 )
3917 .unwrap();
3918 let rpc = forwarding
3919 .begin_module_control_rpc_for(
3920 "corr-shared",
3921 "health.check",
3922 Instant::now() - Duration::from_millis(1),
3923 )
3924 .unwrap();
3925 assert_eq!(rpc.corr, 2);
3926 assert_eq!(
3927 forwarding
3928 .complete_module_control_rpc(
3929 module_connection,
3930 rpc.corr,
3931 Some("health.check"),
3932 ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
3933 status: subc_protocol::session::HealthStatus::Ok,
3934 detail: None,
3935 metrics: None,
3936 }),
3937 )
3938 .unwrap(),
3939 ModuleControlRpcCompletion::Settled
3940 );
3941 assert!(matches!(
3942 rpc.receiver.blocking_recv().unwrap(),
3943 ModuleControlRpcOutcome::DeadlineElapsed
3944 ));
3945 }
3946
3947 #[tokio::test(start_paused = true)]
3948 async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
3949 let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
3950 let probe_started_at = Instant::now();
3951 let rpc = forwarding
3952 .begin_health_probe_rpc_for(
3953 "tombstone-ttl",
3954 "health.check",
3955 probe_started_at,
3956 probe_started_at + Duration::from_secs(5),
3957 )
3958 .unwrap();
3959 assert!(forwarding
3960 .tombstone_health_probe_rpc(endpoint, rpc.corr)
3961 .unwrap());
3962 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
3963
3964 tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
3965 tokio::task::yield_now().await;
3966
3967 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
3968 }
3969
3970 #[test]
3971 fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
3972 let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
3973 let mut close = forwarding.register_connection_close(endpoint.connection_id);
3974 forwarding.inject_control_corr(endpoint, u64::MAX);
3975 let final_rpc = forwarding
3976 .begin_module_control_rpc_for(
3977 "corr-max",
3978 "health.check",
3979 Instant::now() + Duration::from_secs(1),
3980 )
3981 .unwrap();
3982 assert_eq!(final_rpc.corr, u64::MAX);
3983 forwarding
3984 .cancel_module_control_rpc(endpoint, final_rpc.corr)
3985 .unwrap();
3986 assert!(matches!(
3987 forwarding.begin_module_control_rpc_for(
3988 "corr-max",
3989 "health.check",
3990 Instant::now() + Duration::from_secs(1),
3991 ),
3992 Err(ForwardingError::RelayCorrelationExhausted)
3993 ));
3994 assert!(close.try_recv().is_ok());
3995 }
3996
3997 #[test]
3998 fn publication_epoch_controls_delivery_failure_escalation() {
3999 fn setup_successor(
4000 commit_successor: Option<bool>,
4001 ) -> (ForwardingTable, ConnectionId, u16, u32) {
4002 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4003 route_fixture("escalation");
4004 let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
4005 forwarding
4006 .complete_pending_relay(
4007 module_connection,
4008 first.corr,
4009 RouteBindRelayOutcome::Accepted,
4010 )
4011 .unwrap();
4012 client_rx.try_recv().unwrap();
4013 assert!(matches!(
4014 forwarding
4015 .release_client_route(client, first.client_channel, first.client_epoch)
4016 .unwrap(),
4017 RouteRelease::Removed(_)
4018 ));
4019 if let Some(commit_successor) = commit_successor {
4020 forwarding.inject_client_slot_epoch(
4021 client,
4022 first.client_channel,
4023 first.client_epoch,
4024 );
4025 forwarding.inject_module_slot_epoch(
4026 endpoint,
4027 first.module_channel,
4028 first.module_epoch,
4029 );
4030 let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
4031 if commit_successor {
4032 forwarding
4033 .complete_pending_relay(
4034 module_connection,
4035 successor.corr,
4036 RouteBindRelayOutcome::Accepted,
4037 )
4038 .unwrap();
4039 client_rx.try_recv().unwrap();
4040 } else {
4041 forwarding
4042 .abort_pending_relay(
4043 endpoint,
4044 successor.corr,
4045 RouteBindRelayOutcome::ModuleGone("abort".into()),
4046 )
4047 .unwrap();
4048 }
4049 }
4050 (forwarding, client, first.client_channel, first.client_epoch)
4051 }
4052
4053 let probe_sink = FrameSink::new(mpsc::channel(1).0);
4054 let (no_successor, client, channel, epoch) = setup_successor(None);
4055 let mut close = no_successor.register_connection_close(client);
4056 assert!(no_successor
4057 .escalate_client_delivery_failure(
4058 client,
4059 channel,
4060 epoch,
4061 CloseReason::new("delivery", "failed"),
4062 UndeliveredFrame {
4063 module_id: None,
4064 sink: &probe_sink,
4065 },
4066 )
4067 .unwrap());
4068 assert!(close.try_recv().is_ok());
4069
4070 let (aborted, client, channel, epoch) = setup_successor(Some(false));
4071 let mut close = aborted.register_connection_close(client);
4072 assert!(aborted
4073 .escalate_client_delivery_failure(
4074 client,
4075 channel,
4076 epoch,
4077 CloseReason::new("delivery", "failed"),
4078 UndeliveredFrame {
4079 module_id: None,
4080 sink: &probe_sink,
4081 },
4082 )
4083 .unwrap());
4084 assert!(close.try_recv().is_ok());
4085
4086 let (published, client, channel, epoch) = setup_successor(Some(true));
4087 let mut close = published.register_connection_close(client);
4088 assert!(!published
4089 .escalate_client_delivery_failure(
4090 client,
4091 channel,
4092 epoch,
4093 CloseReason::new("delivery", "stale failure"),
4094 UndeliveredFrame {
4095 module_id: None,
4096 sink: &probe_sink,
4097 },
4098 )
4099 .unwrap());
4100 assert!(close.try_recv().is_err());
4101 }
4102
4103 #[test]
4104 fn route_concentration_separates_client_count_from_routes_per_client() {
4105 let (forwarding, module_connection, _, client, sink, _client_rx) =
4109 route_fixture("concentration");
4110 assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
4111
4112 for corr in [70_u64, 71] {
4113 let pending =
4114 begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
4115 forwarding
4116 .complete_pending_relay(
4117 module_connection,
4118 pending.corr,
4119 RouteBindRelayOutcome::Accepted,
4120 )
4121 .unwrap();
4122 }
4123
4124 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4126 assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
4127 }
4128
4129 #[test]
4130 fn cleanup_and_accepted_resolution_have_one_lock_winner() {
4131 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4132 route_fixture("cleanup-race");
4133 let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
4134 forwarding
4135 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4136 .unwrap();
4137 let released = forwarding.cleanup_connection(client).unwrap();
4138 assert_eq!(released.len(), 1);
4139 let completion = forwarding
4140 .complete_pending_relay(
4141 module_connection,
4142 pending.corr,
4143 RouteBindRelayOutcome::Accepted,
4144 )
4145 .unwrap();
4146 assert!(!completion.settled);
4147 assert!(client_rx.try_recv().is_err());
4148 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4149
4150 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4151 route_fixture("accepted-race");
4152 let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
4153 forwarding
4154 .complete_pending_relay(
4155 module_connection,
4156 pending.corr,
4157 RouteBindRelayOutcome::Accepted,
4158 )
4159 .unwrap();
4160 assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
4161 let released = forwarding.cleanup_connection(client).unwrap();
4162 assert_eq!(released.len(), 1);
4163 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4164 }
4165
4166 #[test]
4167 fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
4168 let (forwarding, module_connection, _, client, sink, mut client_rx) =
4169 route_fixture("drain-gap");
4170 let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
4171 forwarding
4172 .complete_pending_relay(
4173 module_connection,
4174 live.corr,
4175 RouteBindRelayOutcome::Accepted,
4176 )
4177 .unwrap();
4178 client_rx.try_recv().unwrap();
4179 let binding = match forwarding
4180 .lookup_data_route(client, live.client_channel, live.client_epoch)
4181 .unwrap()
4182 {
4183 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
4184 other => panic!("expected live route, got {other:?}"),
4185 };
4186
4187 let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
4188 forwarding
4189 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4190 .unwrap();
4191 let control_rpc = forwarding
4192 .begin_module_control_rpc_for(
4193 "drain-gap",
4194 "health.check",
4195 Instant::now() + Duration::from_secs(1),
4196 )
4197 .unwrap();
4198 let target = forwarding
4199 .begin_module_drain("drain-gap", RouteCloseReason::Reload)
4200 .unwrap()
4201 .unwrap();
4202 assert!(matches!(
4203 control_rpc.receiver.blocking_recv().unwrap(),
4204 ModuleControlRpcOutcome::ModuleGone(_)
4205 ));
4206 assert_eq!(target.abandoned_bindings.len(), 1);
4207 assert!(binding.flow.sem.is_closed());
4208 assert!(
4209 !forwarding
4210 .complete_pending_relay(
4211 module_connection,
4212 pending.corr,
4213 RouteBindRelayOutcome::Accepted,
4214 )
4215 .unwrap()
4216 .settled
4217 );
4218 assert!(matches!(
4219 forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
4220 Err(ForwardingError::ModuleReloading { .. })
4221 ));
4222 let released = forwarding
4223 .release_module_endpoint_routes(target.endpoint)
4224 .unwrap();
4225 assert_eq!(released.len(), 1);
4226 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4227 }
4228
4229 #[test]
4237 fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
4238 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
4239 route_fixture("closing-client");
4240
4241 let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
4244 forwarding
4245 .complete_pending_relay(
4246 module_connection,
4247 live.corr,
4248 RouteBindRelayOutcome::Accepted,
4249 )
4250 .unwrap();
4251 client_rx.try_recv().unwrap();
4252
4253 let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
4255 forwarding
4256 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
4257 .unwrap();
4258
4259 assert!(forwarding
4261 .escalate_client_delivery_failure(
4262 client,
4263 live.client_channel,
4264 live.client_epoch,
4265 CloseReason::new(
4266 "module_to_client_delivery_failed",
4267 "client egress refused a module frame",
4268 ),
4269 UndeliveredFrame {
4270 module_id: None,
4271 sink: &sink,
4272 },
4273 )
4274 .unwrap());
4275 assert!(!sink.is_closed());
4276
4277 let completion = forwarding
4278 .complete_pending_relay(
4279 module_connection,
4280 pending.corr,
4281 RouteBindRelayOutcome::Accepted,
4282 )
4283 .expect("a closing client must not turn a module's ack into an error");
4284
4285 assert!(completion.settled);
4286 let abandoned = completion
4287 .abandoned
4288 .expect("the module must be told to drop the binding it just created");
4289 assert_eq!(abandoned.connection_id, module_connection);
4290 assert_eq!(abandoned.channel, pending.module_channel);
4291 assert_eq!(abandoned.epoch, pending.module_epoch);
4292 assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
4293 assert!(matches!(
4294 pending.receiver.blocking_recv().unwrap(),
4295 RouteBindRelayOutcome::ModuleGone(_)
4296 ));
4297 assert!(client_rx.try_recv().is_err());
4300 assert_eq!(forwarding.active_binding_count().unwrap(), 1);
4301
4302 assert!(forwarding
4305 .has_live_module_connection("closing-client")
4306 .unwrap());
4307 let cotenant = ConnectionId::new(201);
4308 let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
4309 let cotenant_route = begin_test_route(
4310 &forwarding,
4311 cotenant,
4312 FrameSink::new(cotenant_tx),
4313 62,
4314 "closing-client",
4315 );
4316 assert_eq!(cotenant_route.endpoint, endpoint);
4317 forwarding
4318 .complete_pending_relay(
4319 module_connection,
4320 cotenant_route.corr,
4321 RouteBindRelayOutcome::Accepted,
4322 )
4323 .unwrap();
4324 assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
4325 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
4326 }
4327
4328 #[test]
4329 fn pending_route_permit_is_released_on_rejection_and_abort() {
4330 let forwarding = ForwardingTable::default();
4331 let module_connection = ConnectionId::new(300);
4332 let client = ConnectionId::new(301);
4333 let (module_tx, _module_rx) = mpsc::channel(1);
4334 let endpoint = forwarding
4335 .register_module_connection(
4336 module_connection,
4337 "permit".into(),
4338 2,
4339 Concurrency::ModuleManaged,
4340 FrameSink::new(module_tx),
4341 )
4342 .unwrap();
4343 let (client_tx, mut client_rx) = mpsc::channel(1);
4344 let sink = FrameSink::new(client_tx);
4345 let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
4346 assert!(sink.try_send(test_ping(999)).is_err());
4347 forwarding
4348 .complete_pending_relay(
4349 module_connection,
4350 rejected.corr,
4351 RouteBindRelayOutcome::Rejected(ErrorBody {
4352 code: "no".into(),
4353 message: "rejected".into(),
4354 detail: None,
4355 }),
4356 )
4357 .unwrap();
4358 sink.try_send(test_ping(1000)).unwrap();
4359 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
4360
4361 let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
4362 assert!(sink.try_send(test_ping(1001)).is_err());
4363 forwarding
4364 .abort_pending_relay(
4365 endpoint,
4366 aborted.corr,
4367 RouteBindRelayOutcome::ModuleGone("abort".into()),
4368 )
4369 .unwrap();
4370 sink.try_send(test_ping(1002)).unwrap();
4371 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
4372
4373 let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
4374 forwarding
4375 .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
4376 .unwrap();
4377 drop(client_rx);
4378 let completion = forwarding
4379 .complete_pending_relay(
4380 module_connection,
4381 receiver_closed.corr,
4382 RouteBindRelayOutcome::Accepted,
4383 )
4384 .unwrap();
4385 assert!(completion.abandoned.is_some());
4386 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
4387 }
4388
4389 #[test]
4395 fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
4396 let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
4397 route_fixture("closing-set-leak");
4398
4399 const CONNECTIONS: u64 = 32;
4400 for index in 0..CONNECTIONS {
4401 let client = ConnectionId::new(1000 + index);
4402 let (client_tx, _client_rx) = mpsc::channel(8);
4403 let route = begin_test_route(
4404 &forwarding,
4405 client,
4406 FrameSink::new(client_tx),
4407 index + 1,
4408 "closing-set-leak",
4409 );
4410 forwarding
4411 .complete_pending_relay(
4412 module_connection,
4413 route.corr,
4414 RouteBindRelayOutcome::Accepted,
4415 )
4416 .unwrap();
4417 forwarding.cleanup_connection(client).unwrap();
4418 }
4419 forwarding.cleanup_connection(module_connection).unwrap();
4420
4421 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
4422 }
4423
4424 #[test]
4431 fn closing_connection_is_refused_new_work_until_cleanup_completes() {
4432 let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
4433 route_fixture("closing-gate");
4434
4435 let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
4438 forwarding
4439 .complete_pending_relay(
4440 module_connection,
4441 live.corr,
4442 RouteBindRelayOutcome::Accepted,
4443 )
4444 .unwrap();
4445 client_rx.try_recv().unwrap();
4446
4447 assert!(forwarding
4450 .escalate_client_delivery_failure(
4451 client,
4452 live.client_channel,
4453 live.client_epoch,
4454 CloseReason::new(
4455 "module_to_client_delivery_failed",
4456 "client egress refused a module frame",
4457 ),
4458 UndeliveredFrame {
4459 module_id: None,
4460 sink: &sink,
4461 },
4462 )
4463 .unwrap());
4464 assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
4465
4466 assert!(matches!(
4468 forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
4469 Err(ForwardingError::ConnectionClosing { connection_id })
4470 if connection_id == client
4471 ));
4472 let (late_tx, _late_rx) = mpsc::channel(1);
4474 assert!(matches!(
4475 forwarding.register_module_connection(
4476 client,
4477 "late-module".into(),
4478 2,
4479 Concurrency::ModuleManaged,
4480 FrameSink::new(late_tx),
4481 ),
4482 Err(ForwardingError::ConnectionClosing { connection_id })
4483 if connection_id == client
4484 ));
4485
4486 forwarding.cleanup_connection(client).unwrap();
4490 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
4491 }
4492}
4493
4494#[cfg(test)]
4497mod swap_slot_tests {
4498 use std::time::Duration;
4499
4500 use super::*;
4501 use tokio::sync::mpsc;
4502
4503 const MODULE_ID: &str = "swapped";
4504
4505 struct SwapFixture {
4506 forwarding: ForwardingTable,
4507 incumbent_connection: ConnectionId,
4508 incumbent: ModuleEndpointId,
4509 candidate_connection: ConnectionId,
4510 candidate: ModuleEndpointId,
4511 _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
4512 }
4513
4514 fn swap_fixture() -> SwapFixture {
4515 let forwarding = ForwardingTable::default();
4516 let incumbent_connection = ConnectionId::new(100);
4517 let candidate_connection = ConnectionId::new(110);
4518 let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
4519 let incumbent = forwarding
4520 .register_module_connection(
4521 incumbent_connection,
4522 MODULE_ID.to_string(),
4523 2,
4524 Concurrency::ModuleManaged,
4525 FrameSink::new(incumbent_tx),
4526 )
4527 .unwrap();
4528 let (candidate_tx, candidate_rx) = mpsc::channel(8);
4529 let candidate = forwarding
4530 .register_candidate_module_connection(
4531 candidate_connection,
4532 MODULE_ID.to_string(),
4533 2,
4534 Concurrency::ModuleManaged,
4535 FrameSink::new(candidate_tx),
4536 )
4537 .unwrap();
4538 SwapFixture {
4539 forwarding,
4540 incumbent_connection,
4541 incumbent,
4542 candidate_connection,
4543 candidate,
4544 _module_rxs: vec![incumbent_rx, candidate_rx],
4545 }
4546 }
4547
4548 fn client(
4549 raw: u64,
4550 ) -> (
4551 ConnectionId,
4552 FrameSink,
4553 mpsc::Receiver<crate::router::OutboundFrame>,
4554 ) {
4555 let (tx, rx) = mpsc::channel(8);
4556 (ConnectionId::new(raw), FrameSink::new(tx), rx)
4557 }
4558
4559 fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
4560 forwarding
4561 .read_inner()
4562 .unwrap()
4563 .client_to_module
4564 .values()
4565 .map(|route| route.module_endpoint)
4566 .collect()
4567 }
4568
4569 #[test]
4570 fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
4571 let fixture = swap_fixture();
4572 let forwarding = &fixture.forwarding;
4573 assert_ne!(fixture.incumbent, fixture.candidate);
4574
4575 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4577 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4578 let (client_connection, client_sink, _client_rx) = client(200);
4579 let pending = forwarding
4580 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4581 .unwrap();
4582 assert_eq!(pending.endpoint, fixture.incumbent);
4583 let rpc = forwarding
4584 .begin_module_control_rpc_for(
4585 MODULE_ID,
4586 "health.check",
4587 Instant::now() + Duration::from_secs(1),
4588 )
4589 .unwrap();
4590 assert_eq!(rpc.endpoint, fixture.incumbent);
4591 let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
4592 assert_eq!(census.len(), 1, "the census lists one endpoint per id");
4593
4594 assert_eq!(
4596 forwarding
4597 .module_endpoint_for_connection(fixture.candidate_connection)
4598 .unwrap(),
4599 Some(fixture.candidate)
4600 );
4601 assert_eq!(
4602 forwarding
4603 .module_id_for_connection(fixture.candidate_connection)
4604 .unwrap()
4605 .as_deref(),
4606 Some(MODULE_ID)
4607 );
4608
4609 let (other_tx, _other_rx) = mpsc::channel(1);
4611 assert_eq!(
4612 forwarding.register_candidate_module_connection(
4613 ConnectionId::new(120),
4614 MODULE_ID.to_string(),
4615 2,
4616 Concurrency::ModuleManaged,
4617 FrameSink::new(other_tx),
4618 ),
4619 Err(ForwardingError::CandidateSlotOccupied {
4620 module_id: MODULE_ID.to_string()
4621 })
4622 );
4623 }
4624
4625 #[test]
4629 fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
4630 let fixture = swap_fixture();
4631 let forwarding = &fixture.forwarding;
4632 let (early_client, early_sink, _early_rx) = client(200);
4633 let mut early = forwarding
4634 .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
4635 .unwrap();
4636 assert_eq!(early.endpoint, fixture.incumbent);
4637 assert!(forwarding
4638 .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
4639 .unwrap());
4640
4641 let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4642 assert_eq!(
4643 cutover,
4644 ForwardingCutover {
4645 promoted: fixture.candidate,
4646 incumbent: Some(fixture.incumbent),
4647 }
4648 );
4649
4650 let (late_client, late_sink, _late_rx) = client(201);
4652 let late = forwarding
4653 .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
4654 .unwrap();
4655 assert_eq!(
4656 late.endpoint, fixture.candidate,
4657 "a route.open after cutover was reserved on the incumbent"
4658 );
4659
4660 let completion = forwarding
4662 .complete_pending_relay(
4663 fixture.incumbent_connection,
4664 early.corr,
4665 RouteBindRelayOutcome::Accepted,
4666 )
4667 .expect("a superseded endpoint's ack is not an error on its connection");
4668 assert!(completion.settled);
4669 assert!(
4670 !committed_endpoints(forwarding).contains(&fixture.incumbent),
4671 "a relay reserved before cutover committed a route on the incumbent"
4672 );
4673 let goodbye = completion
4674 .abandoned
4675 .expect("the incumbent is told to drop the binding it just created");
4676 assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
4677 assert_eq!(goodbye.channel, early.module_channel);
4678 assert_eq!(goodbye.epoch, early.module_epoch);
4679 assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
4680 match early.receiver.try_recv() {
4681 Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
4682 other => panic!("expected a retryable module_reloading answer, got {other:?}"),
4683 }
4684 assert!(matches!(
4685 forwarding
4686 .lookup_data_route(early_client, early.client_channel, early.client_epoch)
4687 .unwrap(),
4688 DataRoute::Client(DataRouteState::Absent)
4689 ));
4690
4691 assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4694 forwarding
4695 .complete_pending_relay(
4696 fixture.candidate_connection,
4697 late.corr,
4698 RouteBindRelayOutcome::Accepted,
4699 )
4700 .unwrap();
4701 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4702 assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
4703 }
4704
4705 #[test]
4706 fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
4707 let fixture = swap_fixture();
4708 let forwarding = &fixture.forwarding;
4709 let (bound_client, bound_sink, _bound_rx) = client(200);
4711 let bound = forwarding
4712 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4713 .unwrap();
4714 forwarding
4715 .complete_pending_relay(
4716 fixture.incumbent_connection,
4717 bound.corr,
4718 RouteBindRelayOutcome::Accepted,
4719 )
4720 .unwrap();
4721 let (pending_client, pending_sink, _pending_rx) = client(201);
4722 let mut in_flight = forwarding
4723 .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
4724 .unwrap();
4725 forwarding
4726 .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
4727 .unwrap();
4728
4729 let incumbent = forwarding
4730 .cutover_candidate(MODULE_ID)
4731 .unwrap()
4732 .unwrap()
4733 .incumbent
4734 .unwrap();
4735 let target = forwarding
4736 .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
4737 .unwrap()
4738 .expect("the superseded incumbent is still registered");
4739
4740 assert_eq!(target.endpoint, fixture.incumbent);
4741 assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
4742 assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
4743 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4744 assert_eq!(target.abandoned_bindings.len(), 1);
4745 assert_eq!(
4746 target.abandoned_bindings[0].channel,
4747 in_flight.module_channel
4748 );
4749 assert!(matches!(
4750 in_flight.receiver.try_recv(),
4751 Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
4752 ));
4753 assert_eq!(
4754 forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
4755 1,
4756 "the incumbent's bound route stays until its drain finishes"
4757 );
4758
4759 let (next_client, next_sink, _next_rx) = client(202);
4760 let next = forwarding
4761 .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
4762 .expect("the promoted candidate keeps accepting routes");
4763 assert_eq!(next.endpoint, fixture.candidate);
4764 }
4765
4766 #[test]
4772 fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
4773 let forwarding = ForwardingTable::default();
4774 let first_connection = ConnectionId::new(70);
4775 let (first_tx, _first_rx) = mpsc::channel(8);
4776 forwarding
4777 .register_module_connection(
4778 first_connection,
4779 MODULE_ID.to_string(),
4780 2,
4781 Concurrency::ModuleManaged,
4782 FrameSink::new(first_tx),
4783 )
4784 .unwrap();
4785 let (client_connection, client_sink, _client_rx) = client(200);
4786 let mut pending = forwarding
4787 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4788 .unwrap();
4789 let (second_tx, _second_rx) = mpsc::channel(8);
4790 forwarding
4791 .register_module_connection(
4792 ConnectionId::new(80),
4793 MODULE_ID.to_string(),
4794 2,
4795 Concurrency::ModuleManaged,
4796 FrameSink::new(second_tx),
4797 )
4798 .unwrap();
4799
4800 assert_eq!(
4801 forwarding
4802 .complete_pending_relay(
4803 first_connection,
4804 pending.corr,
4805 RouteBindRelayOutcome::Accepted
4806 )
4807 .unwrap_err(),
4808 ForwardingError::StaleModuleEndpoint
4809 );
4810 assert!(committed_endpoints(&forwarding).is_empty());
4811 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4812 assert!(matches!(
4813 pending.receiver.try_recv(),
4814 Err(oneshot::error::TryRecvError::Closed)
4815 ));
4816 }
4817
4818 #[test]
4819 fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
4820 let fixture = swap_fixture();
4822 let forwarding = &fixture.forwarding;
4823 assert!(forwarding
4824 .cleanup_connection(fixture.candidate_connection)
4825 .unwrap()
4826 .is_empty());
4827 assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
4828 let (client_connection, client_sink, _client_rx) = client(200);
4829 assert_eq!(
4830 forwarding
4831 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4832 .unwrap()
4833 .endpoint,
4834 fixture.incumbent
4835 );
4836
4837 let fixture = swap_fixture();
4840 let forwarding = &fixture.forwarding;
4841 let (bound_client, bound_sink, _bound_rx) = client(200);
4842 let bound = forwarding
4843 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4844 .unwrap();
4845 forwarding
4846 .complete_pending_relay(
4847 fixture.incumbent_connection,
4848 bound.corr,
4849 RouteBindRelayOutcome::Accepted,
4850 )
4851 .unwrap();
4852 forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4853 let released = forwarding
4854 .cleanup_connection(fixture.incumbent_connection)
4855 .unwrap();
4856 assert_eq!(released.len(), 1);
4857 assert_eq!(released[0].connection_id, bound_client);
4858 assert!(forwarding
4859 .read_inner()
4860 .unwrap()
4861 .superseded_endpoints
4862 .is_empty());
4863 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4864 let (next_client, next_sink, _next_rx) = client(201);
4865 assert_eq!(
4866 forwarding
4867 .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
4868 .unwrap()
4869 .endpoint,
4870 fixture.candidate
4871 );
4872 }
4873}