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