1use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::sync::{Arc, Mutex, MutexGuard};
5use std::time::{Duration, Instant};
6
7use aion_core::{ClusterEvent, InterventionCapabilities, WorkerDeathReason, WorkerTransport};
8use aion_proto::{ProtoActivityTask, ProtoRegisterWorker};
9use aion_store::{NamespaceOrigin, NamespacePlacement, NamespaceStore};
10use tokio::sync::{Notify, mpsc};
11
12use crate::cluster_publisher::ClusterEventPublisher;
13use crate::config::AutoCreate;
14use crate::error::ServerError;
15use crate::namespace::{CallerIdentity, NamespaceGuard, NamespaceMinter, NamespaceOperation};
16use crate::observability::Metrics;
17use crate::worker::admission_audit::AdmissionAudit;
18
19pub use aion_core::DEFAULT_TASK_QUEUE;
30
31pub type WorkerTaskSender = mpsc::Sender<WorkerMessage>;
33
34#[derive(Clone, Debug)]
46pub enum WorkerDelivery {
47 Grpc(WorkerTaskSender),
50 #[cfg(feature = "liminal-transport")]
55 Liminal(crate::worker::liminal_transport::LiminalWorkerDelivery),
56}
57
58#[derive(Clone, Debug, Eq, PartialEq)]
60pub enum WorkerMessage {
61 ActivityTask(Box<ProtoActivityTask>),
63 DrainRequest,
65}
66
67#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
79pub struct PoolAddress {
80 namespace: String,
81 task_queue: String,
82}
83
84impl PoolAddress {
85 #[must_use]
89 pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
90 let task_queue = task_queue.into();
91 let task_queue = if task_queue.is_empty() {
92 String::from(DEFAULT_TASK_QUEUE)
93 } else {
94 task_queue
95 };
96 Self {
97 namespace: namespace.into(),
98 task_queue,
99 }
100 }
101
102 #[must_use]
104 pub fn namespace(&self) -> &str {
105 &self.namespace
106 }
107
108 #[must_use]
110 pub fn task_queue(&self) -> &str {
111 &self.task_queue
112 }
113}
114
115#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
119struct ActivityKey {
120 pool: PoolAddress,
121 activity_type: String,
122}
123
124impl ActivityKey {
125 fn new(pool: PoolAddress, activity_type: impl Into<String>) -> Self {
126 Self {
127 pool,
128 activity_type: activity_type.into(),
129 }
130 }
131}
132
133type WorkerMap = HashMap<WorkerId, WorkerHandle>;
134type RegistryMap = HashMap<ActivityKey, WorkerMap>;
135
136#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
138pub struct WorkerId(u64);
139
140impl WorkerId {
141 #[must_use]
144 pub const fn from_value(value: u64) -> Self {
145 Self(value)
146 }
147
148 #[must_use]
151 pub const fn value(self) -> u64 {
152 self.0
153 }
154}
155
156#[derive(Clone, Debug)]
164pub struct WorkerHandle {
165 id: WorkerId,
166 namespaces: BTreeSet<String>,
167 task_queue: String,
168 node: Option<String>,
169 activity_types: BTreeSet<String>,
170 delivery: WorkerDelivery,
171 intervention_capabilities: InterventionCapabilities,
177}
178
179impl WorkerHandle {
180 #[must_use]
182 pub const fn id(&self) -> WorkerId {
183 self.id
184 }
185
186 #[must_use]
189 pub const fn namespaces(&self) -> &BTreeSet<String> {
190 &self.namespaces
191 }
192
193 #[must_use]
195 pub fn task_queue(&self) -> &str {
196 &self.task_queue
197 }
198
199 #[must_use]
202 pub fn node(&self) -> Option<&str> {
203 self.node.as_deref()
204 }
205
206 #[must_use]
208 pub fn activity_types(&self) -> &BTreeSet<String> {
209 &self.activity_types
210 }
211
212 #[must_use]
214 pub const fn delivery(&self) -> &WorkerDelivery {
215 &self.delivery
216 }
217
218 #[must_use]
224 pub const fn intervention_capabilities(&self) -> &InterventionCapabilities {
225 &self.intervention_capabilities
226 }
227
228 #[must_use]
235 pub fn sender(&self) -> Option<&WorkerTaskSender> {
236 match &self.delivery {
237 WorkerDelivery::Grpc(sender) => Some(sender),
238 #[cfg(feature = "liminal-transport")]
239 WorkerDelivery::Liminal(_) => None,
240 }
241 }
242}
243
244#[derive(Debug)]
245struct RegistryState {
246 next_worker_id: u64,
247 workers: BTreeMap<WorkerId, WorkerHandle>,
248 by_activity: RegistryMap,
249 rotation: HashMap<ActivityKey, usize>,
252 last_departure: HashMap<ActivityKey, BTreeMap<Option<String>, Instant>>,
259 dispatch_ineligible: BTreeSet<WorkerId>,
272}
273
274impl Default for RegistryState {
275 fn default() -> Self {
276 Self {
277 next_worker_id: 1,
278 workers: BTreeMap::new(),
279 by_activity: HashMap::new(),
280 rotation: HashMap::new(),
281 last_departure: HashMap::new(),
282 dispatch_ineligible: BTreeSet::new(),
283 }
284 }
285}
286
287#[derive(Clone, Debug)]
289pub struct ConnectedWorkerRegistry {
290 inner: Arc<Mutex<RegistryState>>,
291 metrics: Option<Metrics>,
292 cluster_publisher: Option<ClusterEventPublisher>,
296 minter: Option<NamespaceMinter>,
300 worker_arrived: Arc<Notify>,
301 audit: Arc<AdmissionAudit>,
309}
310
311impl Default for ConnectedWorkerRegistry {
312 fn default() -> Self {
313 Self {
314 inner: Arc::new(Mutex::new(RegistryState::default())),
315 metrics: None,
316 cluster_publisher: None,
317 minter: None,
318 worker_arrived: Arc::new(Notify::new()),
319 audit: Arc::new(AdmissionAudit::new()),
320 }
321 }
322}
323
324impl ConnectedWorkerRegistry {
325 #[must_use]
327 pub fn with_metrics(metrics: Metrics) -> Self {
328 Self {
329 inner: Arc::new(Mutex::new(RegistryState::default())),
330 metrics: Some(metrics),
331 cluster_publisher: None,
332 minter: None,
333 worker_arrived: Arc::new(Notify::new()),
334 audit: Arc::new(AdmissionAudit::new()),
335 }
336 }
337
338 #[must_use]
342 pub fn admission_audit(&self) -> &AdmissionAudit {
343 &self.audit
344 }
345
346 #[must_use]
349 pub fn with_cluster_publisher(mut self, publisher: ClusterEventPublisher) -> Self {
350 self.cluster_publisher = Some(publisher);
351 self
352 }
353
354 #[must_use]
369 pub fn with_namespace_minting(
370 mut self,
371 store: Arc<dyn NamespaceStore>,
372 policy: AutoCreate,
373 ) -> Self {
374 let minter = NamespaceMinter::new(store, policy);
375 let minter = match &self.cluster_publisher {
376 Some(publisher) => minter.with_cluster_publisher(publisher.clone()),
377 None => minter,
378 };
379 self.minter = Some(minter);
380 self
381 }
382
383 #[must_use]
395 pub fn with_namespace_routing(mut self, routing: crate::namespace::NamespaceRouting) -> Self {
396 self.minter = self.minter.map(|minter| minter.with_routing(routing));
397 self
398 }
399
400 pub async fn accept_registration(
406 &self,
407 guard: &NamespaceGuard,
408 caller: &CallerIdentity,
409 registration: &ProtoRegisterWorker,
410 sender: WorkerTaskSender,
411 ) -> Result<WorkerRegistration, ServerError> {
412 guard
418 .scope(caller, &NamespaceOperation::register_worker(registration))
419 .await?;
420 let namespaces = guard.scope_worker_namespaces(caller, ®istration.namespaces)?;
421 self.mint_or_gate_namespaces(&namespaces).await?;
429 let node = optional_node(®istration.node);
430 self.enforce_pinned_placement(&namespaces, node.as_deref())
440 .await?;
441 self.register_namespaces(
442 namespaces,
443 registration.task_queue.clone(),
444 node,
445 registration.activity_types.iter(),
446 sender,
447 )
448 }
449
450 async fn mint_or_gate_namespaces(&self, namespaces: &[String]) -> Result<(), ServerError> {
467 let Some(minter) = &self.minter else {
468 return Ok(());
469 };
470 minter
471 .mint_or_gate(namespaces, NamespaceOrigin::WorkerMint)
472 .await
473 }
474
475 async fn enforce_pinned_placement(
498 &self,
499 namespaces: &[String],
500 node: Option<&str>,
501 ) -> Result<(), ServerError> {
502 let Some(minter) = &self.minter else {
503 return Ok(());
504 };
505 for namespace in namespaces {
506 let NamespacePlacement::Pinned { nodes } = minter.placement_of(namespace).await? else {
507 continue;
508 };
509 let admitted = node.is_some_and(|n| nodes.contains(n));
510 if !admitted {
511 return Err(ServerError::placement_admission_denied(
512 namespace, node, &nodes,
513 ));
514 }
515 }
516 Ok(())
517 }
518
519 pub fn register<'a>(
530 &self,
531 namespace: impl Into<String>,
532 activity_types: impl IntoIterator<Item = &'a String>,
533 sender: WorkerTaskSender,
534 ) -> Result<WorkerRegistration, ServerError> {
535 self.register_namespaces(
536 [namespace.into()],
537 String::from(DEFAULT_TASK_QUEUE),
538 None,
539 activity_types,
540 sender,
541 )
542 }
543
544 pub fn register_pool<'a>(
551 &self,
552 pool: PoolAddress,
553 activity_types: impl IntoIterator<Item = &'a String>,
554 sender: WorkerTaskSender,
555 ) -> Result<WorkerRegistration, ServerError> {
556 let PoolAddress {
557 namespace,
558 task_queue,
559 } = pool;
560 self.register_namespaces([namespace], task_queue, None, activity_types, sender)
561 }
562
563 pub fn register_namespaces<'a>(
575 &self,
576 namespaces: impl IntoIterator<Item = String>,
577 task_queue: impl Into<String>,
578 node: Option<String>,
579 activity_types: impl IntoIterator<Item = &'a String>,
580 sender: WorkerTaskSender,
581 ) -> Result<WorkerRegistration, ServerError> {
582 self.register_delivery(
583 namespaces,
584 task_queue,
585 node,
586 activity_types,
587 WorkerDelivery::Grpc(sender),
588 )
589 }
590
591 pub fn register_delivery<'a>(
604 &self,
605 namespaces: impl IntoIterator<Item = String>,
606 task_queue: impl Into<String>,
607 node: Option<String>,
608 activity_types: impl IntoIterator<Item = &'a String>,
609 delivery: WorkerDelivery,
610 ) -> Result<WorkerRegistration, ServerError> {
611 self.register_delivery_with_capabilities(
612 namespaces,
613 task_queue,
614 node,
615 activity_types,
616 delivery,
617 InterventionCapabilities::none(),
618 )
619 }
620
621 pub fn register_delivery_with_capabilities<'a>(
635 &self,
636 namespaces: impl IntoIterator<Item = String>,
637 task_queue: impl Into<String>,
638 node: Option<String>,
639 activity_types: impl IntoIterator<Item = &'a String>,
640 delivery: WorkerDelivery,
641 intervention_capabilities: InterventionCapabilities,
642 ) -> Result<WorkerRegistration, ServerError> {
643 let namespaces = namespaces.into_iter().collect::<BTreeSet<_>>();
644 let task_queue = task_queue.into();
645 let activity_types = activity_types.into_iter().cloned().collect::<BTreeSet<_>>();
646 let mut state = self.state()?;
647 let worker_id = WorkerId(state.next_worker_id);
648 state.next_worker_id = state.next_worker_id.saturating_add(1);
649
650 let node_for_event = node.clone();
653 let handle = WorkerHandle {
654 id: worker_id,
655 namespaces: namespaces.clone(),
656 task_queue: task_queue.clone(),
657 node,
658 activity_types: activity_types.clone(),
659 delivery,
660 intervention_capabilities,
661 };
662
663 for namespace in &namespaces {
664 let pool = PoolAddress::new(namespace.clone(), task_queue.clone());
665 for activity_type in &activity_types {
666 let key = ActivityKey::new(pool.clone(), activity_type.clone());
667 if let Some(by_node) = state.last_departure.get_mut(&key) {
673 by_node.remove(&handle.node);
674 if by_node.is_empty() {
675 state.last_departure.remove(&key);
676 }
677 }
678 state
679 .by_activity
680 .entry(key)
681 .or_default()
682 .insert(worker_id, handle.clone());
683 }
684 }
685 let transport = transport_of(&handle.delivery);
686 state.workers.insert(worker_id, handle);
687 drop(state);
688
689 if let Some(metrics) = &self.metrics {
690 for namespace in &namespaces {
691 metrics.worker_connected(namespace);
692 }
693 }
694
695 if let Some(publisher) = &self.cluster_publisher {
699 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
700 let task_queue_owned = task_queue.clone();
701 drop(publisher.emit(|meta| ClusterEvent::WorkerConnected {
702 meta,
703 worker_id: worker_id.value().to_string(),
704 namespaces: namespaces_vec,
705 task_queue: task_queue_owned,
706 transport,
707 node: node_for_event,
708 }));
709 }
710
711 self.worker_arrived.notify_waiters();
712
713 Ok(WorkerRegistration {
714 registry: self.clone(),
715 parts: Some(WorkerRegistrationParts {
716 worker_id,
717 namespaces,
718 task_queue,
719 activity_types,
720 }),
721 })
722 }
723
724 pub async fn wait_for_worker(&self) {
730 self.worker_arrived.notified().await;
731 }
732
733 pub fn workers_for(
755 &self,
756 namespace: &str,
757 task_queue: &str,
758 activity_type: &str,
759 node: Option<&str>,
760 ) -> Result<Vec<WorkerHandle>, ServerError> {
761 let mut state = self.state()?;
762 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
763 let mut workers: Vec<WorkerHandle> = state
764 .by_activity
765 .get(&key)
766 .map(|workers| {
767 workers
768 .values()
769 .filter(|worker| worker_matches_node(worker, node))
770 .cloned()
771 .collect()
772 })
773 .unwrap_or_default();
774 if workers.is_empty() {
775 return Ok(workers);
776 }
777 workers.sort_by_key(WorkerHandle::id);
778 let idx = state.rotation.entry(key).or_insert(0);
779 let start = *idx % workers.len();
780 *idx = idx.wrapping_add(1);
781 let mut rotated = Vec::with_capacity(workers.len());
782 rotated.extend_from_slice(&workers[start..]);
783 rotated.extend_from_slice(&workers[..start]);
784 Ok(rotated)
785 }
786
787 pub fn all_workers(&self) -> Result<Vec<WorkerHandle>, ServerError> {
793 let state = self.state()?;
794 Ok(state.workers.values().cloned().collect())
795 }
796
797 pub fn worker_by_id(&self, worker_id: WorkerId) -> Result<Option<WorkerHandle>, ServerError> {
809 Ok(self.state()?.workers.get(&worker_id).cloned())
810 }
811
812 pub fn set_intervention_capabilities(
823 &self,
824 worker_id: WorkerId,
825 capabilities: &InterventionCapabilities,
826 ) -> Result<bool, ServerError> {
827 let mut state = self.state()?;
828 if !state.workers.contains_key(&worker_id) {
829 return Ok(false);
830 }
831 if let Some(handle) = state.workers.get_mut(&worker_id) {
832 handle.intervention_capabilities = capabilities.clone();
833 }
834 for workers in state.by_activity.values_mut() {
837 if let Some(handle) = workers.get_mut(&worker_id) {
838 handle.intervention_capabilities = capabilities.clone();
839 }
840 }
841 Ok(true)
842 }
843
844 pub fn broadcast_drain(&self) -> Result<usize, ServerError> {
854 let workers = self.all_workers()?;
855 let mut delivered = 0usize;
856 for worker in workers {
857 if self.drain_worker(worker.id())? {
858 delivered = delivered.saturating_add(1);
859 }
860 }
861 Ok(delivered)
862 }
863
864 pub fn drain_worker(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
874 let worker = {
875 let mut state = self.state()?;
876 let Some(worker) = state.workers.get(&worker_id).cloned() else {
877 return Ok(false);
878 };
879 Self::remove_worker_from_service(&mut state, &worker);
880 worker
881 };
882 match worker.delivery() {
883 WorkerDelivery::Grpc(sender) => {
884 if sender.try_send(WorkerMessage::DrainRequest).is_ok() {
885 tracing::info!(worker_id = worker_id.value(), "worker drain requested");
886 Ok(true)
887 } else {
888 tracing::error!(
889 worker_id = worker_id.value(),
890 "worker drain signal failed; force-deregistering closed transport"
891 );
892 self.deregister(worker_id)?;
893 Ok(false)
894 }
895 }
896 #[cfg(feature = "liminal-transport")]
897 WorkerDelivery::Liminal(delivery) => {
898 tracing::error!(
899 worker_id = worker_id.value(),
900 connection_pid = delivery.pid(),
901 "liminal transport has no drain control channel; worker fenced and \
902 deregistered at drain start"
903 );
904 self.deregister(worker_id)?;
905 Ok(false)
906 }
907 }
908 }
909
910 pub fn select_worker(
920 &self,
921 namespace: &str,
922 task_queue: &str,
923 activity_type: &str,
924 node: Option<&str>,
925 ) -> Result<Option<WorkerHandle>, ServerError> {
926 let state = self.state()?;
927 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
928 Ok(state.by_activity.get(&key).and_then(|workers| {
929 workers
930 .values()
931 .filter(|worker| worker_matches_node(worker, node))
932 .filter(|worker| !state.dispatch_ineligible.contains(&worker.id))
939 .min_by_key(|worker| worker.id)
940 .cloned()
941 }))
942 }
943
944 pub fn set_dispatch_ineligible(
956 &self,
957 unreachable: BTreeSet<WorkerId>,
958 ) -> Result<(), ServerError> {
959 self.state()?.dispatch_ineligible = unreachable;
960 Ok(())
961 }
962
963 pub fn is_dispatch_ineligible(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
970 Ok(self.state()?.dispatch_ineligible.contains(&worker_id))
971 }
972
973 pub fn dispatch_ineligible(&self) -> Result<BTreeSet<WorkerId>, ServerError> {
985 Ok(self.state()?.dispatch_ineligible.clone())
986 }
987
988 pub fn pool_census(
1004 &self,
1005 namespace: &str,
1006 task_queue: &str,
1007 activity_type: &str,
1008 node: Option<&str>,
1009 ) -> Result<super::queue_service::PoolCensus, ServerError> {
1010 let state = self.state()?;
1011 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1012 let workers_in_pool = state
1013 .workers
1014 .values()
1015 .filter(|worker| {
1016 worker.task_queue == task_queue && worker.namespaces.contains(namespace)
1017 })
1018 .count();
1019 let serving = state.by_activity.get(&key);
1020 let workers_serving_activity = serving.map_or(0, HashMap::len);
1021 let compatible_workers = serving.map_or(0, |workers| {
1022 workers
1023 .values()
1024 .filter(|worker| worker_matches_node(worker, node))
1025 .count()
1026 });
1027 let last_compatible_poller_age = if compatible_workers > 0 {
1028 Some(Duration::ZERO)
1029 } else {
1030 state
1031 .last_departure
1032 .get(&key)
1033 .and_then(|by_node| {
1034 by_node
1035 .iter()
1036 .filter(|(departed_node, _)| match node {
1037 None => true,
1038 Some(node) => departed_node.as_deref() == Some(node),
1039 })
1040 .map(|(_, departed_at)| *departed_at)
1041 .max()
1042 })
1043 .map(|departed_at| departed_at.elapsed())
1044 };
1045 Ok(super::queue_service::PoolCensus {
1046 workers_in_pool,
1047 workers_serving_activity,
1048 compatible_workers,
1049 last_compatible_poller_age,
1050 })
1051 }
1052
1053 pub fn is_registered(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1064 Ok(self.state()?.workers.contains_key(&worker_id))
1065 }
1066
1067 pub fn deregister(&self, worker_id: WorkerId) -> Result<(), ServerError> {
1079 self.deregister_with_reason(worker_id, WorkerDeathReason::Disconnect)
1080 }
1081
1082 pub fn deregister_with_reason(
1090 &self,
1091 worker_id: WorkerId,
1092 reason: WorkerDeathReason,
1093 ) -> Result<(), ServerError> {
1094 let mut state = self.state()?;
1095 let removed_namespaces = Self::remove_worker(&mut state, worker_id);
1096 drop(state);
1097
1098 let Some(namespaces) = removed_namespaces else {
1099 return Ok(());
1101 };
1102
1103 if let Some(metrics) = &self.metrics {
1104 for namespace in &namespaces {
1105 metrics.worker_disconnected(namespace);
1106 }
1107 }
1108 self.emit_worker_disconnected(worker_id, &namespaces, reason);
1109
1110 Ok(())
1111 }
1112
1113 fn emit_worker_disconnected(
1115 &self,
1116 worker_id: WorkerId,
1117 namespaces: &BTreeSet<String>,
1118 reason: WorkerDeathReason,
1119 ) {
1120 if let Some(publisher) = &self.cluster_publisher {
1121 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
1122 drop(publisher.emit(|meta| ClusterEvent::WorkerDisconnected {
1123 meta,
1124 worker_id: worker_id.value().to_string(),
1125 namespaces: namespaces_vec,
1126 reason,
1127 }));
1128 }
1129 }
1130
1131 fn remove_worker(state: &mut RegistryState, worker_id: WorkerId) -> Option<BTreeSet<String>> {
1135 let handle = state.workers.remove(&worker_id)?;
1136 Self::remove_worker_from_service(state, &handle);
1137 Some(handle.namespaces)
1138 }
1139
1140 fn remove_worker_from_service(state: &mut RegistryState, handle: &WorkerHandle) {
1143 let departed_at = Instant::now();
1144 for namespace in &handle.namespaces {
1145 let pool = PoolAddress::new(namespace.clone(), handle.task_queue.clone());
1146 for activity_type in &handle.activity_types {
1147 let key = ActivityKey::new(pool.clone(), activity_type.clone());
1148 state
1152 .last_departure
1153 .entry(key.clone())
1154 .or_default()
1155 .insert(handle.node.clone(), departed_at);
1156 if let Some(workers) = state.by_activity.get_mut(&key) {
1157 workers.remove(&handle.id);
1158 if workers.is_empty() {
1159 state.by_activity.remove(&key);
1160 state.rotation.remove(&key);
1167 }
1168 }
1169 }
1170 }
1171 }
1172
1173 fn state(&self) -> Result<MutexGuard<'_, RegistryState>, ServerError> {
1174 self.inner
1175 .lock()
1176 .map_err(|_| ServerError::lock_poisoned("connected worker registry"))
1177 }
1178}
1179
1180const fn transport_of(delivery: &WorkerDelivery) -> WorkerTransport {
1183 match delivery {
1184 WorkerDelivery::Grpc(_) => WorkerTransport::Grpc,
1185 #[cfg(feature = "liminal-transport")]
1186 WorkerDelivery::Liminal(_) => WorkerTransport::Liminal,
1187 }
1188}
1189
1190pub(crate) fn optional_node(node: &str) -> Option<String> {
1200 if node.is_empty() {
1201 None
1202 } else {
1203 Some(node.to_owned())
1204 }
1205}
1206
1207fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
1212 match node {
1213 None => true,
1214 Some(node) => worker.node() == Some(node),
1215 }
1216}
1217
1218#[derive(Clone, Debug)]
1219struct WorkerRegistrationParts {
1220 worker_id: WorkerId,
1221 namespaces: BTreeSet<String>,
1222 task_queue: String,
1223 activity_types: BTreeSet<String>,
1224}
1225
1226#[derive(Debug)]
1231pub struct WorkerRegistration {
1232 registry: ConnectedWorkerRegistry,
1233 parts: Option<WorkerRegistrationParts>,
1234}
1235
1236impl WorkerRegistration {
1237 #[must_use]
1239 pub fn worker_id(&self) -> Option<WorkerId> {
1240 self.parts.as_ref().map(|parts| parts.worker_id)
1241 }
1242
1243 #[must_use]
1245 pub fn namespaces(&self) -> Option<&BTreeSet<String>> {
1246 self.parts.as_ref().map(|parts| &parts.namespaces)
1247 }
1248
1249 #[must_use]
1251 pub fn task_queue(&self) -> Option<&str> {
1252 self.parts.as_ref().map(|parts| parts.task_queue.as_str())
1253 }
1254
1255 #[must_use]
1257 pub fn activity_types(&self) -> Option<&BTreeSet<String>> {
1258 self.parts.as_ref().map(|parts| &parts.activity_types)
1259 }
1260
1261 pub fn deregister(mut self) -> Result<(), ServerError> {
1267 let Some(parts) = self.parts.take() else {
1268 return Ok(());
1269 };
1270 self.registry.deregister(parts.worker_id)
1271 }
1272}
1273
1274impl Drop for WorkerRegistration {
1275 fn drop(&mut self) {
1276 let Some(parts) = self.parts.take() else {
1277 return;
1278 };
1279 let removed_namespaces = self.registry.inner.lock().ok().and_then(|mut state| {
1280 ConnectedWorkerRegistry::remove_worker(&mut state, parts.worker_id)
1281 });
1282 if let Some(namespaces) = removed_namespaces {
1283 if let Some(metrics) = &self.registry.metrics {
1284 for namespace in &namespaces {
1285 metrics.worker_disconnected(namespace);
1286 }
1287 }
1288 self.registry.emit_worker_disconnected(
1292 parts.worker_id,
1293 &namespaces,
1294 WorkerDeathReason::Disconnect,
1295 );
1296 }
1297 }
1298}
1299
1300#[cfg(test)]
1301mod tests {
1302 use crate::config::NamespaceMode;
1303 use crate::namespace::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
1304 use crate::worker::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
1305
1306 use super::*;
1307
1308 fn guard() -> NamespaceGuard {
1309 NamespaceGuard::new(NamespaceResolver::authorization_only(
1310 NamespaceMode::SharedEngine,
1311 StaticWorkflowNamespaces::default(),
1312 StaticScheduleNamespaces::default(),
1313 ))
1314 }
1315
1316 fn caller(namespace: &str) -> CallerIdentity {
1317 CallerIdentity::new("worker", [namespace.to_owned()])
1318 }
1319
1320 fn test_failure(message: &str) -> ServerError {
1323 ServerError::worker_dispatch("default".to_owned(), "test".to_owned(), message.to_owned())
1324 }
1325
1326 fn registration(namespace: &str, activity_types: &[&str]) -> ProtoRegisterWorker {
1327 registration_with_queue(namespace, "", activity_types)
1328 }
1329
1330 fn registration_with_queue(
1331 namespace: &str,
1332 task_queue: &str,
1333 activity_types: &[&str],
1334 ) -> ProtoRegisterWorker {
1335 registration_full(&[namespace], task_queue, "", activity_types)
1336 }
1337
1338 fn registration_full(
1339 namespaces: &[&str],
1340 task_queue: &str,
1341 node: &str,
1342 activity_types: &[&str],
1343 ) -> ProtoRegisterWorker {
1344 ProtoRegisterWorker {
1345 namespaces: namespaces.iter().map(|value| (*value).to_owned()).collect(),
1346 activity_types: activity_types
1347 .iter()
1348 .map(|value| (*value).to_owned())
1349 .collect(),
1350 task_queue: task_queue.to_owned(),
1351 node: node.to_owned(),
1352 activities: Vec::new(),
1353 identity: String::new(),
1354 }
1355 }
1356
1357 fn multi_caller(namespaces: &[&str]) -> CallerIdentity {
1358 CallerIdentity::new("worker", namespaces.iter().map(|value| (*value).to_owned()))
1359 }
1360
1361 #[tokio::test]
1367 async fn set_intervention_capabilities_updates_live_worker() -> Result<(), ServerError> {
1368 let registry = ConnectedWorkerRegistry::default();
1369 let (sender, _receiver) = mpsc::channel(1);
1370 let types = ["scout".to_owned()];
1371 let guard = registry.register_delivery_with_capabilities(
1372 ["default".to_owned()],
1373 "default",
1374 None,
1375 types.iter(),
1376 WorkerDelivery::Grpc(sender),
1377 InterventionCapabilities::none(),
1378 )?;
1379 let Some(worker_id) = guard.worker_id() else {
1380 return Err(test_failure("registration carries an id"));
1381 };
1382
1383 let announced = InterventionCapabilities {
1384 supported: vec![aion_core::InterventionPrimitive::InjectMessage],
1385 };
1386 assert!(
1387 registry.set_intervention_capabilities(worker_id, &announced)?,
1388 "a live worker's capabilities must be updatable"
1389 );
1390 let Some(handle) = registry.worker_by_id(worker_id)? else {
1391 return Err(test_failure("worker stays registered"));
1392 };
1393 assert_eq!(handle.intervention_capabilities(), &announced);
1394
1395 assert!(
1396 !registry.set_intervention_capabilities(WorkerId(u64::MAX), &announced)?,
1397 "an unknown worker reports false, never an error"
1398 );
1399 Ok(())
1400 }
1401
1402 #[tokio::test]
1403 async fn register_and_deregister_are_namespace_isolated() -> Result<(), ServerError> {
1404 let registry = ConnectedWorkerRegistry::default();
1405 let (tenant_a_tx, _tenant_a_rx) = mpsc::channel(1);
1406 let (tenant_b_tx, _tenant_b_rx) = mpsc::channel(1);
1407
1408 let tenant_a = registry
1409 .accept_registration(
1410 &guard(),
1411 &caller("tenant-a"),
1412 ®istration("tenant-a", &["charge", "charge"]),
1413 tenant_a_tx,
1414 )
1415 .await?;
1416 let tenant_b = registry
1417 .accept_registration(
1418 &guard(),
1419 &caller("tenant-b"),
1420 ®istration("tenant-b", &["charge"]),
1421 tenant_b_tx,
1422 )
1423 .await?;
1424
1425 let tq = DEFAULT_TASK_QUEUE;
1426 assert_eq!(
1427 registry.workers_for("tenant-a", tq, "charge", None)?.len(),
1428 1
1429 );
1430 assert_eq!(
1431 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1432 1
1433 );
1434 assert!(
1435 registry
1436 .workers_for("tenant-a", tq, "missing", None)?
1437 .is_empty()
1438 );
1439
1440 let tenant_a_id = tenant_a.worker_id();
1441 tenant_a.deregister()?;
1442
1443 assert!(
1444 registry
1445 .workers_for("tenant-a", tq, "charge", None)?
1446 .is_empty()
1447 );
1448 assert_eq!(
1449 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1450 1
1451 );
1452 assert_ne!(tenant_a_id, tenant_b.worker_id());
1453
1454 tenant_b.deregister()?;
1455 assert!(
1456 registry
1457 .workers_for("tenant-b", tq, "charge", None)?
1458 .is_empty()
1459 );
1460 Ok(())
1461 }
1462
1463 #[tokio::test]
1464 async fn denied_namespace_is_not_registered() -> Result<(), ServerError> {
1465 let registry = ConnectedWorkerRegistry::default();
1466 let (tx, _rx) = mpsc::channel(1);
1467 let denied = registry
1468 .accept_registration(
1469 &guard(),
1470 &caller("tenant-a"),
1471 ®istration("tenant-b", &["charge"]),
1472 tx,
1473 )
1474 .await;
1475
1476 assert!(denied.is_err());
1477 assert!(
1478 registry
1479 .workers_for("tenant-b", DEFAULT_TASK_QUEUE, "charge", None)?
1480 .is_empty()
1481 );
1482 Ok(())
1483 }
1484
1485 #[tokio::test]
1486 async fn task_queues_partition_disjoint_pools_within_one_namespace() -> Result<(), ServerError>
1487 {
1488 let registry = ConnectedWorkerRegistry::default();
1492 let (norn_tx, _norn_rx) = mpsc::channel(1);
1493 let (claude_a_tx, _claude_a_rx) = mpsc::channel(1);
1494 let (claude_b_tx, _claude_b_rx) = mpsc::channel(1);
1495
1496 let norn = registry
1497 .accept_registration(
1498 &guard(),
1499 &caller("local"),
1500 ®istration_with_queue("local", "norn", &["dev"]),
1501 norn_tx,
1502 )
1503 .await?;
1504 let claude_a = registry
1506 .accept_registration(
1507 &guard(),
1508 &caller("local"),
1509 ®istration_with_queue("local", "claude", &["dev"]),
1510 claude_a_tx,
1511 )
1512 .await?;
1513 let claude_b = registry
1514 .accept_registration(
1515 &guard(),
1516 &caller("local"),
1517 ®istration_with_queue("local", "claude", &["dev"]),
1518 claude_b_tx,
1519 )
1520 .await?;
1521
1522 let norn_pool = registry.workers_for("local", "norn", "dev", None)?;
1523 assert_eq!(norn_pool.len(), 1, "norn pool has exactly its one worker");
1524 let norn_id = norn.worker_id().ok_or_else(missing_id)?;
1525 assert_eq!(norn_pool[0].id(), norn_id);
1526
1527 let claude_pool = registry.workers_for("local", "claude", "dev", None)?;
1528 assert_eq!(
1529 claude_pool.len(),
1530 2,
1531 "claude pool sees only its two workers"
1532 );
1533 let claude_ids: BTreeSet<WorkerId> = claude_pool.iter().map(WorkerHandle::id).collect();
1534 assert!(
1535 !claude_ids.contains(&norn_id),
1536 "the norn worker must never appear in the claude pool"
1537 );
1538
1539 assert!(
1542 !registry
1543 .workers_for("local", "norn", "dev", None)?
1544 .iter()
1545 .any(|worker| claude_ids.contains(&worker.id()))
1546 );
1547
1548 let first = registry.workers_for("local", "claude", "dev", None)?[0].id();
1552 let second = registry.workers_for("local", "claude", "dev", None)?[0].id();
1553 assert_ne!(
1554 first, second,
1555 "claude pool round-robins across both workers"
1556 );
1557 assert_eq!(
1558 registry.workers_for("local", "norn", "dev", None)?[0].id(),
1559 norn_id,
1560 "the norn pool rotation is unaffected by claude traffic"
1561 );
1562
1563 norn.deregister()?;
1564 claude_a.deregister()?;
1565 claude_b.deregister()?;
1566 Ok(())
1567 }
1568
1569 #[tokio::test]
1570 async fn same_task_queue_in_different_namespaces_is_isolated() -> Result<(), ServerError> {
1571 let registry = ConnectedWorkerRegistry::default();
1574 let (local_tx, _local_rx) = mpsc::channel(1);
1575 let (remote_tx, _remote_rx) = mpsc::channel(1);
1576
1577 let local = registry
1578 .accept_registration(
1579 &guard(),
1580 &caller("local"),
1581 ®istration_with_queue("local", "gpu", &["render"]),
1582 local_tx,
1583 )
1584 .await?;
1585 let remote = registry
1586 .accept_registration(
1587 &guard(),
1588 &caller("remote"),
1589 ®istration_with_queue("remote", "gpu", &["render"]),
1590 remote_tx,
1591 )
1592 .await?;
1593
1594 let local_pool = registry.workers_for("local", "gpu", "render", None)?;
1595 let remote_pool = registry.workers_for("remote", "gpu", "render", None)?;
1596 assert_eq!(local_pool.len(), 1);
1597 assert_eq!(remote_pool.len(), 1);
1598 assert_ne!(
1599 local_pool[0].id(),
1600 remote_pool[0].id(),
1601 "a shared task_queue string does not merge two namespaces"
1602 );
1603
1604 local.deregister()?;
1605 assert!(
1606 registry
1607 .workers_for("local", "gpu", "render", None)?
1608 .is_empty(),
1609 "deregistering the local worker leaves the remote namespace untouched"
1610 );
1611 assert_eq!(
1612 registry.workers_for("remote", "gpu", "render", None)?.len(),
1613 1
1614 );
1615
1616 remote.deregister()?;
1617 Ok(())
1618 }
1619
1620 #[tokio::test]
1621 async fn worker_serving_a_namespace_set_is_reachable_in_each() -> Result<(), ServerError> {
1622 let registry = ConnectedWorkerRegistry::default();
1625 let (ab_tx, _ab_rx) = mpsc::channel(1);
1626 let (a_tx, _a_rx) = mpsc::channel(1);
1627
1628 let worker_ab = registry
1629 .accept_registration(
1630 &guard(),
1631 &multi_caller(&["a", "b"]),
1632 ®istration_full(&["a", "b"], "default", "", &["dev"]),
1633 ab_tx,
1634 )
1635 .await?;
1636 let worker_a = registry
1637 .accept_registration(
1638 &guard(),
1639 &caller("a"),
1640 ®istration_full(&["a"], "default", "", &["dev"]),
1641 a_tx,
1642 )
1643 .await?;
1644
1645 let in_a = registry.workers_for("a", "default", "dev", None)?;
1646 let in_b = registry.workers_for("b", "default", "dev", None)?;
1647 let both_id = worker_ab.worker_id().ok_or_else(missing_id)?;
1648 let only_a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1649
1650 let a_ids: BTreeSet<WorkerId> = in_a.iter().map(WorkerHandle::id).collect();
1652 assert_eq!(a_ids, BTreeSet::from([both_id, only_a_id]));
1653 assert_eq!(in_b.len(), 1, "only the {{a, b}} worker is reachable in b");
1654 assert_eq!(in_b[0].id(), both_id);
1655 assert!(
1656 !in_b.iter().any(|worker| worker.id() == only_a_id),
1657 "the {{a}}-only worker must not be reachable in b"
1658 );
1659
1660 worker_ab.deregister()?;
1662 assert!(
1663 registry
1664 .workers_for("b", "default", "dev", None)?
1665 .is_empty()
1666 );
1667 assert_eq!(registry.workers_for("a", "default", "dev", None)?.len(), 1);
1668
1669 worker_a.deregister()?;
1670 Ok(())
1671 }
1672
1673 #[tokio::test]
1674 async fn node_pin_filters_within_pool() -> Result<(), ServerError> {
1675 let registry = ConnectedWorkerRegistry::default();
1679 let (n1_tx, _n1_rx) = mpsc::channel(1);
1680 let (n2_tx, _n2_rx) = mpsc::channel(1);
1681
1682 let on_n1 = registry
1683 .accept_registration(
1684 &guard(),
1685 &caller("ns"),
1686 ®istration_full(&["ns"], "tq", "n1", &["dev"]),
1687 n1_tx,
1688 )
1689 .await?;
1690 let on_n2 = registry
1691 .accept_registration(
1692 &guard(),
1693 &caller("ns"),
1694 ®istration_full(&["ns"], "tq", "n2", &["dev"]),
1695 n2_tx,
1696 )
1697 .await?;
1698 let n1_id = on_n1.worker_id().ok_or_else(missing_id)?;
1699 let n2_id = on_n2.worker_id().ok_or_else(missing_id)?;
1700
1701 let unpinned = registry.workers_for("ns", "tq", "dev", None)?;
1703 assert_eq!(unpinned.len(), 2, "unpinned reaches the whole pool");
1704 let first = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1705 let second = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1706 assert_ne!(first, second, "unpinned round-robins across both nodes");
1707
1708 let pinned_n1 = registry.workers_for("ns", "tq", "dev", Some("n1"))?;
1710 assert_eq!(pinned_n1.len(), 1);
1711 assert_eq!(pinned_n1[0].id(), n1_id);
1712 let pinned_n2 = registry.workers_for("ns", "tq", "dev", Some("n2"))?;
1713 assert_eq!(pinned_n2.len(), 1);
1714 assert_eq!(pinned_n2[0].id(), n2_id);
1715
1716 assert_eq!(
1718 registry
1719 .select_worker("ns", "tq", "dev", Some("n1"))?
1720 .map(|worker| worker.id()),
1721 Some(n1_id)
1722 );
1723
1724 assert!(
1727 registry
1728 .workers_for("ns", "tq", "dev", Some("absent"))?
1729 .is_empty(),
1730 "a pin to a node with no worker yields no candidate"
1731 );
1732 assert!(
1733 registry
1734 .select_worker("ns", "tq", "dev", Some("absent"))?
1735 .is_none()
1736 );
1737
1738 on_n1.deregister()?;
1739 on_n2.deregister()?;
1740 Ok(())
1741 }
1742
1743 #[tokio::test]
1744 async fn shared_node_id_round_robins_across_workers() -> Result<(), ServerError> {
1745 let registry = ConnectedWorkerRegistry::default();
1748 let (a_tx, _a_rx) = mpsc::channel(1);
1749 let (b_tx, _b_rx) = mpsc::channel(1);
1750
1751 let worker_a = registry
1752 .accept_registration(
1753 &guard(),
1754 &caller("ns"),
1755 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1756 a_tx,
1757 )
1758 .await?;
1759 let worker_b = registry
1760 .accept_registration(
1761 &guard(),
1762 &caller("ns"),
1763 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1764 b_tx,
1765 )
1766 .await?;
1767 let a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1768 let b_id = worker_b.worker_id().ok_or_else(missing_id)?;
1769
1770 let pinned = registry.workers_for("ns", "tq", "dev", Some("shared"))?;
1771 assert_eq!(
1772 pinned.len(),
1773 2,
1774 "both workers on the shared node are candidates"
1775 );
1776 let pinned_ids: BTreeSet<WorkerId> = pinned.iter().map(WorkerHandle::id).collect();
1777 assert_eq!(pinned_ids, BTreeSet::from([a_id, b_id]));
1778
1779 let first = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1780 let second = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1781 assert_ne!(
1782 first, second,
1783 "a pin to a shared node round-robins across both workers on it"
1784 );
1785
1786 worker_a.deregister()?;
1787 worker_b.deregister()?;
1788 Ok(())
1789 }
1790
1791 #[tokio::test]
1792 async fn rotation_cursor_is_pruned_when_last_worker_leaves() -> Result<(), ServerError> {
1793 let registry = ConnectedWorkerRegistry::default();
1798 let (tx, _rx) = mpsc::channel(1);
1799 let worker = registry
1800 .accept_registration(
1801 &guard(),
1802 &caller("ns"),
1803 ®istration_full(&["ns"], "tq", "", &["dev"]),
1804 tx,
1805 )
1806 .await?;
1807
1808 let _ = registry.workers_for("ns", "tq", "dev", None)?;
1810 let key = ActivityKey::new(PoolAddress::new("ns", "tq"), "dev");
1811 assert!(
1812 registry.state()?.rotation.contains_key(&key),
1813 "a lookup must have created the rotation cursor"
1814 );
1815
1816 worker.deregister()?;
1817 let state = registry.state()?;
1818 assert!(
1819 !state.rotation.contains_key(&key),
1820 "the rotation cursor must be pruned once the last worker leaves"
1821 );
1822 assert!(
1823 !state.by_activity.contains_key(&key),
1824 "the activity bucket must also be gone"
1825 );
1826 Ok(())
1827 }
1828
1829 fn missing_id() -> ServerError {
1830 ServerError::lock_poisoned("registration unexpectedly missing a worker id")
1831 }
1832
1833 fn namespace_store() -> Arc<dyn NamespaceStore> {
1836 Arc::new(aion_store::InMemoryStore::default())
1837 }
1838
1839 fn minting_registry(
1840 store: &Arc<dyn NamespaceStore>,
1841 policy: AutoCreate,
1842 ) -> ConnectedWorkerRegistry {
1843 ConnectedWorkerRegistry::default().with_namespace_minting(Arc::clone(store), policy)
1844 }
1845
1846 #[tokio::test]
1847 async fn open_register_mints_durable_record_and_is_idempotent() -> Result<(), ServerError> {
1848 let store = namespace_store();
1849 let registry = minting_registry(&store, AutoCreate::Open);
1850
1851 let (tx_one, _rx_one) = mpsc::channel(1);
1853 let first = registry
1854 .accept_registration(
1855 &guard(),
1856 &caller("orders"),
1857 ®istration("orders", &["charge"]),
1858 tx_one,
1859 )
1860 .await?;
1861 let record = store
1862 .get_namespace("orders")
1863 .await?
1864 .ok_or_else(|| ServerError::namespace_denied("expected a minted record"))?;
1865 assert_eq!(record.name, "orders");
1866 assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
1867
1868 let (tx_two, _rx_two) = mpsc::channel(1);
1871 registry
1872 .accept_registration(
1873 &guard(),
1874 &caller("orders"),
1875 ®istration("orders", &["refund"]),
1876 tx_two,
1877 )
1878 .await?;
1879 let all = store.list_namespaces().await?;
1880 assert_eq!(
1881 all.iter().filter(|r| r.name == "orders").count(),
1882 1,
1883 "re-register must not create a duplicate namespace row"
1884 );
1885 drop(first);
1886 Ok(())
1887 }
1888
1889 #[tokio::test]
1890 async fn open_register_mints_each_namespace_in_a_multi_namespace_worker()
1891 -> Result<(), ServerError> {
1892 let store = namespace_store();
1893 let registry = minting_registry(&store, AutoCreate::Open);
1894 let (tx, _rx) = mpsc::channel(1);
1895
1896 registry
1897 .accept_registration(
1898 &guard(),
1899 &multi_caller(&["alpha", "beta"]),
1900 ®istration_full(&["alpha", "beta"], "", "", &["charge"]),
1901 tx,
1902 )
1903 .await?;
1904
1905 assert!(store.get_namespace("alpha").await?.is_some());
1906 assert!(store.get_namespace("beta").await?.is_some());
1907 Ok(())
1908 }
1909
1910 async fn pinned_registry(
1916 store: &Arc<dyn NamespaceStore>,
1917 namespace: &str,
1918 nodes: &[&str],
1919 ) -> Result<ConnectedWorkerRegistry, ServerError> {
1920 store
1921 .register_namespace(namespace, NamespaceOrigin::Explicit)
1922 .await?;
1923 store
1924 .set_namespace_placement(
1925 namespace,
1926 NamespacePlacement::Pinned {
1927 nodes: nodes.iter().map(|n| (*n).to_owned()).collect(),
1928 },
1929 )
1930 .await?;
1931 Ok(minting_registry(store, AutoCreate::Open))
1932 }
1933
1934 #[tokio::test]
1937 async fn pinned_admits_a_worker_on_a_required_node() -> Result<(), ServerError> {
1938 let store = namespace_store();
1939 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1940 let (tx, _rx) = mpsc::channel(1);
1941
1942 let _registration = registry
1943 .accept_registration(
1944 &guard(),
1945 &caller("iso"),
1946 ®istration_full(&["iso"], "", "n1", &["charge"]),
1947 tx,
1948 )
1949 .await?;
1950
1951 assert_eq!(
1952 registry
1953 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", Some("n1"))?
1954 .len(),
1955 1,
1956 "an n1 worker must be admitted into the Pinned{{n1}} namespace's pool"
1957 );
1958 Ok(())
1959 }
1960
1961 #[tokio::test]
1965 async fn pinned_rejects_a_wrong_node_worker() -> Result<(), ServerError> {
1966 let store = namespace_store();
1967 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1968 let (tx, _rx) = mpsc::channel(1);
1969
1970 let denied = registry
1971 .accept_registration(
1972 &guard(),
1973 &caller("iso"),
1974 ®istration_full(&["iso"], "", "n2", &["charge"]),
1975 tx,
1976 )
1977 .await;
1978 assert!(
1979 matches!(denied, Err(ServerError::Namespace { .. })),
1980 "a wrong-node (n2) worker must be rejected from a Pinned{{n1}} namespace"
1981 );
1982 assert!(
1983 registry
1984 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
1985 .is_empty(),
1986 "a rejected registration must not insert a worker on any node"
1987 );
1988 Ok(())
1989 }
1990
1991 #[tokio::test]
1994 async fn pinned_rejects_a_node_less_worker() -> Result<(), ServerError> {
1995 let store = namespace_store();
1996 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1997 let (tx, _rx) = mpsc::channel(1);
1998
1999 let denied = registry
2000 .accept_registration(
2001 &guard(),
2002 &caller("iso"),
2003 ®istration_full(&["iso"], "", "", &["charge"]),
2004 tx,
2005 )
2006 .await;
2007 assert!(
2008 matches!(denied, Err(ServerError::Namespace { .. })),
2009 "a node-less worker must be rejected from a Pinned{{n1}} namespace"
2010 );
2011 assert!(
2012 registry
2013 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
2014 .is_empty(),
2015 "a rejected node-less registration must not insert a worker"
2016 );
2017 Ok(())
2018 }
2019
2020 #[tokio::test]
2024 async fn pinned_violation_rejects_the_whole_multi_namespace_registration()
2025 -> Result<(), ServerError> {
2026 let store = namespace_store();
2027 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2028 let (tx, _rx) = mpsc::channel(1);
2029
2030 let denied = registry
2031 .accept_registration(
2032 &guard(),
2033 &multi_caller(&["free", "iso"]),
2034 ®istration_full(&["free", "iso"], "", "n2", &["charge"]),
2035 tx,
2036 )
2037 .await;
2038 assert!(
2039 matches!(denied, Err(ServerError::Namespace { .. })),
2040 "a wrong-node worker serving a Pinned namespace fails the WHOLE registration"
2041 );
2042 assert!(
2043 registry
2044 .workers_for("free", DEFAULT_TASK_QUEUE, "charge", None)?
2045 .is_empty(),
2046 "the compliant namespace must NOT be partially admitted"
2047 );
2048 Ok(())
2049 }
2050
2051 #[tokio::test]
2054 async fn unplaced_and_prefer_admission_is_unaffected_by_the_pinned_gate()
2055 -> Result<(), ServerError> {
2056 let store = namespace_store();
2057 store
2060 .register_namespace("pref", NamespaceOrigin::Explicit)
2061 .await?;
2062 store
2063 .set_namespace_placement(
2064 "pref",
2065 NamespacePlacement::Prefer {
2066 nodes: ["n1".to_owned()].into_iter().collect(),
2067 },
2068 )
2069 .await?;
2070 let registry = minting_registry(&store, AutoCreate::Open);
2071
2072 let (tx_a, _rx_a) = mpsc::channel(1);
2073 let _reg_a = registry
2074 .accept_registration(
2075 &guard(),
2076 &caller("unpl"),
2077 ®istration_full(&["unpl"], "", "", &["charge"]),
2078 tx_a,
2079 )
2080 .await?;
2081 let (tx_b, _rx_b) = mpsc::channel(1);
2082 let _reg_b = registry
2083 .accept_registration(
2084 &guard(),
2085 &caller("pref"),
2086 ®istration_full(&["pref"], "", "", &["charge"]),
2087 tx_b,
2088 )
2089 .await?;
2090
2091 assert_eq!(
2092 registry
2093 .workers_for("unpl", DEFAULT_TASK_QUEUE, "charge", None)?
2094 .len(),
2095 1,
2096 "an Unplaced namespace admits a node-less worker unchanged"
2097 );
2098 assert_eq!(
2099 registry
2100 .workers_for("pref", DEFAULT_TASK_QUEUE, "charge", None)?
2101 .len(),
2102 1,
2103 "a Prefer namespace admits a node-less worker unchanged (only Pinned gates)"
2104 );
2105 Ok(())
2106 }
2107
2108 #[tokio::test]
2112 async fn no_minter_registry_skips_the_placement_gate() -> Result<(), ServerError> {
2113 let registry = ConnectedWorkerRegistry::default();
2114 let (tx, _rx) = mpsc::channel(1);
2115 let _registration = registry
2116 .accept_registration(
2117 &guard(),
2118 &caller("plain"),
2119 ®istration_full(&["plain"], "", "", &["charge"]),
2120 tx,
2121 )
2122 .await?;
2123 assert_eq!(
2124 registry
2125 .workers_for("plain", DEFAULT_TASK_QUEUE, "charge", None)?
2126 .len(),
2127 1,
2128 "with no minter the placement gate is a no-op — registration is unchanged"
2129 );
2130 Ok(())
2131 }
2132
2133 #[tokio::test]
2134 async fn concurrent_registrations_for_a_new_namespace_create_exactly_one_record()
2135 -> Result<(), ServerError> {
2136 let store = namespace_store();
2137 let registry = minting_registry(&store, AutoCreate::Open);
2138
2139 let mut handles = Vec::new();
2140 for _ in 0..8 {
2141 let registry = registry.clone();
2142 handles.push(tokio::spawn(async move {
2143 let (tx, rx) = mpsc::channel(1);
2144 let outcome = registry
2145 .accept_registration(
2146 &guard(),
2147 &caller("rush"),
2148 ®istration("rush", &["charge"]),
2149 tx,
2150 )
2151 .await;
2152 drop(rx);
2154 outcome.map(|registration| registration.worker_id())
2155 }));
2156 }
2157 for handle in handles {
2158 handle
2159 .await
2160 .map_err(|_| ServerError::lock_poisoned("registration task panicked"))??;
2161 }
2162
2163 let all = store.list_namespaces().await?;
2164 assert_eq!(
2165 all.iter().filter(|r| r.name == "rush").count(),
2166 1,
2167 "racing registrations must converge on exactly one durable record"
2168 );
2169 Ok(())
2170 }
2171
2172 #[tokio::test]
2173 async fn closed_rejects_unknown_namespace_and_does_not_create_it() -> Result<(), ServerError> {
2174 let store = namespace_store();
2175 let registry = minting_registry(&store, AutoCreate::Closed);
2176 let (tx, _rx) = mpsc::channel(1);
2177
2178 let denied = registry
2179 .accept_registration(
2180 &guard(),
2181 &caller("ghost"),
2182 ®istration("ghost", &["charge"]),
2183 tx,
2184 )
2185 .await;
2186 assert!(
2187 matches!(denied, Err(ServerError::Namespace { .. })),
2188 "closed policy must reject an unknown namespace"
2189 );
2190 assert!(
2191 store.get_namespace("ghost").await?.is_none(),
2192 "closed policy must NOT create the namespace it rejected"
2193 );
2194 let tq = DEFAULT_TASK_QUEUE;
2195 assert!(
2196 registry
2197 .workers_for("ghost", tq, "charge", None)?
2198 .is_empty(),
2199 "a rejected registration must not insert a worker"
2200 );
2201 Ok(())
2202 }
2203
2204 #[tokio::test]
2205 async fn closed_admits_a_known_namespace() -> Result<(), ServerError> {
2206 let store = namespace_store();
2207 store
2209 .register_namespace("known", NamespaceOrigin::Explicit)
2210 .await?;
2211 let registry = minting_registry(&store, AutoCreate::Closed);
2212 let (tx, _rx) = mpsc::channel(1);
2213
2214 let _registration = registry
2216 .accept_registration(
2217 &guard(),
2218 &caller("known"),
2219 ®istration("known", &["charge"]),
2220 tx,
2221 )
2222 .await?;
2223 let tq = DEFAULT_TASK_QUEUE;
2224 assert_eq!(
2225 registry.workers_for("known", tq, "charge", None)?.len(),
2226 1,
2227 "a known namespace must register under closed policy"
2228 );
2229 Ok(())
2230 }
2231
2232 #[tokio::test]
2233 async fn no_minter_leaves_registration_untouched() -> Result<(), ServerError> {
2234 let registry = ConnectedWorkerRegistry::default();
2237 let (tx, _rx) = mpsc::channel(1);
2238 let _registration = registry
2239 .accept_registration(
2240 &guard(),
2241 &caller("orders"),
2242 ®istration("orders", &["charge"]),
2243 tx,
2244 )
2245 .await?;
2246 let tq = DEFAULT_TASK_QUEUE;
2247 assert_eq!(registry.workers_for("orders", tq, "charge", None)?.len(), 1);
2248 Ok(())
2249 }
2250
2251 #[test]
2254 fn a_never_served_address_censuses_empty_with_no_poller_age() -> Result<(), ServerError> {
2255 let registry = ConnectedWorkerRegistry::default();
2256 let census = registry.pool_census("default", "general", "greet", None)?;
2257 assert_eq!(census.workers_in_pool, 0);
2258 assert_eq!(census.workers_serving_activity, 0);
2259 assert_eq!(census.compatible_workers, 0);
2260 assert_eq!(census.last_compatible_poller_age, None);
2261 assert!(!census.is_served());
2262 Ok(())
2263 }
2264
2265 #[test]
2269 fn the_census_separates_pool_activity_and_node_coverage() -> Result<(), ServerError> {
2270 let registry = ConnectedWorkerRegistry::default();
2271 let (tx, _rx) = mpsc::channel(1);
2272 let _worker = registry.register_namespaces(
2273 [String::from("default")],
2274 "general",
2275 Some(String::from("n1")),
2276 [String::from("greet")].iter(),
2277 tx,
2278 )?;
2279
2280 let unpinned = registry.pool_census("default", "general", "greet", None)?;
2281 assert_eq!(unpinned.workers_in_pool, 1);
2282 assert_eq!(unpinned.workers_serving_activity, 1);
2283 assert_eq!(unpinned.compatible_workers, 1);
2284 assert_eq!(unpinned.last_compatible_poller_age, Some(Duration::ZERO));
2285
2286 let other_activity = registry.pool_census("default", "general", "settle", None)?;
2289 assert_eq!(other_activity.workers_in_pool, 1);
2290 assert_eq!(other_activity.workers_serving_activity, 0);
2291 assert_eq!(other_activity.compatible_workers, 0);
2292
2293 let wrong_node = registry.pool_census("default", "general", "greet", Some("n2"))?;
2296 assert_eq!(wrong_node.workers_serving_activity, 1);
2297 assert_eq!(wrong_node.compatible_workers, 0);
2298 assert_eq!(wrong_node.last_compatible_poller_age, None);
2299 Ok(())
2300 }
2301
2302 #[test]
2323 fn a_pump_alive_worker_the_server_cannot_reach_is_not_selected() -> Result<(), ServerError> {
2324 const WINDOW: Duration = Duration::from_secs(30);
2325
2326 let registry = ConnectedWorkerRegistry::default();
2327 let tracker = HeartbeatTracker::new(WINDOW);
2328 let (tx, _rx) = mpsc::channel(1);
2329 let worker = registry.register_namespaces(
2330 [String::from("default")],
2331 "general",
2332 None,
2333 [String::from("greet")].iter(),
2334 tx,
2335 )?;
2336 let Some(worker_id) = worker.worker_id() else {
2337 return Err(test_failure("registration carries an id"));
2338 };
2339 let start = Instant::now();
2340 tracker.register_connection(worker_id, start)?;
2341
2342 let publish = |now: Instant| -> Result<(), ServerError> {
2347 registry.set_dispatch_ineligible(
2348 tracker
2349 .unreachable_workers(now)?
2350 .into_iter()
2351 .map(|excluded| excluded.worker_id)
2352 .collect(),
2353 )
2354 };
2355
2356 for _ in 0..DISPATCH_PROBATION_PINGS {
2361 assert!(
2362 tracker.record_dispatch_reachability(worker_id, start)?,
2363 "the worker is tracked while it serves its probation"
2364 );
2365 }
2366 publish(start)?;
2367 assert!(
2368 registry
2369 .select_worker("default", "general", "greet", None)?
2370 .is_some(),
2371 "a worker that has answered a full run of pings must be selectable: this is the \
2372 control, and without it a registry selecting NOBODY would satisfy every assertion \
2373 below"
2374 );
2375
2376 let much_later = start + WINDOW * 4;
2382 assert!(
2383 tracker.record_connection_activity(worker_id, much_later)?,
2384 "the worker is still tracked; its process is plainly alive"
2385 );
2386 assert!(
2387 tracker.record_dispatch_unreachable(worker_id)?,
2388 "the probe fired and went unanswered"
2389 );
2390 publish(much_later)?;
2391
2392 assert!(
2393 registry.is_dispatch_ineligible(worker_id)?,
2394 "a worker the server cannot reach must be marked ineligible however alive its \
2395 process looks"
2396 );
2397 assert!(
2398 registry
2399 .select_worker("default", "general", "greet", None)?
2400 .is_none(),
2401 "an ineligible worker must not be SELECTED: selecting one produces a dispatch that \
2402 can only fail, and on the liminal transport it fails by consuming connection \
2403 capacity — making the unreachability worse"
2404 );
2405
2406 assert!(
2411 tracker.record_dispatch_reachability(worker_id, much_later)?,
2412 "the worker is still tracked"
2413 );
2414 publish(much_later)?;
2415 assert!(
2416 registry.is_dispatch_ineligible(worker_id)?,
2417 "one answer part-way through a fresh probation must NOT restore eligibility: a link \
2418 answering one probe in three would otherwise flap in and out of selection"
2419 );
2420
2421 for _ in 1..DISPATCH_PROBATION_PINGS {
2422 assert!(
2423 tracker.record_dispatch_reachability(worker_id, much_later)?,
2424 "the worker is still tracked"
2425 );
2426 }
2427 publish(much_later)?;
2428 assert!(
2429 !registry.is_dispatch_ineligible(worker_id)?,
2430 "an answered ping must clear the exclusion"
2431 );
2432 assert!(
2433 registry
2434 .select_worker("default", "general", "greet", None)?
2435 .is_some(),
2436 "and the worker must be selectable again"
2437 );
2438 Ok(())
2439 }
2440
2441 #[test]
2445 fn a_departed_worker_leaves_a_last_compatible_poller_age() -> Result<(), ServerError> {
2446 let registry = ConnectedWorkerRegistry::default();
2447 let (tx, _rx) = mpsc::channel(1);
2448 let worker = registry.register_namespaces(
2449 [String::from("default")],
2450 "general",
2451 None,
2452 [String::from("greet")].iter(),
2453 tx,
2454 )?;
2455 worker.deregister()?;
2456
2457 let census = registry.pool_census("default", "general", "greet", None)?;
2458 assert_eq!(census.workers_in_pool, 0);
2459 assert_eq!(census.compatible_workers, 0);
2460 let age = census
2461 .last_compatible_poller_age
2462 .ok_or_else(|| test_failure("a departed worker must leave an age behind"))?;
2463 assert!(
2464 age < Duration::from_secs(60),
2465 "the recorded departure is implausibly old: {age:?}"
2466 );
2467
2468 let (tx, _rx) = mpsc::channel(1);
2471 let _back = registry.register_namespaces(
2472 [String::from("default")],
2473 "general",
2474 None,
2475 [String::from("greet")].iter(),
2476 tx,
2477 )?;
2478 let served = registry.pool_census("default", "general", "greet", None)?;
2479 assert_eq!(served.last_compatible_poller_age, Some(Duration::ZERO));
2480 assert!(served.is_served());
2481 Ok(())
2482 }
2483}