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}
1765
1766impl ForwardingTable {
1767 pub(crate) fn endpoint_drain_holdouts(
1769 &self,
1770 endpoint: ModuleEndpointId,
1771 ) -> Result<DrainHoldouts, ForwardingError> {
1772 let inner = self.read_inner()?;
1773 let mut holdouts = DrainHoldouts::default();
1774 let mut by_connection: HashMap<u64, usize> = HashMap::new();
1775 for (key, route) in &inner.client_to_module {
1776 if route.module_endpoint != endpoint {
1777 continue;
1778 }
1779 holdouts.total_routes += 1;
1780 let held = route.flow.drain_in_flight();
1781 if held == 0 {
1782 continue;
1783 }
1784 holdouts.requests += held;
1785 holdouts.routes += 1;
1786 *by_connection.entry(key.connection_id.get()).or_default() += held;
1787 }
1788 let mut connections = by_connection.into_iter().collect::<Vec<_>>();
1789 connections.sort_by(|left, right| right.1.cmp(&left.1).then(left.0.cmp(&right.0)));
1790 connections.truncate(3);
1791 holdouts.top_connections = connections;
1792 Ok(holdouts)
1793 }
1794
1795 pub(crate) fn endpoint_in_flight_count(
1796 &self,
1797 endpoint: ModuleEndpointId,
1798 ) -> Result<usize, ForwardingError> {
1799 let inner = self.read_inner()?;
1800 Ok(inner
1801 .client_to_module
1802 .values()
1803 .filter(|route| route.module_endpoint == endpoint)
1804 .map(|route| route.flow.drain_in_flight())
1805 .sum())
1806 }
1807
1808 pub(crate) fn endpoint_is_draining(
1809 &self,
1810 endpoint: ModuleEndpointId,
1811 ) -> Result<bool, ForwardingError> {
1812 Ok(self
1813 .read_inner()?
1814 .draining_endpoints
1815 .contains_key(&endpoint))
1816 }
1817
1818 pub(crate) fn module_is_draining(&self, module_id: &str) -> Result<bool, ForwardingError> {
1819 let inner = self.read_inner()?;
1820 Ok(inner
1821 .modules_by_id
1822 .get(module_id)
1823 .is_some_and(|module| inner.draining_endpoints.contains_key(&module.endpoint)))
1824 }
1825
1826 pub(crate) fn release_module_endpoint_routes(
1827 &self,
1828 endpoint: ModuleEndpointId,
1829 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1830 let mut inner = self.write_inner()?;
1831 let routes = inner
1832 .module_to_client
1833 .iter()
1834 .filter(|(module_key, _)| module_key.endpoint == endpoint)
1835 .map(|(module_key, route)| (*module_key, route.module_epoch))
1836 .collect::<Vec<_>>();
1837 let mut released = Vec::with_capacity(routes.len());
1838 for (module_key, epoch) in routes {
1839 if let RouteRelease::Removed(target) =
1840 release_module_route_locked(&mut inner, module_key, epoch)
1841 {
1842 released.push(target);
1843 }
1844 }
1845 Ok(released)
1846 }
1847
1848 pub(crate) fn endpoint_routes(
1854 &self,
1855 endpoint: ModuleEndpointId,
1856 ) -> Result<Vec<EndpointRoute>, ForwardingError> {
1857 let inner = self.read_inner()?;
1858 Ok(endpoint_routes_locked(&inner, endpoint))
1859 }
1860
1861 pub(crate) fn route_census(
1863 &self,
1864 module_id: Option<&str>,
1865 ) -> Result<Vec<(String, Vec<EndpointRoute>)>, ForwardingError> {
1866 let inner = self.read_inner()?;
1867 let mut endpoints = inner
1868 .modules_by_id
1869 .iter()
1870 .filter(|(id, _)| module_id.is_none_or(|requested| requested == id.as_str()))
1871 .map(|(id, module)| (id.clone(), module.endpoint))
1872 .collect::<Vec<_>>();
1873 endpoints.sort_by(|left, right| left.0.cmp(&right.0));
1874 Ok(endpoints
1875 .into_iter()
1876 .map(|(id, endpoint)| (id, endpoint_routes_locked(&inner, endpoint)))
1877 .collect())
1878 }
1879
1880 pub(crate) fn live_roots(
1882 &self,
1883 module_id: &str,
1884 ) -> Result<ModuleControlResponseToModule, ForwardingError> {
1885 let inner = self.read_inner()?;
1886 let endpoint = inner
1887 .modules_by_id
1888 .get(module_id)
1889 .map(|module| module.endpoint);
1890 let mut roots = BTreeMap::new();
1891 let mut unknown_root_bindings = 0;
1892 let mut total_bindings = 0;
1893 if let Some(endpoint) = endpoint {
1894 for binding in inner
1895 .module_to_client
1896 .values()
1897 .filter(|binding| binding.module_endpoint == endpoint)
1898 {
1899 total_bindings += 1;
1900 if let Some(root) = &binding.project_root {
1901 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1902 entry.0 += 1;
1903 } else {
1904 unknown_root_bindings += 1;
1905 }
1906 }
1907 for pending in inner
1908 .pending_relays
1909 .values()
1910 .filter(|pending| pending.reservation.module_key.endpoint == endpoint)
1911 {
1912 total_bindings += 1;
1913 if let Some(root) = &pending.reservation.project_root {
1914 let entry = roots.entry(root.as_path().to_path_buf()).or_insert((0, 0));
1915 entry.1 += 1;
1916 } else {
1917 unknown_root_bindings += 1;
1918 }
1919 }
1920 }
1921 Ok(ModuleControlResponseToModule::LiveRoots {
1922 roots: roots
1923 .into_iter()
1924 .map(|(project_root, (bound, pending))| LiveRoot {
1925 project_root,
1926 bound,
1927 pending,
1928 })
1929 .collect(),
1930 unknown_root_bindings,
1931 total_bindings,
1932 })
1933 }
1934
1935 pub(crate) fn connection_has_client_routes(
1941 &self,
1942 connection_id: ConnectionId,
1943 ) -> Result<bool, ForwardingError> {
1944 let inner = self.read_inner()?;
1945 let has = inner
1946 .client_to_module
1947 .keys()
1948 .any(|key| key.connection_id == connection_id)
1949 || inner
1950 .reserved_client
1951 .keys()
1952 .any(|key| key.connection_id == connection_id);
1953 Ok(has)
1954 }
1955
1956 pub(crate) fn cleanup_connection(
1957 &self,
1958 connection_id: ConnectionId,
1959 ) -> Result<Vec<GoodbyeTarget>, ForwardingError> {
1960 let mut inner = self.write_inner()?;
1961 inner.closing_connections.insert(connection_id);
1962 let released = if let Some(endpoint) = inner.endpoint_by_connection.remove(&connection_id) {
1963 remove_module_connection_locked(&mut inner, endpoint)
1964 } else {
1965 Self::cleanup_client_connection_locked(&mut inner, connection_id)
1966 };
1967 inner.closing_connections.remove(&connection_id);
1976 Ok(released)
1977 }
1978
1979 fn cleanup_client_connection_locked(
1980 inner: &mut ForwardingInner,
1981 connection_id: ConnectionId,
1982 ) -> Vec<GoodbyeTarget> {
1983 let routes = inner
1984 .client_to_module
1985 .iter()
1986 .filter(|(key, _)| key.connection_id == connection_id)
1987 .map(|(key, route)| (*key, route.client_epoch))
1988 .collect::<Vec<_>>();
1989 let mut released = Vec::with_capacity(routes.len());
1990 for (client_key, epoch) in routes {
1991 if let RouteRelease::Removed(target) =
1992 release_client_route_locked(inner, client_key, epoch)
1993 {
1994 released.push(target);
1995 }
1996 }
1997
1998 let pending_keys = inner
1999 .pending_relays
2000 .iter()
2001 .filter(|(_, pending)| pending.reservation.client_key.connection_id == connection_id)
2002 .map(|(key, _)| *key)
2003 .collect::<Vec<_>>();
2004 for key in pending_keys {
2005 let Some(pending) = inner.pending_relays.remove(&key) else {
2006 continue;
2007 };
2008 release_reserved_route_locked(
2009 inner,
2010 pending.reservation.client_key,
2011 pending.reservation.module_key,
2012 );
2013 if pending.relay_enqueued {
2014 if let Some(target) = abandoned_route_target(inner, &pending.reservation) {
2015 released.push(target);
2016 }
2017 }
2018 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2019 "client connection closed during route.bind relay".to_string(),
2020 ));
2021 }
2022
2023 let orphaned = inner
2024 .reserved_client
2025 .iter()
2026 .filter(|(key, _)| key.connection_id == connection_id)
2027 .map(|(client, module)| (*client, *module))
2028 .collect::<Vec<_>>();
2029 for (client_key, module_key) in orphaned {
2030 release_reserved_route_locked(inner, client_key, module_key);
2031 }
2032 inner.next_client_channel.remove(&connection_id);
2033 inner
2034 .client_slot_epochs
2035 .retain(|key, _| key.connection_id != connection_id);
2036 inner
2037 .last_published_epoch
2038 .retain(|key, _| key.connection_id != connection_id);
2039 inner
2040 .status
2041 .retain(|(key, _), _| key.connection_id != connection_id);
2042
2043 released
2044 }
2045
2046 pub(crate) fn escalate_client_delivery_failure(
2047 &self,
2048 connection_id: ConnectionId,
2049 channel: u16,
2050 expected_epoch: u32,
2051 reason: CloseReason,
2052 ) -> Result<bool, ForwardingError> {
2053 let should_close = {
2054 let mut inner = self.write_inner()?;
2055 let key = ClientRouteKey {
2056 connection_id,
2057 channel,
2058 };
2059 if inner.last_published_epoch.get(&key).copied() != Some(expected_epoch) {
2060 false
2061 } else {
2062 inner.closing_connections.insert(connection_id);
2063 true
2064 }
2065 };
2066 if should_close {
2067 self.request_connection_close(connection_id, reason);
2068 }
2069 Ok(should_close)
2070 }
2071
2072 fn record_route_release(&self, release: &RouteRelease) {
2073 match release {
2074 RouteRelease::Removed(_) => self.counters.increment_route_released_epoch_fenced(),
2075 RouteRelease::Stale => self.counters.increment_route_release_stale_skipped(),
2076 RouteRelease::Absent => {}
2077 }
2078 }
2079
2080 fn read_inner(&self) -> Result<RwLockReadGuard<'_, ForwardingInner>, ForwardingError> {
2081 self.inner.read().map_err(|_| ForwardingError::Poisoned)
2082 }
2083
2084 fn write_inner(&self) -> Result<RwLockWriteGuard<'_, ForwardingInner>, ForwardingError> {
2085 self.inner.write().map_err(|_| ForwardingError::Poisoned)
2086 }
2087
2088 fn lock_close_registry(
2089 &self,
2090 ) -> MutexGuard<'_, HashMap<ConnectionId, oneshot::Sender<CloseReason>>> {
2091 self.close_registry
2092 .lock()
2093 .unwrap_or_else(|poisoned| poisoned.into_inner())
2094 }
2095}
2096
2097impl ForwardingInner {
2098 fn allocate_route_slots(
2099 &mut self,
2100 connection_id: ConnectionId,
2101 endpoint: ModuleEndpointId,
2102 ) -> Result<(u16, u32, u16, u32), ForwardingError> {
2103 let client_start = *self.next_client_channel.entry(connection_id).or_insert(1);
2104 let mut client_channel = client_start;
2105 let client_channel = loop {
2106 let key = ClientRouteKey {
2107 connection_id,
2108 channel: client_channel,
2109 };
2110 let eligible = !self.client_to_module.contains_key(&key)
2111 && !self.reserved_client.contains_key(&key)
2112 && self.client_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2113 if eligible {
2114 break client_channel;
2115 }
2116 client_channel = next_channel(client_channel);
2117 if client_channel == client_start {
2118 return Err(ForwardingError::ClientRouteChannelExhausted { connection_id });
2119 }
2120 };
2121
2122 let module_start = *self.next_module_channel.entry(endpoint).or_insert(1);
2123 let mut module_channel = module_start;
2124 let module_channel = loop {
2125 let key = ModuleRouteKey {
2126 endpoint,
2127 channel: module_channel,
2128 };
2129 let eligible = !self.module_to_client.contains_key(&key)
2130 && !self.reserved_module.contains_key(&key)
2131 && self.module_slot_epochs.get(&key).copied().unwrap_or(0) < u32::MAX;
2132 if eligible {
2133 break module_channel;
2134 }
2135 module_channel = next_channel(module_channel);
2136 if module_channel == module_start {
2137 return Err(ForwardingError::ModuleRouteChannelExhausted { endpoint });
2138 }
2139 };
2140
2141 let client_key = ClientRouteKey {
2142 connection_id,
2143 channel: client_channel,
2144 };
2145 let module_key = ModuleRouteKey {
2146 endpoint,
2147 channel: module_channel,
2148 };
2149 let client_epoch = self
2150 .client_slot_epochs
2151 .get(&client_key)
2152 .copied()
2153 .unwrap_or(0)
2154 + 1;
2155 let module_epoch = self
2156 .module_slot_epochs
2157 .get(&module_key)
2158 .copied()
2159 .unwrap_or(0)
2160 + 1;
2161 self.client_slot_epochs.insert(client_key, client_epoch);
2162 self.module_slot_epochs.insert(module_key, module_epoch);
2163 self.next_client_channel
2164 .insert(connection_id, next_channel(client_channel));
2165 self.next_module_channel
2166 .insert(endpoint, next_channel(module_channel));
2167 Ok((client_channel, client_epoch, module_channel, module_epoch))
2168 }
2169
2170 fn allocate_control_corr(
2171 &mut self,
2172 endpoint: ModuleEndpointId,
2173 ) -> Result<u64, ForwardingError> {
2174 let candidate = self.next_control_corr.get(&endpoint).copied().unwrap_or(1);
2175 if candidate == 0 {
2176 self.closing_connections.insert(endpoint.connection_id);
2177 return Err(ForwardingError::RelayCorrelationExhausted);
2178 }
2179 self.next_control_corr.insert(
2180 endpoint,
2181 if candidate == u64::MAX {
2182 0
2183 } else {
2184 candidate + 1
2185 },
2186 );
2187 Ok(candidate)
2188 }
2189}
2190
2191fn next_channel(channel: u16) -> u16 {
2192 let next = channel.wrapping_add(1);
2193 if next == 0 {
2194 1
2195 } else {
2196 next
2197 }
2198}
2199
2200fn endpoint_routes_locked(
2201 inner: &ForwardingInner,
2202 endpoint: ModuleEndpointId,
2203) -> Vec<EndpointRoute> {
2204 let drain_reason = inner.draining_endpoints.get(&endpoint).copied();
2205 let draining = drain_reason.is_some();
2206 let mut routes = inner
2207 .module_to_client
2208 .iter()
2209 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2210 .map(|(_, route)| EndpointRoute {
2211 goodbye_target: GoodbyeTarget {
2212 connection_id: route.client_connection_id,
2213 sink: route.client_sink.clone(),
2214 negotiated_ver: route.client_negotiated_ver,
2215 channel: route.client_channel,
2216 epoch: route.client_epoch,
2217 kind: GoodbyeTargetKind::Client,
2218 module_id: None,
2219 },
2220 principal: route.principal.clone(),
2221 bound_at: route.bound_at,
2222 draining,
2223 drain_reason,
2224 })
2225 .collect::<Vec<_>>();
2226 routes.sort_by_key(|route| {
2227 (
2228 route.goodbye_target.connection_id.get(),
2229 route.goodbye_target.channel,
2230 route.goodbye_target.epoch,
2231 )
2232 });
2233 routes
2234}
2235
2236fn release_reserved_route_locked(
2237 inner: &mut ForwardingInner,
2238 client_key: ClientRouteKey,
2239 module_key: ModuleRouteKey,
2240) {
2241 if inner.reserved_client.get(&client_key).copied() == Some(module_key) {
2242 inner.reserved_client.remove(&client_key);
2243 }
2244 if inner.reserved_module.get(&module_key).copied() == Some(client_key) {
2245 inner.reserved_module.remove(&module_key);
2246 }
2247 inner.status.retain(|(key, _), _| *key != client_key);
2248}
2249
2250fn release_client_route_locked(
2251 inner: &mut ForwardingInner,
2252 client_key: ClientRouteKey,
2253 expected_epoch: u32,
2254) -> RouteRelease {
2255 let Some(route) = inner.client_to_module.get(&client_key) else {
2256 return RouteRelease::Absent;
2257 };
2258 if route.client_epoch != expected_epoch {
2259 return RouteRelease::Stale;
2260 }
2261 let route = inner
2262 .client_to_module
2263 .remove(&client_key)
2264 .expect("route checked under the same forwarding lock");
2265 route.flow.close();
2266 inner.module_to_client.remove(&ModuleRouteKey {
2267 endpoint: route.module_endpoint,
2268 channel: route.module_channel,
2269 });
2270 inner.status.remove(&(client_key, expected_epoch));
2271 RouteRelease::Removed(GoodbyeTarget {
2272 connection_id: route.module_endpoint.connection_id,
2273 sink: route.module_sink.clone(),
2274 negotiated_ver: route.module_negotiated_ver,
2275 channel: route.module_channel,
2276 epoch: route.module_epoch,
2277 kind: GoodbyeTargetKind::Module,
2278 module_id: Some(route.module_id.clone()),
2279 })
2280}
2281
2282fn release_module_route_locked(
2283 inner: &mut ForwardingInner,
2284 module_key: ModuleRouteKey,
2285 expected_epoch: u32,
2286) -> RouteRelease {
2287 let Some(route) = inner.module_to_client.get(&module_key) else {
2288 return RouteRelease::Absent;
2289 };
2290 if route.module_epoch != expected_epoch {
2291 return RouteRelease::Stale;
2292 }
2293 let route = inner
2294 .module_to_client
2295 .remove(&module_key)
2296 .expect("route checked under the same forwarding lock");
2297 route.flow.close();
2298 let client_key = ClientRouteKey {
2299 connection_id: route.client_connection_id,
2300 channel: route.client_channel,
2301 };
2302 inner.client_to_module.remove(&client_key);
2303 inner.status.remove(&(client_key, route.client_epoch));
2304 RouteRelease::Removed(GoodbyeTarget {
2305 connection_id: route.client_connection_id,
2306 sink: route.client_sink.clone(),
2307 negotiated_ver: route.client_negotiated_ver,
2308 channel: route.client_channel,
2309 epoch: route.client_epoch,
2310 kind: GoodbyeTargetKind::Client,
2311 module_id: None,
2312 })
2313}
2314
2315fn commit_route_locked(
2316 inner: &mut ForwardingInner,
2317 pending: PendingRouteBindRelayEntry,
2318) -> Result<Option<GoodbyeTarget>, ForwardingError> {
2319 let reservation = pending.reservation;
2320 if inner
2321 .closing_connections
2322 .contains(&reservation.client_key.connection_id)
2323 {
2324 return Err(ForwardingError::ConnectionClosing {
2325 connection_id: reservation.client_key.connection_id,
2326 });
2327 }
2328 let module_id = inner
2329 .module_id_by_endpoint
2330 .get(&reservation.module_key.endpoint)
2331 .cloned()
2332 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2333 if inner
2334 .draining_endpoints
2335 .contains_key(&reservation.module_key.endpoint)
2336 {
2337 return Err(ForwardingError::ModuleReloading { module_id });
2338 }
2339 if inner.reserved_client.remove(&reservation.client_key) != Some(reservation.module_key)
2340 || inner.reserved_module.remove(&reservation.module_key) != Some(reservation.client_key)
2341 {
2342 return Err(ForwardingError::UnknownReservation {
2343 client_channel: reservation.client_key.channel,
2344 module_channel: reservation.module_key.channel,
2345 });
2346 }
2347 let module = inner
2348 .modules_by_id
2349 .get(&module_id)
2350 .filter(|module| module.endpoint == reservation.module_key.endpoint)
2351 .cloned()
2352 .ok_or(ForwardingError::StaleModuleEndpoint)?;
2353 let binding = Arc::new(RouteBinding {
2354 client_connection_id: reservation.client_key.connection_id,
2355 client_sink: pending.client_sink,
2356 client_negotiated_ver: pending.client_negotiated_ver,
2357 client_channel: reservation.client_key.channel,
2358 client_epoch: reservation.client_epoch,
2359 module_id,
2360 module_endpoint: reservation.module_key.endpoint,
2361 module_sink: module.sink,
2362 module_negotiated_ver: module.negotiated_ver,
2363 module_channel: reservation.module_key.channel,
2364 module_epoch: reservation.module_epoch,
2365 principal: pending.principal,
2366 project_root: reservation.project_root.clone(),
2367 bound_at: Instant::now(),
2368 flow: Arc::new(ChannelFlow::new(window_for(&module.concurrency))),
2369 });
2370 inner
2371 .client_to_module
2372 .insert(reservation.client_key, Arc::clone(&binding));
2373 inner
2374 .module_to_client
2375 .insert(reservation.module_key, binding);
2376 let previous_published = inner
2377 .last_published_epoch
2378 .insert(reservation.client_key, reservation.client_epoch);
2379
2380 let client_sender = pending.client_permit.send(crate::router::OutboundFrame {
2385 frame: pending.route_open_frame,
2386 enqueued_at: std::time::Instant::now(),
2387 flushed: None,
2388 });
2389 if client_sender.is_closed() {
2390 let abandoned = pending
2391 .relay_enqueued
2392 .then(|| abandoned_route_target(inner, &reservation))
2393 .flatten();
2394 if let Some(route) = inner.client_to_module.remove(&reservation.client_key) {
2395 route.flow.close();
2396 }
2397 inner.module_to_client.remove(&reservation.module_key);
2398 inner
2399 .status
2400 .remove(&(reservation.client_key, reservation.client_epoch));
2401 match previous_published {
2402 Some(epoch) => {
2403 inner
2404 .last_published_epoch
2405 .insert(reservation.client_key, epoch);
2406 }
2407 None => {
2408 inner.last_published_epoch.remove(&reservation.client_key);
2409 }
2410 }
2411 let _ = pending.sender.send(RouteBindRelayOutcome::ModuleGone(
2412 "client egress closed during route publication".to_string(),
2413 ));
2414 return Ok(abandoned);
2415 }
2416
2417 let _ = pending.sender.send(RouteBindRelayOutcome::Accepted);
2418 Ok(None)
2419}
2420
2421fn module_connection_for_endpoint_locked(
2429 inner: &ForwardingInner,
2430 endpoint: ModuleEndpointId,
2431) -> Option<&ModuleConnection> {
2432 let module_id = inner.module_id_by_endpoint.get(&endpoint)?;
2433 inner
2434 .modules_by_id
2435 .get(module_id)
2436 .filter(|module| module.endpoint == endpoint)
2437 .or_else(|| {
2438 inner
2439 .candidates_by_id
2440 .get(module_id)
2441 .filter(|module| module.endpoint == endpoint)
2442 })
2443 .or_else(|| inner.superseded_endpoints.get(&endpoint))
2444}
2445
2446fn abandoned_route_target(
2447 inner: &ForwardingInner,
2448 reservation: &RouteReservation,
2449) -> Option<GoodbyeTarget> {
2450 let module_id = inner
2451 .module_id_by_endpoint
2452 .get(&reservation.module_key.endpoint)?;
2453 let module = module_connection_for_endpoint_locked(inner, reservation.module_key.endpoint)?;
2454 (module.endpoint == reservation.module_key.endpoint).then(|| GoodbyeTarget {
2455 connection_id: module.endpoint.connection_id,
2456 sink: module.sink.clone(),
2457 negotiated_ver: module.negotiated_ver,
2458 channel: reservation.module_key.channel,
2459 epoch: reservation.module_epoch,
2460 kind: GoodbyeTargetKind::Module,
2461 module_id: Some(module_id.clone()),
2462 })
2463}
2464
2465fn remove_module_connection_locked(
2466 inner: &mut ForwardingInner,
2467 endpoint: ModuleEndpointId,
2468) -> Vec<GoodbyeTarget> {
2469 inner.draining_endpoints.remove(&endpoint);
2470 let module_id = inner.module_id_by_endpoint.remove(&endpoint);
2471 if let Some(module_id) = module_id.as_ref() {
2472 if inner
2473 .modules_by_id
2474 .get(module_id)
2475 .is_some_and(|module| module.endpoint == endpoint)
2476 {
2477 inner.modules_by_id.remove(module_id);
2478 }
2479 if inner
2480 .candidates_by_id
2481 .get(module_id)
2482 .is_some_and(|module| module.endpoint == endpoint)
2483 {
2484 inner.candidates_by_id.remove(module_id);
2485 }
2486 }
2487 inner.superseded_endpoints.remove(&endpoint);
2488 inner.endpoint_by_connection.remove(&endpoint.connection_id);
2489 inner.next_module_channel.remove(&endpoint);
2490 inner.next_control_corr.remove(&endpoint);
2491 inner
2492 .health_probe_tombstones
2493 .retain(|(pending_endpoint, _), _| *pending_endpoint != endpoint);
2494 inner
2495 .module_slot_epochs
2496 .retain(|key, _| key.endpoint != endpoint);
2497 let reserved_module_keys: Vec<ModuleRouteKey> = inner
2498 .reserved_module
2499 .keys()
2500 .filter(|module_key| module_key.endpoint == endpoint)
2501 .copied()
2502 .collect();
2503 for module_key in reserved_module_keys {
2504 if let Some(client_key) = inner.reserved_module.get(&module_key).copied() {
2505 release_reserved_route_locked(inner, client_key, module_key);
2506 }
2507 }
2508
2509 let pending_keys: Vec<_> = inner
2510 .pending_relays
2511 .keys()
2512 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2513 .copied()
2514 .collect();
2515 let pending: Vec<_> = pending_keys
2516 .into_iter()
2517 .filter_map(|key| inner.pending_relays.remove(&key))
2518 .collect();
2519 for pending in pending {
2520 let module_label = module_id.as_deref().unwrap_or("unknown");
2521 let _ = pending
2522 .sender
2523 .send(RouteBindRelayOutcome::ModuleGone(format!(
2524 "module '{module_label}' connection closed during route.bind relay"
2525 )));
2526 }
2527
2528 let pending_control_keys: Vec<_> = inner
2529 .pending_control_rpcs
2530 .keys()
2531 .filter(|(pending_endpoint, _)| *pending_endpoint == endpoint)
2532 .copied()
2533 .collect();
2534 let pending_control: Vec<_> = pending_control_keys
2535 .into_iter()
2536 .filter_map(|key| inner.pending_control_rpcs.remove(&key))
2537 .collect();
2538 for pending in pending_control {
2539 let module_label = module_id.as_deref().unwrap_or("unknown");
2540 let _ = pending
2541 .sender
2542 .send(ModuleControlRpcOutcome::ModuleGone(format!(
2543 "module '{module_label}' connection closed during module-control RPC"
2544 )));
2545 }
2546
2547 let module_routes = inner
2548 .module_to_client
2549 .iter()
2550 .filter(|(module_key, _)| module_key.endpoint == endpoint)
2551 .map(|(module_key, route)| (*module_key, route.module_epoch))
2552 .collect::<Vec<_>>();
2553 let mut released = Vec::with_capacity(module_routes.len());
2554 for (module_key, epoch) in module_routes {
2555 if let RouteRelease::Removed(target) = release_module_route_locked(inner, module_key, epoch)
2556 {
2557 released.push(target);
2558 }
2559 }
2560 released
2561}
2562
2563#[derive(Debug, Clone, Copy)]
2564struct RequestCredit {
2565 subscription: bool,
2566 excluded_from_drain: bool,
2567}
2568
2569#[derive(Debug, Default)]
2570struct CreditLedger {
2571 by_corr: HashMap<u64, Vec<RequestCredit>>,
2572}
2573
2574impl CreditLedger {
2575 fn acquire(&mut self, corr: u64, subscription: bool) {
2576 self.by_corr.entry(corr).or_default().push(RequestCredit {
2577 subscription,
2578 excluded_from_drain: false,
2579 });
2580 }
2581
2582 fn release(&mut self, corr: u64) -> bool {
2583 let Some(credits) = self.by_corr.get_mut(&corr) else {
2584 return false;
2585 };
2586 let released = credits.pop().is_some();
2587 if credits.is_empty() {
2588 self.by_corr.remove(&corr);
2589 }
2590 released
2591 }
2592
2593 fn capture_subscription_exclusions(&mut self) -> u32 {
2594 let mut excluded = 0u32;
2595 for credit in self.by_corr.values_mut().flatten() {
2596 if credit.subscription && !credit.excluded_from_drain {
2597 credit.excluded_from_drain = true;
2598 excluded = excluded.saturating_add(1);
2599 }
2600 }
2601 excluded
2602 }
2603
2604 #[cfg(test)]
2605 fn in_flight(&self) -> usize {
2606 self.by_corr.values().map(Vec::len).sum()
2607 }
2608
2609 fn drain_in_flight(&self) -> usize {
2610 self.by_corr
2611 .values()
2612 .flatten()
2613 .filter(|credit| !credit.excluded_from_drain)
2614 .count()
2615 }
2616}
2617
2618#[derive(Debug, Default)]
2619struct ChannelFlowState {
2620 closed: bool,
2621 credits: CreditLedger,
2622}
2623
2624#[derive(Debug)]
2626pub(crate) struct ChannelFlow {
2627 sem: Semaphore,
2628 window: usize,
2629 state: Mutex<ChannelFlowState>,
2630}
2631
2632impl ChannelFlow {
2633 pub(crate) fn new(window: usize) -> Self {
2634 debug_assert!(window > 0, "flow-control window must be non-zero");
2635 Self {
2636 sem: Semaphore::new(window),
2637 window,
2638 state: Mutex::new(ChannelFlowState::default()),
2639 }
2640 }
2641
2642 #[cfg(test)]
2643 pub(crate) async fn acquire(&self) -> Result<(), ChannelFlowClosed> {
2644 self.acquire_tagged(0, false).await
2645 }
2646
2647 pub(crate) async fn acquire_tagged(
2648 &self,
2649 corr: u64,
2650 subscription: bool,
2651 ) -> Result<(), ChannelFlowClosed> {
2652 let permit = self.sem.acquire().await.map_err(|_| ChannelFlowClosed)?;
2653 let mut state = self
2654 .state
2655 .lock()
2656 .unwrap_or_else(|poisoned| poisoned.into_inner());
2657 if state.closed {
2658 return Err(ChannelFlowClosed);
2659 }
2660 state.credits.acquire(corr, subscription);
2661 permit.forget();
2662 Ok(())
2663 }
2664
2665 #[cfg(test)]
2666 pub(crate) fn release(&self) {
2667 self.release_corr(0);
2668 }
2669
2670 pub(crate) fn release_corr(&self, corr: u64) {
2671 let released = self
2672 .state
2673 .lock()
2674 .unwrap_or_else(|poisoned| poisoned.into_inner())
2675 .credits
2676 .release(corr);
2677 if !released {
2678 warn!(
2682 window = self.window,
2683 available = self.sem.available_permits(),
2684 "flow-control over-release ignored"
2685 );
2686 return;
2687 }
2688 if !self.sem.is_closed() {
2689 self.sem.add_permits(1);
2690 }
2691 }
2692
2693 #[cfg(test)]
2694 pub(crate) fn in_flight(&self) -> usize {
2695 self.state
2696 .lock()
2697 .unwrap_or_else(|poisoned| poisoned.into_inner())
2698 .credits
2699 .in_flight()
2700 }
2701
2702 pub(crate) fn drain_in_flight(&self) -> usize {
2703 self.state
2704 .lock()
2705 .unwrap_or_else(|poisoned| poisoned.into_inner())
2706 .credits
2707 .drain_in_flight()
2708 }
2709
2710 #[cfg(test)]
2711 pub(crate) fn available_permits(&self) -> usize {
2712 self.sem.available_permits()
2713 }
2714
2715 pub(crate) fn begin_drain(&self) -> u32 {
2716 let mut state = self
2717 .state
2718 .lock()
2719 .unwrap_or_else(|poisoned| poisoned.into_inner());
2720 state.closed = true;
2721 self.sem.close();
2722 state.credits.capture_subscription_exclusions()
2723 }
2724
2725 pub(crate) fn close(&self) {
2726 self.state
2727 .lock()
2728 .unwrap_or_else(|poisoned| poisoned.into_inner())
2729 .closed = true;
2730 self.sem.close();
2731 }
2732}
2733
2734#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2735pub(crate) struct ChannelFlowClosed;
2736
2737impl fmt::Display for ChannelFlowClosed {
2738 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2739 write!(f, "flow-control window closed")
2740 }
2741}
2742
2743impl Error for ChannelFlowClosed {}
2744
2745fn window_for(concurrency: &Concurrency) -> usize {
2746 match concurrency {
2747 Concurrency::Serial => 1,
2748 Concurrency::ModuleManaged => DEFAULT_MODULE_MANAGED_WINDOW,
2749 Concurrency::StatelessParallel => STATELESS_PARALLEL_WINDOW,
2750 }
2751}
2752
2753#[derive(Debug, Clone, PartialEq, Eq)]
2754pub enum ForwardingError {
2755 NoModuleConnection,
2756 ModuleReloading {
2757 module_id: String,
2758 },
2759 StaleModuleEndpoint,
2760 UnknownReservation {
2761 client_channel: u16,
2762 module_channel: u16,
2763 },
2764 ClientRouteChannelExhausted {
2765 connection_id: ConnectionId,
2766 },
2767 ModuleRouteChannelExhausted {
2768 endpoint: ModuleEndpointId,
2769 },
2770 RelayCorrelationExhausted,
2771 ConnectionClosing {
2772 connection_id: ConnectionId,
2773 },
2774 ClientEgressClosed {
2775 connection_id: ConnectionId,
2776 },
2777 RouteOpenBuild(String),
2778 CandidateSlotOccupied {
2780 module_id: String,
2781 },
2782 Poisoned,
2783}
2784
2785impl fmt::Display for ForwardingError {
2786 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2787 match self {
2788 Self::NoModuleConnection => write!(f, "no module connection is registered"),
2789 Self::ModuleReloading { module_id } => {
2790 write!(f, "module_id '{module_id}' is reloading")
2791 }
2792 Self::StaleModuleEndpoint => write!(f, "module connection generation is stale"),
2793 Self::UnknownReservation {
2794 client_channel,
2795 module_channel,
2796 } => write!(
2797 f,
2798 "route reservation client channel {client_channel} / module channel {module_channel} was not found"
2799 ),
2800 Self::ClientRouteChannelExhausted { connection_id } => write!(
2801 f,
2802 "no client route channels are available for connection {}",
2803 connection_id.get()
2804 ),
2805 Self::ModuleRouteChannelExhausted { endpoint } => write!(
2806 f,
2807 "no module route channels are available for endpoint generation {} on connection {}",
2808 endpoint.generation,
2809 endpoint.connection_id.get()
2810 ),
2811 Self::RelayCorrelationExhausted => {
2812 write!(f, "module control correlation ids are exhausted")
2813 }
2814 Self::ConnectionClosing { connection_id } => write!(
2815 f,
2816 "connection {} is closing and cannot accept route allocation",
2817 connection_id.get()
2818 ),
2819 Self::ClientEgressClosed { connection_id } => write!(
2820 f,
2821 "client connection {} egress is closed",
2822 connection_id.get()
2823 ),
2824 Self::RouteOpenBuild(message) => {
2825 write!(f, "failed to prebuild route.open response: {message}")
2826 }
2827 Self::CandidateSlotOccupied { module_id } => write!(
2828 f,
2829 "module_id '{module_id}' already has a swap candidate registered"
2830 ),
2831 Self::Poisoned => write!(f, "forwarding table lock was poisoned"),
2832 }
2833 }
2834}
2835
2836impl Error for ForwardingError {}
2837
2838#[cfg(test)]
2839mod tests {
2840 use std::time::Duration;
2841
2842 use super::*;
2843 use tokio::sync::mpsc;
2844
2845 #[test]
2846 fn ordinary_long_running_request_is_not_excluded_from_drain() {
2847 let mut ledger = CreditLedger::default();
2848 ledger.acquire(1, false);
2849
2850 assert_eq!(ledger.capture_subscription_exclusions(), 0);
2851 assert_eq!(ledger.drain_in_flight(), 1);
2852 }
2853
2854 #[test]
2855 fn bit_set_subscription_is_excluded_and_counted() {
2856 let mut ledger = CreditLedger::default();
2857 ledger.acquire(1, true);
2858
2859 assert_eq!(ledger.capture_subscription_exclusions(), 1);
2860 assert_eq!(ledger.drain_in_flight(), 0);
2861 }
2862
2863 #[test]
2864 fn subscription_opened_after_drain_snapshot_is_not_excluded() {
2865 let mut ledger = CreditLedger::default();
2866 ledger.acquire(1, true);
2867 assert_eq!(ledger.capture_subscription_exclusions(), 1);
2868
2869 ledger.acquire(2, true);
2870
2871 assert_eq!(ledger.drain_in_flight(), 1);
2872 }
2873
2874 #[test]
2875 fn drain_with_no_subscriptions_reports_zero_excluded() {
2876 let mut ledger = CreditLedger::default();
2877 assert_eq!(ledger.capture_subscription_exclusions(), 0);
2878 }
2879
2880 #[test]
2881 fn multi_provider_route_limit_reports_per_client_exhaustion_without_affecting_second_client() {
2882 let forwarding = ForwardingTable::default();
2883 let module_connection = ConnectionId::new(10);
2884 let exhausted_client = ConnectionId::new(20);
2885 let second_client = ConnectionId::new(30);
2886 let (module_tx, _module_rx) = mpsc::channel(1);
2887 let endpoint = forwarding
2888 .register_module_connection(
2889 module_connection,
2890 "route-limit-provider".to_string(),
2891 1,
2892 Concurrency::ModuleManaged,
2893 FrameSink::new(module_tx),
2894 )
2895 .unwrap();
2896
2897 {
2898 let mut inner = forwarding.inner.write().unwrap();
2899 for channel in 1..=u16::MAX {
2900 inner.reserved_client.insert(
2901 ClientRouteKey {
2902 connection_id: exhausted_client,
2903 channel,
2904 },
2905 ModuleRouteKey {
2906 endpoint,
2907 channel: 1,
2908 },
2909 );
2910 }
2911 }
2912
2913 let (exhausted_tx, _exhausted_rx) = mpsc::channel(1);
2914 let err = forwarding
2915 .begin_route_bind_relay_for_test(
2916 exhausted_client,
2917 FrameSink::new(exhausted_tx),
2918 1,
2919 "route-limit-provider",
2920 )
2921 .unwrap_err();
2922 assert!(matches!(
2923 err,
2924 ForwardingError::ClientRouteChannelExhausted { connection_id }
2925 if connection_id == exhausted_client
2926 ));
2927
2928 let (second_tx, _second_rx) = mpsc::channel(1);
2929 let pending = forwarding
2930 .begin_route_bind_relay_for_test(
2931 second_client,
2932 FrameSink::new(second_tx),
2933 2,
2934 "route-limit-provider",
2935 )
2936 .unwrap();
2937 assert_eq!(pending.client_channel, 1);
2938 }
2939
2940 #[test]
2941 fn released_module_channels_are_reused_after_wrap_without_slot_leak() {
2942 let forwarding = ForwardingTable::default();
2943 let module_connection = ConnectionId::new(40);
2944 let client = ConnectionId::new(50);
2945 let (module_tx, _module_rx) = mpsc::channel(1);
2946 forwarding
2947 .register_module_connection(
2948 module_connection,
2949 "slot-reuse-provider".to_string(),
2950 1,
2951 Concurrency::ModuleManaged,
2952 FrameSink::new(module_tx),
2953 )
2954 .unwrap();
2955
2956 let (client_tx, _client_rx) = mpsc::channel(1);
2957 let client_sink = FrameSink::new(client_tx);
2958 let mut wrapped_channel = None;
2959 for index in 0..=usize::from(u16::MAX) {
2960 let pending = forwarding
2961 .begin_route_bind_relay_for_test(
2962 client,
2963 client_sink.clone(),
2964 index as u64 + 1,
2965 "slot-reuse-provider",
2966 )
2967 .unwrap();
2968 if index == usize::from(u16::MAX) {
2969 wrapped_channel = Some(pending.module_channel);
2970 }
2971 forwarding
2972 .abort_pending_relay(
2973 pending.endpoint,
2974 pending.corr,
2975 RouteBindRelayOutcome::ModuleGone("test abort".to_string()),
2976 )
2977 .unwrap();
2978 }
2979
2980 assert_eq!(wrapped_channel, Some(1));
2981 }
2982
2983 #[test]
2984 fn cleanup_connection_prunes_stale_next_client_channel_cursor() {
2985 let forwarding = ForwardingTable::default();
2986 let client = ConnectionId::new(60);
2987 forwarding
2988 .inner
2989 .write()
2990 .unwrap()
2991 .next_client_channel
2992 .insert(client, 41);
2993
2994 let released = forwarding.cleanup_connection(client).unwrap();
2995
2996 assert!(released.is_empty());
2997 assert!(!forwarding
2998 .inner
2999 .read()
3000 .unwrap()
3001 .next_client_channel
3002 .contains_key(&client));
3003 }
3004
3005 #[test]
3006 fn stale_module_cleanup_preserves_fast_reconnect_successor() {
3007 let forwarding = ForwardingTable::default();
3008 let module_id = "fast-reconnect-provider";
3009 let first_connection = ConnectionId::new(70);
3010 let second_connection = ConnectionId::new(80);
3011 let (first_tx, _first_rx) = mpsc::channel(1);
3012 let first_endpoint = forwarding
3013 .register_module_connection(
3014 first_connection,
3015 module_id.to_string(),
3016 1,
3017 Concurrency::ModuleManaged,
3018 FrameSink::new(first_tx),
3019 )
3020 .unwrap();
3021 let (second_tx, _second_rx) = mpsc::channel(1);
3022 let second_endpoint = forwarding
3023 .register_module_connection(
3024 second_connection,
3025 module_id.to_string(),
3026 1,
3027 Concurrency::ModuleManaged,
3028 FrameSink::new(second_tx),
3029 )
3030 .unwrap();
3031 assert_ne!(first_endpoint, second_endpoint);
3032
3033 let released = forwarding.cleanup_connection(first_connection).unwrap();
3034
3035 assert!(released.is_empty());
3036 assert_eq!(
3037 forwarding
3038 .inner
3039 .read()
3040 .unwrap()
3041 .modules_by_id
3042 .get(module_id)
3043 .map(|module| module.endpoint),
3044 Some(second_endpoint)
3045 );
3046 assert!(forwarding.has_live_module_connection(module_id).unwrap());
3047 let control_rpc = forwarding
3048 .begin_module_control_rpc_for(
3049 module_id,
3050 "health.check",
3051 Instant::now() + Duration::from_secs(1),
3052 )
3053 .unwrap();
3054 assert_eq!(control_rpc.endpoint, second_endpoint);
3055 }
3056
3057 fn route_fixture(
3058 module_id: &str,
3059 ) -> (
3060 ForwardingTable,
3061 ConnectionId,
3062 ModuleEndpointId,
3063 ConnectionId,
3064 FrameSink,
3065 mpsc::Receiver<crate::router::OutboundFrame>,
3066 ) {
3067 let forwarding = ForwardingTable::default();
3068 let module_connection = ConnectionId::new(100);
3069 let client_connection = ConnectionId::new(200);
3070 let (module_tx, _module_rx) = mpsc::channel(8);
3071 let endpoint = forwarding
3072 .register_module_connection(
3073 module_connection,
3074 module_id.to_string(),
3075 2,
3076 Concurrency::ModuleManaged,
3077 FrameSink::new(module_tx),
3078 )
3079 .unwrap();
3080 let (client_tx, client_rx) = mpsc::channel(8);
3081 (
3082 forwarding,
3083 module_connection,
3084 endpoint,
3085 client_connection,
3086 FrameSink::new(client_tx),
3087 client_rx,
3088 )
3089 }
3090
3091 #[test]
3092 #[cfg(unix)]
3093 fn daemon_drain_gates_current_and_racing_provider_registrations() {
3094 let (forwarding, _, endpoint, _, sink, _) = route_fixture("provider");
3095 assert_eq!(forwarding.begin_daemon_drain().unwrap(), ["provider"]);
3096 assert!(forwarding.endpoint_is_draining(endpoint).unwrap());
3097 assert!(matches!(
3098 forwarding.register_module_connection(
3099 ConnectionId::new(300),
3100 "late-provider".into(),
3101 2,
3102 Concurrency::ModuleManaged,
3103 sink,
3104 ),
3105 Err(ForwardingError::ConnectionClosing { .. })
3106 ));
3107 }
3108
3109 fn test_ping(corr: u64) -> Frame {
3110 Frame::build(
3111 FrameType::Ping,
3112 Flags::new(false, Priority::Passive, false),
3113 0,
3114 0,
3115 corr,
3116 Vec::new(),
3117 )
3118 .unwrap()
3119 }
3120
3121 fn begin_test_route(
3122 forwarding: &ForwardingTable,
3123 client_connection: ConnectionId,
3124 client_sink: FrameSink,
3125 corr: u64,
3126 module_id: &str,
3127 ) -> PendingRouteBindRelay {
3128 forwarding
3129 .begin_route_bind_relay_for_test(client_connection, client_sink, corr, module_id)
3130 .unwrap()
3131 }
3132
3133 #[tokio::test]
3138 async fn drain_holdouts_count_held_requests_and_name_the_connection() {
3139 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3140 route_fixture("holdouts");
3141 let mut bound = |corr| {
3142 let route = begin_test_route(&forwarding, client, sink.clone(), corr, "holdouts");
3143 forwarding
3144 .complete_pending_relay(
3145 module_connection,
3146 route.corr,
3147 RouteBindRelayOutcome::Accepted,
3148 )
3149 .unwrap();
3150 client_rx.try_recv().unwrap();
3151 match forwarding
3152 .lookup_data_route(client, route.client_channel, route.client_epoch)
3153 .unwrap()
3154 {
3155 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3156 other => panic!("expected live route, got {other:?}"),
3157 }
3158 };
3159 let holding = bound(61);
3160 let _idle = bound(62);
3161 holding.flow.acquire_tagged(1, false).await.unwrap();
3162 holding.flow.acquire_tagged(2, false).await.unwrap();
3163 holding.flow.acquire_tagged(3, true).await.unwrap();
3164 forwarding
3165 .begin_module_drain("holdouts", RouteCloseReason::Restart)
3166 .unwrap();
3167
3168 let holdouts = forwarding.endpoint_drain_holdouts(endpoint).unwrap();
3169 assert_eq!(
3170 holdouts,
3171 DrainHoldouts {
3172 requests: 2,
3173 routes: 1,
3174 total_routes: 2,
3175 top_connections: vec![(client.get(), 2)],
3176 }
3177 );
3178 }
3179
3180 #[test]
3181 fn endpoint_routes_keep_goodbye_targets_and_mark_draining_routes() {
3182 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3183 route_fixture("census");
3184 let pending = begin_test_route(&forwarding, client, sink, 1, "census");
3185 forwarding
3186 .complete_pending_relay(
3187 module_connection,
3188 pending.corr,
3189 RouteBindRelayOutcome::Accepted,
3190 )
3191 .unwrap();
3192
3193 let routes = forwarding.endpoint_routes(endpoint).unwrap();
3194 assert_eq!(routes.len(), 1);
3195 assert!(matches!(routes[0].principal, Principal::Direct));
3196 assert_eq!(routes[0].goodbye_target.connection_id, client);
3197 assert_eq!(routes[0].goodbye_target.channel, pending.client_channel);
3198 assert_eq!(routes[0].goodbye_target.epoch, pending.client_epoch);
3199 assert!(!routes[0].draining);
3200
3201 forwarding
3202 .begin_module_drain("census", RouteCloseReason::Restart)
3203 .unwrap();
3204 let draining_routes = forwarding.endpoint_routes(endpoint).unwrap();
3205 assert_eq!(draining_routes.len(), 1);
3206 assert!(draining_routes[0].draining);
3207 }
3208
3209 #[test]
3210 fn aborted_reservation_consumes_both_epochs_and_reuse_advances_them() {
3211 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-abort");
3212 let first = begin_test_route(&forwarding, client, sink.clone(), 1, "epoch-abort");
3213 assert_eq!((first.client_epoch, first.module_epoch), (1, 1));
3214 forwarding
3215 .abort_pending_relay(
3216 first.endpoint,
3217 first.corr,
3218 RouteBindRelayOutcome::ModuleGone("abort".into()),
3219 )
3220 .unwrap();
3221 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3222 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3223
3224 let second = begin_test_route(&forwarding, client, sink, 2, "epoch-abort");
3225 assert_eq!(second.client_channel, first.client_channel);
3226 assert_eq!(second.module_channel, first.module_channel);
3227 assert_eq!((second.client_epoch, second.module_epoch), (2, 2));
3228 }
3229
3230 #[test]
3231 fn stale_release_cannot_remove_reused_successor_and_status_is_epoch_fenced() {
3232 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3233 route_fixture("epoch-release");
3234 let first = begin_test_route(&forwarding, client, sink.clone(), 10, "epoch-release");
3235 forwarding
3236 .complete_pending_relay(
3237 module_connection,
3238 first.corr,
3239 RouteBindRelayOutcome::Accepted,
3240 )
3241 .unwrap();
3242 assert_eq!(client_rx.try_recv().unwrap().header.corr, 10);
3243 assert!(matches!(
3244 forwarding
3245 .release_client_route(client, first.client_channel, first.client_epoch)
3246 .unwrap(),
3247 RouteRelease::Removed(_)
3248 ));
3249 forwarding.inject_client_slot_epoch(client, first.client_channel, first.client_epoch);
3250 forwarding.inject_module_slot_epoch(endpoint, first.module_channel, first.module_epoch);
3251
3252 let second = begin_test_route(&forwarding, client, sink, 11, "epoch-release");
3253 forwarding
3254 .complete_pending_relay(
3255 module_connection,
3256 second.corr,
3257 RouteBindRelayOutcome::Accepted,
3258 )
3259 .unwrap();
3260 assert_eq!(client_rx.try_recv().unwrap().header.corr, 11);
3261 assert!(matches!(
3262 forwarding
3263 .release_client_route(client, second.client_channel, first.client_epoch)
3264 .unwrap(),
3265 RouteRelease::Stale
3266 ));
3267 assert!(!forwarding
3268 .cache_status(
3269 endpoint,
3270 second.module_channel,
3271 first.module_epoch,
3272 "stale".into(),
3273 )
3274 .unwrap());
3275 assert!(forwarding
3276 .cache_status(
3277 endpoint,
3278 second.module_channel,
3279 second.module_epoch,
3280 "current".into(),
3281 )
3282 .unwrap());
3283 match forwarding
3284 .route_poll_snapshot(client, second.client_channel, second.client_epoch)
3285 .unwrap()
3286 {
3287 RoutePollSnapshot::Bound { status, .. } => {
3288 assert_eq!(status.as_deref(), Some("current"));
3289 }
3290 RoutePollSnapshot::Absent => panic!("successor binding was removed"),
3291 }
3292 let counters = forwarding.counters().snapshot();
3293 assert_eq!(counters["route_released_epoch_fenced"], 1);
3294 assert_eq!(counters["route_release_stale_skipped"], 1);
3295 }
3296
3297 #[test]
3298 fn max_epoch_reservation_retires_only_that_slot() {
3299 let (forwarding, _, endpoint, client, sink, _client_rx) = route_fixture("epoch-max");
3300 forwarding.inject_client_slot_epoch(client, 7, u32::MAX - 1);
3301 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX - 1);
3302 let final_use = begin_test_route(&forwarding, client, sink.clone(), 20, "epoch-max");
3303 assert_eq!(
3304 (final_use.client_channel, final_use.client_epoch),
3305 (7, u32::MAX)
3306 );
3307 assert_eq!(
3308 (final_use.module_channel, final_use.module_epoch),
3309 (9, u32::MAX)
3310 );
3311 forwarding
3312 .abort_pending_relay(
3313 endpoint,
3314 final_use.corr,
3315 RouteBindRelayOutcome::ModuleGone("abort".into()),
3316 )
3317 .unwrap();
3318 forwarding.inject_client_slot_epoch(client, 7, u32::MAX);
3319 forwarding.inject_module_slot_epoch(endpoint, 9, u32::MAX);
3320 let next = begin_test_route(&forwarding, client, sink, 21, "epoch-max");
3321 assert_ne!(next.client_channel, 7);
3322 assert_ne!(next.module_channel, 9);
3323 assert_eq!((next.client_epoch, next.module_epoch), (1, 1));
3324 }
3325
3326 #[test]
3327 fn bind_and_module_control_share_monotonic_corr_and_deadline_arbitration() {
3328 let (forwarding, module_connection, endpoint, client, sink, _client_rx) =
3329 route_fixture("corr-shared");
3330 let bind = begin_test_route(&forwarding, client, sink, 30, "corr-shared");
3331 assert_eq!(bind.corr, 1);
3332 forwarding
3333 .abort_pending_relay(
3334 endpoint,
3335 bind.corr,
3336 RouteBindRelayOutcome::ModuleGone("abort".into()),
3337 )
3338 .unwrap();
3339 let rpc = forwarding
3340 .begin_module_control_rpc_for(
3341 "corr-shared",
3342 "health.check",
3343 Instant::now() - Duration::from_millis(1),
3344 )
3345 .unwrap();
3346 assert_eq!(rpc.corr, 2);
3347 assert_eq!(
3348 forwarding
3349 .complete_module_control_rpc(
3350 module_connection,
3351 rpc.corr,
3352 Some("health.check"),
3353 ModuleControlRpcOutcome::Response(ModuleControlResponse::HealthCheck {
3354 status: subc_protocol::session::HealthStatus::Ok,
3355 detail: None,
3356 metrics: None,
3357 }),
3358 )
3359 .unwrap(),
3360 ModuleControlRpcCompletion::Settled
3361 );
3362 assert!(matches!(
3363 rpc.receiver.blocking_recv().unwrap(),
3364 ModuleControlRpcOutcome::DeadlineElapsed
3365 ));
3366 }
3367
3368 #[tokio::test(start_paused = true)]
3369 async fn health_probe_tombstone_ttl_removes_an_endpoint_that_stops_probing() {
3370 let (forwarding, _, endpoint, _, _, _) = route_fixture("tombstone-ttl");
3371 let probe_started_at = Instant::now();
3372 let rpc = forwarding
3373 .begin_health_probe_rpc_for(
3374 "tombstone-ttl",
3375 "health.check",
3376 probe_started_at,
3377 probe_started_at + Duration::from_secs(5),
3378 )
3379 .unwrap();
3380 assert!(forwarding
3381 .tombstone_health_probe_rpc(endpoint, rpc.corr)
3382 .unwrap());
3383 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 1);
3384
3385 tokio::time::advance(HEALTH_PROBE_TOMBSTONE_TTL).await;
3386 tokio::task::yield_now().await;
3387
3388 assert_eq!(forwarding.health_probe_tombstone_count().unwrap(), 0);
3389 }
3390
3391 #[test]
3392 fn correlation_exhaustion_emits_max_once_then_closes_endpoint() {
3393 let (forwarding, _, endpoint, _, _, _) = route_fixture("corr-max");
3394 let mut close = forwarding.register_connection_close(endpoint.connection_id);
3395 forwarding.inject_control_corr(endpoint, u64::MAX);
3396 let final_rpc = forwarding
3397 .begin_module_control_rpc_for(
3398 "corr-max",
3399 "health.check",
3400 Instant::now() + Duration::from_secs(1),
3401 )
3402 .unwrap();
3403 assert_eq!(final_rpc.corr, u64::MAX);
3404 forwarding
3405 .cancel_module_control_rpc(endpoint, final_rpc.corr)
3406 .unwrap();
3407 assert!(matches!(
3408 forwarding.begin_module_control_rpc_for(
3409 "corr-max",
3410 "health.check",
3411 Instant::now() + Duration::from_secs(1),
3412 ),
3413 Err(ForwardingError::RelayCorrelationExhausted)
3414 ));
3415 assert!(close.try_recv().is_ok());
3416 }
3417
3418 #[test]
3419 fn publication_epoch_controls_delivery_failure_escalation() {
3420 fn setup_successor(
3421 commit_successor: Option<bool>,
3422 ) -> (ForwardingTable, ConnectionId, u16, u32) {
3423 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3424 route_fixture("escalation");
3425 let first = begin_test_route(&forwarding, client, sink.clone(), 40, "escalation");
3426 forwarding
3427 .complete_pending_relay(
3428 module_connection,
3429 first.corr,
3430 RouteBindRelayOutcome::Accepted,
3431 )
3432 .unwrap();
3433 client_rx.try_recv().unwrap();
3434 assert!(matches!(
3435 forwarding
3436 .release_client_route(client, first.client_channel, first.client_epoch)
3437 .unwrap(),
3438 RouteRelease::Removed(_)
3439 ));
3440 if let Some(commit_successor) = commit_successor {
3441 forwarding.inject_client_slot_epoch(
3442 client,
3443 first.client_channel,
3444 first.client_epoch,
3445 );
3446 forwarding.inject_module_slot_epoch(
3447 endpoint,
3448 first.module_channel,
3449 first.module_epoch,
3450 );
3451 let successor = begin_test_route(&forwarding, client, sink, 41, "escalation");
3452 if commit_successor {
3453 forwarding
3454 .complete_pending_relay(
3455 module_connection,
3456 successor.corr,
3457 RouteBindRelayOutcome::Accepted,
3458 )
3459 .unwrap();
3460 client_rx.try_recv().unwrap();
3461 } else {
3462 forwarding
3463 .abort_pending_relay(
3464 endpoint,
3465 successor.corr,
3466 RouteBindRelayOutcome::ModuleGone("abort".into()),
3467 )
3468 .unwrap();
3469 }
3470 }
3471 (forwarding, client, first.client_channel, first.client_epoch)
3472 }
3473
3474 let (no_successor, client, channel, epoch) = setup_successor(None);
3475 let mut close = no_successor.register_connection_close(client);
3476 assert!(no_successor
3477 .escalate_client_delivery_failure(
3478 client,
3479 channel,
3480 epoch,
3481 CloseReason::new("delivery", "failed"),
3482 )
3483 .unwrap());
3484 assert!(close.try_recv().is_ok());
3485
3486 let (aborted, client, channel, epoch) = setup_successor(Some(false));
3487 let mut close = aborted.register_connection_close(client);
3488 assert!(aborted
3489 .escalate_client_delivery_failure(
3490 client,
3491 channel,
3492 epoch,
3493 CloseReason::new("delivery", "failed"),
3494 )
3495 .unwrap());
3496 assert!(close.try_recv().is_ok());
3497
3498 let (published, client, channel, epoch) = setup_successor(Some(true));
3499 let mut close = published.register_connection_close(client);
3500 assert!(!published
3501 .escalate_client_delivery_failure(
3502 client,
3503 channel,
3504 epoch,
3505 CloseReason::new("delivery", "stale failure"),
3506 )
3507 .unwrap());
3508 assert!(close.try_recv().is_err());
3509 }
3510
3511 #[test]
3512 fn route_concentration_separates_client_count_from_routes_per_client() {
3513 let (forwarding, module_connection, _, client, sink, _client_rx) =
3517 route_fixture("concentration");
3518 assert_eq!(forwarding.client_route_concentration().unwrap(), (0, 0));
3519
3520 for corr in [70_u64, 71] {
3521 let pending =
3522 begin_test_route(&forwarding, client, sink.clone(), corr, "concentration");
3523 forwarding
3524 .complete_pending_relay(
3525 module_connection,
3526 pending.corr,
3527 RouteBindRelayOutcome::Accepted,
3528 )
3529 .unwrap();
3530 }
3531
3532 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3534 assert_eq!(forwarding.client_route_concentration().unwrap(), (1, 2));
3535 }
3536
3537 #[test]
3538 fn cleanup_and_accepted_resolution_have_one_lock_winner() {
3539 let (forwarding, module_connection, _, client, sink, mut client_rx) =
3540 route_fixture("cleanup-race");
3541 let pending = begin_test_route(&forwarding, client, sink, 45, "cleanup-race");
3542 forwarding
3543 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3544 .unwrap();
3545 let released = forwarding.cleanup_connection(client).unwrap();
3546 assert_eq!(released.len(), 1);
3547 let completion = forwarding
3548 .complete_pending_relay(
3549 module_connection,
3550 pending.corr,
3551 RouteBindRelayOutcome::Accepted,
3552 )
3553 .unwrap();
3554 assert!(!completion.settled);
3555 assert!(client_rx.try_recv().is_err());
3556 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3557
3558 let (forwarding, module_connection, _, client, sink, mut client_rx) =
3559 route_fixture("accepted-race");
3560 let pending = begin_test_route(&forwarding, client, sink, 46, "accepted-race");
3561 forwarding
3562 .complete_pending_relay(
3563 module_connection,
3564 pending.corr,
3565 RouteBindRelayOutcome::Accepted,
3566 )
3567 .unwrap();
3568 assert_eq!(client_rx.try_recv().unwrap().header.corr, 46);
3569 let released = forwarding.cleanup_connection(client).unwrap();
3570 assert_eq!(released.len(), 1);
3571 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3572 }
3573
3574 #[test]
3575 fn drain_marks_block_reservation_commit_and_live_request_admission_until_phase_two() {
3576 let (forwarding, module_connection, _, client, sink, mut client_rx) =
3577 route_fixture("drain-gap");
3578 let live = begin_test_route(&forwarding, client, sink.clone(), 47, "drain-gap");
3579 forwarding
3580 .complete_pending_relay(
3581 module_connection,
3582 live.corr,
3583 RouteBindRelayOutcome::Accepted,
3584 )
3585 .unwrap();
3586 client_rx.try_recv().unwrap();
3587 let binding = match forwarding
3588 .lookup_data_route(client, live.client_channel, live.client_epoch)
3589 .unwrap()
3590 {
3591 DataRoute::Client(DataRouteState::Bound(binding)) => binding,
3592 other => panic!("expected live route, got {other:?}"),
3593 };
3594
3595 let pending = begin_test_route(&forwarding, client, sink.clone(), 48, "drain-gap");
3596 forwarding
3597 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3598 .unwrap();
3599 let control_rpc = forwarding
3600 .begin_module_control_rpc_for(
3601 "drain-gap",
3602 "health.check",
3603 Instant::now() + Duration::from_secs(1),
3604 )
3605 .unwrap();
3606 let target = forwarding
3607 .begin_module_drain("drain-gap", RouteCloseReason::Reload)
3608 .unwrap()
3609 .unwrap();
3610 assert!(matches!(
3611 control_rpc.receiver.blocking_recv().unwrap(),
3612 ModuleControlRpcOutcome::ModuleGone(_)
3613 ));
3614 assert_eq!(target.abandoned_bindings.len(), 1);
3615 assert!(binding.flow.sem.is_closed());
3616 assert!(
3617 !forwarding
3618 .complete_pending_relay(
3619 module_connection,
3620 pending.corr,
3621 RouteBindRelayOutcome::Accepted,
3622 )
3623 .unwrap()
3624 .settled
3625 );
3626 assert!(matches!(
3627 forwarding.begin_route_bind_relay_for_test(client, sink, 49, "drain-gap"),
3628 Err(ForwardingError::ModuleReloading { .. })
3629 ));
3630 let released = forwarding
3631 .release_module_endpoint_routes(target.endpoint)
3632 .unwrap();
3633 assert_eq!(released.len(), 1);
3634 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3635 }
3636
3637 #[test]
3645 fn accepted_bind_for_a_closing_client_releases_the_route_instead_of_failing_the_module() {
3646 let (forwarding, module_connection, endpoint, client, sink, mut client_rx) =
3647 route_fixture("closing-client");
3648
3649 let live = begin_test_route(&forwarding, client, sink.clone(), 60, "closing-client");
3652 forwarding
3653 .complete_pending_relay(
3654 module_connection,
3655 live.corr,
3656 RouteBindRelayOutcome::Accepted,
3657 )
3658 .unwrap();
3659 client_rx.try_recv().unwrap();
3660
3661 let pending = begin_test_route(&forwarding, client, sink.clone(), 61, "closing-client");
3663 forwarding
3664 .mark_route_bind_relay_enqueued(pending.endpoint, pending.corr)
3665 .unwrap();
3666
3667 assert!(forwarding
3669 .escalate_client_delivery_failure(
3670 client,
3671 live.client_channel,
3672 live.client_epoch,
3673 CloseReason::new(
3674 "module_to_client_delivery_failed",
3675 "client egress refused a module frame",
3676 ),
3677 )
3678 .unwrap());
3679 assert!(!sink.is_closed());
3680
3681 let completion = forwarding
3682 .complete_pending_relay(
3683 module_connection,
3684 pending.corr,
3685 RouteBindRelayOutcome::Accepted,
3686 )
3687 .expect("a closing client must not turn a module's ack into an error");
3688
3689 assert!(completion.settled);
3690 let abandoned = completion
3691 .abandoned
3692 .expect("the module must be told to drop the binding it just created");
3693 assert_eq!(abandoned.connection_id, module_connection);
3694 assert_eq!(abandoned.channel, pending.module_channel);
3695 assert_eq!(abandoned.epoch, pending.module_epoch);
3696 assert!(matches!(abandoned.kind, GoodbyeTargetKind::Module));
3697 assert!(matches!(
3698 pending.receiver.blocking_recv().unwrap(),
3699 RouteBindRelayOutcome::ModuleGone(_)
3700 ));
3701 assert!(client_rx.try_recv().is_err());
3704 assert_eq!(forwarding.active_binding_count().unwrap(), 1);
3705
3706 assert!(forwarding
3709 .has_live_module_connection("closing-client")
3710 .unwrap());
3711 let cotenant = ConnectionId::new(201);
3712 let (cotenant_tx, mut cotenant_rx) = mpsc::channel(8);
3713 let cotenant_route = begin_test_route(
3714 &forwarding,
3715 cotenant,
3716 FrameSink::new(cotenant_tx),
3717 62,
3718 "closing-client",
3719 );
3720 assert_eq!(cotenant_route.endpoint, endpoint);
3721 forwarding
3722 .complete_pending_relay(
3723 module_connection,
3724 cotenant_route.corr,
3725 RouteBindRelayOutcome::Accepted,
3726 )
3727 .unwrap();
3728 assert_eq!(cotenant_rx.try_recv().unwrap().header.corr, 62);
3729 assert_eq!(forwarding.active_binding_count().unwrap(), 2);
3730 }
3731
3732 #[test]
3733 fn pending_route_permit_is_released_on_rejection_and_abort() {
3734 let forwarding = ForwardingTable::default();
3735 let module_connection = ConnectionId::new(300);
3736 let client = ConnectionId::new(301);
3737 let (module_tx, _module_rx) = mpsc::channel(1);
3738 let endpoint = forwarding
3739 .register_module_connection(
3740 module_connection,
3741 "permit".into(),
3742 2,
3743 Concurrency::ModuleManaged,
3744 FrameSink::new(module_tx),
3745 )
3746 .unwrap();
3747 let (client_tx, mut client_rx) = mpsc::channel(1);
3748 let sink = FrameSink::new(client_tx);
3749 let rejected = begin_test_route(&forwarding, client, sink.clone(), 50, "permit");
3750 assert!(sink.try_send(test_ping(999)).is_err());
3751 forwarding
3752 .complete_pending_relay(
3753 module_connection,
3754 rejected.corr,
3755 RouteBindRelayOutcome::Rejected(ErrorBody {
3756 code: "no".into(),
3757 message: "rejected".into(),
3758 detail: None,
3759 }),
3760 )
3761 .unwrap();
3762 sink.try_send(test_ping(1000)).unwrap();
3763 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1000);
3764
3765 let aborted = begin_test_route(&forwarding, client, sink.clone(), 51, "permit");
3766 assert!(sink.try_send(test_ping(1001)).is_err());
3767 forwarding
3768 .abort_pending_relay(
3769 endpoint,
3770 aborted.corr,
3771 RouteBindRelayOutcome::ModuleGone("abort".into()),
3772 )
3773 .unwrap();
3774 sink.try_send(test_ping(1002)).unwrap();
3775 assert_eq!(client_rx.try_recv().unwrap().header.corr, 1002);
3776
3777 let receiver_closed = begin_test_route(&forwarding, client, sink, 52, "permit");
3778 forwarding
3779 .mark_route_bind_relay_enqueued(endpoint, receiver_closed.corr)
3780 .unwrap();
3781 drop(client_rx);
3782 let completion = forwarding
3783 .complete_pending_relay(
3784 module_connection,
3785 receiver_closed.corr,
3786 RouteBindRelayOutcome::Accepted,
3787 )
3788 .unwrap();
3789 assert!(completion.abandoned.is_some());
3790 assert_eq!(forwarding.active_binding_count().unwrap(), 0);
3791 }
3792
3793 #[test]
3799 fn cleaned_up_connections_do_not_stay_in_the_closing_set() {
3800 let (forwarding, module_connection, _endpoint, _fixture_client, _sink, _rx) =
3801 route_fixture("closing-set-leak");
3802
3803 const CONNECTIONS: u64 = 32;
3804 for index in 0..CONNECTIONS {
3805 let client = ConnectionId::new(1000 + index);
3806 let (client_tx, _client_rx) = mpsc::channel(8);
3807 let route = begin_test_route(
3808 &forwarding,
3809 client,
3810 FrameSink::new(client_tx),
3811 index + 1,
3812 "closing-set-leak",
3813 );
3814 forwarding
3815 .complete_pending_relay(
3816 module_connection,
3817 route.corr,
3818 RouteBindRelayOutcome::Accepted,
3819 )
3820 .unwrap();
3821 forwarding.cleanup_connection(client).unwrap();
3822 }
3823 forwarding.cleanup_connection(module_connection).unwrap();
3824
3825 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3826 }
3827
3828 #[test]
3835 fn closing_connection_is_refused_new_work_until_cleanup_completes() {
3836 let (forwarding, module_connection, _endpoint, client, sink, mut client_rx) =
3837 route_fixture("closing-gate");
3838
3839 let live = begin_test_route(&forwarding, client, sink.clone(), 80, "closing-gate");
3842 forwarding
3843 .complete_pending_relay(
3844 module_connection,
3845 live.corr,
3846 RouteBindRelayOutcome::Accepted,
3847 )
3848 .unwrap();
3849 client_rx.try_recv().unwrap();
3850
3851 assert!(forwarding
3854 .escalate_client_delivery_failure(
3855 client,
3856 live.client_channel,
3857 live.client_epoch,
3858 CloseReason::new(
3859 "module_to_client_delivery_failed",
3860 "client egress refused a module frame",
3861 ),
3862 )
3863 .unwrap());
3864 assert_eq!(forwarding.closing_connection_count().unwrap(), 1);
3865
3866 assert!(matches!(
3868 forwarding.begin_route_bind_relay_for_test(client, sink, 81, "closing-gate"),
3869 Err(ForwardingError::ConnectionClosing { connection_id })
3870 if connection_id == client
3871 ));
3872 let (late_tx, _late_rx) = mpsc::channel(1);
3874 assert!(matches!(
3875 forwarding.register_module_connection(
3876 client,
3877 "late-module".into(),
3878 2,
3879 Concurrency::ModuleManaged,
3880 FrameSink::new(late_tx),
3881 ),
3882 Err(ForwardingError::ConnectionClosing { connection_id })
3883 if connection_id == client
3884 ));
3885
3886 forwarding.cleanup_connection(client).unwrap();
3890 assert_eq!(forwarding.closing_connection_count().unwrap(), 0);
3891 }
3892}
3893
3894#[cfg(test)]
3897mod swap_slot_tests {
3898 use std::time::Duration;
3899
3900 use super::*;
3901 use tokio::sync::mpsc;
3902
3903 const MODULE_ID: &str = "swapped";
3904
3905 struct SwapFixture {
3906 forwarding: ForwardingTable,
3907 incumbent_connection: ConnectionId,
3908 incumbent: ModuleEndpointId,
3909 candidate_connection: ConnectionId,
3910 candidate: ModuleEndpointId,
3911 _module_rxs: Vec<mpsc::Receiver<crate::router::OutboundFrame>>,
3912 }
3913
3914 fn swap_fixture() -> SwapFixture {
3915 let forwarding = ForwardingTable::default();
3916 let incumbent_connection = ConnectionId::new(100);
3917 let candidate_connection = ConnectionId::new(110);
3918 let (incumbent_tx, incumbent_rx) = mpsc::channel(8);
3919 let incumbent = forwarding
3920 .register_module_connection(
3921 incumbent_connection,
3922 MODULE_ID.to_string(),
3923 2,
3924 Concurrency::ModuleManaged,
3925 FrameSink::new(incumbent_tx),
3926 )
3927 .unwrap();
3928 let (candidate_tx, candidate_rx) = mpsc::channel(8);
3929 let candidate = forwarding
3930 .register_candidate_module_connection(
3931 candidate_connection,
3932 MODULE_ID.to_string(),
3933 2,
3934 Concurrency::ModuleManaged,
3935 FrameSink::new(candidate_tx),
3936 )
3937 .unwrap();
3938 SwapFixture {
3939 forwarding,
3940 incumbent_connection,
3941 incumbent,
3942 candidate_connection,
3943 candidate,
3944 _module_rxs: vec![incumbent_rx, candidate_rx],
3945 }
3946 }
3947
3948 fn client(
3949 raw: u64,
3950 ) -> (
3951 ConnectionId,
3952 FrameSink,
3953 mpsc::Receiver<crate::router::OutboundFrame>,
3954 ) {
3955 let (tx, rx) = mpsc::channel(8);
3956 (ConnectionId::new(raw), FrameSink::new(tx), rx)
3957 }
3958
3959 fn committed_endpoints(forwarding: &ForwardingTable) -> Vec<ModuleEndpointId> {
3960 forwarding
3961 .read_inner()
3962 .unwrap()
3963 .client_to_module
3964 .values()
3965 .map(|route| route.module_endpoint)
3966 .collect()
3967 }
3968
3969 #[test]
3970 fn candidate_is_unroutable_until_cutover_and_by_id_lookups_resolve_the_active_slot() {
3971 let fixture = swap_fixture();
3972 let forwarding = &fixture.forwarding;
3973 assert_ne!(fixture.incumbent, fixture.candidate);
3974
3975 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
3977 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
3978 let (client_connection, client_sink, _client_rx) = client(200);
3979 let pending = forwarding
3980 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
3981 .unwrap();
3982 assert_eq!(pending.endpoint, fixture.incumbent);
3983 let rpc = forwarding
3984 .begin_module_control_rpc_for(
3985 MODULE_ID,
3986 "health.check",
3987 Instant::now() + Duration::from_secs(1),
3988 )
3989 .unwrap();
3990 assert_eq!(rpc.endpoint, fixture.incumbent);
3991 let census = forwarding.route_census(Some(MODULE_ID)).unwrap();
3992 assert_eq!(census.len(), 1, "the census lists one endpoint per id");
3993
3994 assert_eq!(
3996 forwarding
3997 .module_endpoint_for_connection(fixture.candidate_connection)
3998 .unwrap(),
3999 Some(fixture.candidate)
4000 );
4001 assert_eq!(
4002 forwarding
4003 .module_id_for_connection(fixture.candidate_connection)
4004 .unwrap()
4005 .as_deref(),
4006 Some(MODULE_ID)
4007 );
4008
4009 let (other_tx, _other_rx) = mpsc::channel(1);
4011 assert_eq!(
4012 forwarding.register_candidate_module_connection(
4013 ConnectionId::new(120),
4014 MODULE_ID.to_string(),
4015 2,
4016 Concurrency::ModuleManaged,
4017 FrameSink::new(other_tx),
4018 ),
4019 Err(ForwardingError::CandidateSlotOccupied {
4020 module_id: MODULE_ID.to_string()
4021 })
4022 );
4023 }
4024
4025 #[test]
4029 fn relay_reserved_before_cutover_never_commits_and_later_relays_land_on_the_candidate() {
4030 let fixture = swap_fixture();
4031 let forwarding = &fixture.forwarding;
4032 let (early_client, early_sink, _early_rx) = client(200);
4033 let mut early = forwarding
4034 .begin_route_bind_relay_for_test(early_client, early_sink, 1, MODULE_ID)
4035 .unwrap();
4036 assert_eq!(early.endpoint, fixture.incumbent);
4037 assert!(forwarding
4038 .mark_route_bind_relay_enqueued(early.endpoint, early.corr)
4039 .unwrap());
4040
4041 let cutover = forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4042 assert_eq!(
4043 cutover,
4044 ForwardingCutover {
4045 promoted: fixture.candidate,
4046 incumbent: Some(fixture.incumbent),
4047 }
4048 );
4049
4050 let (late_client, late_sink, _late_rx) = client(201);
4052 let late = forwarding
4053 .begin_route_bind_relay_for_test(late_client, late_sink, 2, MODULE_ID)
4054 .unwrap();
4055 assert_eq!(
4056 late.endpoint, fixture.candidate,
4057 "a route.open after cutover was reserved on the incumbent"
4058 );
4059
4060 let completion = forwarding
4062 .complete_pending_relay(
4063 fixture.incumbent_connection,
4064 early.corr,
4065 RouteBindRelayOutcome::Accepted,
4066 )
4067 .expect("a superseded endpoint's ack is not an error on its connection");
4068 assert!(completion.settled);
4069 assert!(
4070 !committed_endpoints(forwarding).contains(&fixture.incumbent),
4071 "a relay reserved before cutover committed a route on the incumbent"
4072 );
4073 let goodbye = completion
4074 .abandoned
4075 .expect("the incumbent is told to drop the binding it just created");
4076 assert_eq!(goodbye.connection_id, fixture.incumbent_connection);
4077 assert_eq!(goodbye.channel, early.module_channel);
4078 assert_eq!(goodbye.epoch, early.module_epoch);
4079 assert_eq!(goodbye.kind, GoodbyeTargetKind::Module);
4080 match early.receiver.try_recv() {
4081 Ok(RouteBindRelayOutcome::Rejected(body)) => assert_eq!(body.code, "module_reloading"),
4082 other => panic!("expected a retryable module_reloading answer, got {other:?}"),
4083 }
4084 assert!(matches!(
4085 forwarding
4086 .lookup_data_route(early_client, early.client_channel, early.client_epoch)
4087 .unwrap(),
4088 DataRoute::Client(DataRouteState::Absent)
4089 ));
4090
4091 assert_eq!(forwarding.reserved_route_count().unwrap(), (1, 1));
4094 forwarding
4095 .complete_pending_relay(
4096 fixture.candidate_connection,
4097 late.corr,
4098 RouteBindRelayOutcome::Accepted,
4099 )
4100 .unwrap();
4101 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4102 assert_eq!(committed_endpoints(forwarding), vec![fixture.candidate]);
4103 }
4104
4105 #[test]
4106 fn endpoint_drain_after_cutover_drains_the_incumbent_not_the_promoted_candidate() {
4107 let fixture = swap_fixture();
4108 let forwarding = &fixture.forwarding;
4109 let (bound_client, bound_sink, _bound_rx) = client(200);
4111 let bound = forwarding
4112 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4113 .unwrap();
4114 forwarding
4115 .complete_pending_relay(
4116 fixture.incumbent_connection,
4117 bound.corr,
4118 RouteBindRelayOutcome::Accepted,
4119 )
4120 .unwrap();
4121 let (pending_client, pending_sink, _pending_rx) = client(201);
4122 let mut in_flight = forwarding
4123 .begin_route_bind_relay_for_test(pending_client, pending_sink, 2, MODULE_ID)
4124 .unwrap();
4125 forwarding
4126 .mark_route_bind_relay_enqueued(in_flight.endpoint, in_flight.corr)
4127 .unwrap();
4128
4129 let incumbent = forwarding
4130 .cutover_candidate(MODULE_ID)
4131 .unwrap()
4132 .unwrap()
4133 .incumbent
4134 .unwrap();
4135 let target = forwarding
4136 .begin_endpoint_drain(incumbent, RouteCloseReason::Restart)
4137 .unwrap()
4138 .expect("the superseded incumbent is still registered");
4139
4140 assert_eq!(target.endpoint, fixture.incumbent);
4141 assert!(forwarding.endpoint_is_draining(fixture.incumbent).unwrap());
4142 assert!(!forwarding.endpoint_is_draining(fixture.candidate).unwrap());
4143 assert!(!forwarding.module_is_draining(MODULE_ID).unwrap());
4144 assert_eq!(target.abandoned_bindings.len(), 1);
4145 assert_eq!(
4146 target.abandoned_bindings[0].channel,
4147 in_flight.module_channel
4148 );
4149 assert!(matches!(
4150 in_flight.receiver.try_recv(),
4151 Ok(RouteBindRelayOutcome::Rejected(body)) if body.code == "module_reloading"
4152 ));
4153 assert_eq!(
4154 forwarding.endpoint_routes(fixture.incumbent).unwrap().len(),
4155 1,
4156 "the incumbent's bound route stays until its drain finishes"
4157 );
4158
4159 let (next_client, next_sink, _next_rx) = client(202);
4160 let next = forwarding
4161 .begin_route_bind_relay_for_test(next_client, next_sink, 3, MODULE_ID)
4162 .expect("the promoted candidate keeps accepting routes");
4163 assert_eq!(next.endpoint, fixture.candidate);
4164 }
4165
4166 #[test]
4172 fn stale_endpoint_ack_without_a_promotion_still_fails_as_before() {
4173 let forwarding = ForwardingTable::default();
4174 let first_connection = ConnectionId::new(70);
4175 let (first_tx, _first_rx) = mpsc::channel(8);
4176 forwarding
4177 .register_module_connection(
4178 first_connection,
4179 MODULE_ID.to_string(),
4180 2,
4181 Concurrency::ModuleManaged,
4182 FrameSink::new(first_tx),
4183 )
4184 .unwrap();
4185 let (client_connection, client_sink, _client_rx) = client(200);
4186 let mut pending = forwarding
4187 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4188 .unwrap();
4189 let (second_tx, _second_rx) = mpsc::channel(8);
4190 forwarding
4191 .register_module_connection(
4192 ConnectionId::new(80),
4193 MODULE_ID.to_string(),
4194 2,
4195 Concurrency::ModuleManaged,
4196 FrameSink::new(second_tx),
4197 )
4198 .unwrap();
4199
4200 assert_eq!(
4201 forwarding
4202 .complete_pending_relay(
4203 first_connection,
4204 pending.corr,
4205 RouteBindRelayOutcome::Accepted
4206 )
4207 .unwrap_err(),
4208 ForwardingError::StaleModuleEndpoint
4209 );
4210 assert!(committed_endpoints(&forwarding).is_empty());
4211 assert_eq!(forwarding.reserved_route_count().unwrap(), (0, 0));
4212 assert!(matches!(
4213 pending.receiver.try_recv(),
4214 Err(oneshot::error::TryRecvError::Closed)
4215 ));
4216 }
4217
4218 #[test]
4219 fn cleanup_releases_candidate_and_superseded_slots_without_touching_the_active_one() {
4220 let fixture = swap_fixture();
4222 let forwarding = &fixture.forwarding;
4223 assert!(forwarding
4224 .cleanup_connection(fixture.candidate_connection)
4225 .unwrap()
4226 .is_empty());
4227 assert_eq!(forwarding.cutover_candidate(MODULE_ID).unwrap(), None);
4228 let (client_connection, client_sink, _client_rx) = client(200);
4229 assert_eq!(
4230 forwarding
4231 .begin_route_bind_relay_for_test(client_connection, client_sink, 1, MODULE_ID)
4232 .unwrap()
4233 .endpoint,
4234 fixture.incumbent
4235 );
4236
4237 let fixture = swap_fixture();
4240 let forwarding = &fixture.forwarding;
4241 let (bound_client, bound_sink, _bound_rx) = client(200);
4242 let bound = forwarding
4243 .begin_route_bind_relay_for_test(bound_client, bound_sink, 1, MODULE_ID)
4244 .unwrap();
4245 forwarding
4246 .complete_pending_relay(
4247 fixture.incumbent_connection,
4248 bound.corr,
4249 RouteBindRelayOutcome::Accepted,
4250 )
4251 .unwrap();
4252 forwarding.cutover_candidate(MODULE_ID).unwrap().unwrap();
4253 let released = forwarding
4254 .cleanup_connection(fixture.incumbent_connection)
4255 .unwrap();
4256 assert_eq!(released.len(), 1);
4257 assert_eq!(released[0].connection_id, bound_client);
4258 assert!(forwarding
4259 .read_inner()
4260 .unwrap()
4261 .superseded_endpoints
4262 .is_empty());
4263 assert!(forwarding.has_live_module_connection(MODULE_ID).unwrap());
4264 let (next_client, next_sink, _next_rx) = client(201);
4265 assert_eq!(
4266 forwarding
4267 .begin_route_bind_relay_for_test(next_client, next_sink, 2, MODULE_ID)
4268 .unwrap()
4269 .endpoint,
4270 fixture.candidate
4271 );
4272 }
4273}