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