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