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