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;
17
18pub use aion_core::DEFAULT_TASK_QUEUE;
29
30pub type WorkerTaskSender = mpsc::Sender<WorkerMessage>;
32
33#[derive(Clone, Debug)]
45pub enum WorkerDelivery {
46 Grpc(WorkerTaskSender),
49 #[cfg(feature = "liminal-transport")]
54 Liminal(crate::worker::liminal_transport::LiminalWorkerDelivery),
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
59pub enum WorkerMessage {
60 ActivityTask(Box<ProtoActivityTask>),
62 DrainRequest,
64}
65
66#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
78pub struct PoolAddress {
79 namespace: String,
80 task_queue: String,
81}
82
83impl PoolAddress {
84 #[must_use]
88 pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
89 let task_queue = task_queue.into();
90 let task_queue = if task_queue.is_empty() {
91 String::from(DEFAULT_TASK_QUEUE)
92 } else {
93 task_queue
94 };
95 Self {
96 namespace: namespace.into(),
97 task_queue,
98 }
99 }
100
101 #[must_use]
103 pub fn namespace(&self) -> &str {
104 &self.namespace
105 }
106
107 #[must_use]
109 pub fn task_queue(&self) -> &str {
110 &self.task_queue
111 }
112}
113
114#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
118struct ActivityKey {
119 pool: PoolAddress,
120 activity_type: String,
121}
122
123impl ActivityKey {
124 fn new(pool: PoolAddress, activity_type: impl Into<String>) -> Self {
125 Self {
126 pool,
127 activity_type: activity_type.into(),
128 }
129 }
130}
131
132type WorkerMap = HashMap<WorkerId, WorkerHandle>;
133type RegistryMap = HashMap<ActivityKey, WorkerMap>;
134
135#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
137pub struct WorkerId(u64);
138
139impl WorkerId {
140 #[must_use]
143 pub const fn from_value(value: u64) -> Self {
144 Self(value)
145 }
146
147 #[must_use]
150 pub const fn value(self) -> u64 {
151 self.0
152 }
153}
154
155#[derive(Clone, Debug)]
163pub struct WorkerHandle {
164 id: WorkerId,
165 namespaces: BTreeSet<String>,
166 task_queue: String,
167 node: Option<String>,
168 activity_types: BTreeSet<String>,
169 delivery: WorkerDelivery,
170 intervention_capabilities: InterventionCapabilities,
176}
177
178impl WorkerHandle {
179 #[must_use]
181 pub const fn id(&self) -> WorkerId {
182 self.id
183 }
184
185 #[must_use]
188 pub const fn namespaces(&self) -> &BTreeSet<String> {
189 &self.namespaces
190 }
191
192 #[must_use]
194 pub fn task_queue(&self) -> &str {
195 &self.task_queue
196 }
197
198 #[must_use]
201 pub fn node(&self) -> Option<&str> {
202 self.node.as_deref()
203 }
204
205 #[must_use]
207 pub fn activity_types(&self) -> &BTreeSet<String> {
208 &self.activity_types
209 }
210
211 #[must_use]
213 pub const fn delivery(&self) -> &WorkerDelivery {
214 &self.delivery
215 }
216
217 #[must_use]
223 pub const fn intervention_capabilities(&self) -> &InterventionCapabilities {
224 &self.intervention_capabilities
225 }
226
227 #[must_use]
234 pub fn sender(&self) -> Option<&WorkerTaskSender> {
235 match &self.delivery {
236 WorkerDelivery::Grpc(sender) => Some(sender),
237 #[cfg(feature = "liminal-transport")]
238 WorkerDelivery::Liminal(_) => None,
239 }
240 }
241}
242
243#[derive(Debug)]
244struct RegistryState {
245 next_worker_id: u64,
246 workers: BTreeMap<WorkerId, WorkerHandle>,
247 by_activity: RegistryMap,
248 rotation: HashMap<ActivityKey, usize>,
251 last_departure: HashMap<ActivityKey, BTreeMap<Option<String>, Instant>>,
258 dispatch_ineligible: BTreeSet<WorkerId>,
271}
272
273impl Default for RegistryState {
274 fn default() -> Self {
275 Self {
276 next_worker_id: 1,
277 workers: BTreeMap::new(),
278 by_activity: HashMap::new(),
279 rotation: HashMap::new(),
280 last_departure: HashMap::new(),
281 dispatch_ineligible: BTreeSet::new(),
282 }
283 }
284}
285
286#[derive(Clone, Debug)]
288pub struct ConnectedWorkerRegistry {
289 inner: Arc<Mutex<RegistryState>>,
290 metrics: Option<Metrics>,
291 cluster_publisher: Option<ClusterEventPublisher>,
295 minter: Option<NamespaceMinter>,
299 worker_arrived: Arc<Notify>,
300}
301
302impl Default for ConnectedWorkerRegistry {
303 fn default() -> Self {
304 Self {
305 inner: Arc::new(Mutex::new(RegistryState::default())),
306 metrics: None,
307 cluster_publisher: None,
308 minter: None,
309 worker_arrived: Arc::new(Notify::new()),
310 }
311 }
312}
313
314impl ConnectedWorkerRegistry {
315 #[must_use]
317 pub fn with_metrics(metrics: Metrics) -> Self {
318 Self {
319 inner: Arc::new(Mutex::new(RegistryState::default())),
320 metrics: Some(metrics),
321 cluster_publisher: None,
322 minter: None,
323 worker_arrived: Arc::new(Notify::new()),
324 }
325 }
326
327 #[must_use]
330 pub fn with_cluster_publisher(mut self, publisher: ClusterEventPublisher) -> Self {
331 self.cluster_publisher = Some(publisher);
332 self
333 }
334
335 #[must_use]
350 pub fn with_namespace_minting(
351 mut self,
352 store: Arc<dyn NamespaceStore>,
353 policy: AutoCreate,
354 ) -> Self {
355 let minter = NamespaceMinter::new(store, policy);
356 let minter = match &self.cluster_publisher {
357 Some(publisher) => minter.with_cluster_publisher(publisher.clone()),
358 None => minter,
359 };
360 self.minter = Some(minter);
361 self
362 }
363
364 #[must_use]
376 pub fn with_namespace_routing(mut self, routing: crate::namespace::NamespaceRouting) -> Self {
377 self.minter = self.minter.map(|minter| minter.with_routing(routing));
378 self
379 }
380
381 pub async fn accept_registration(
387 &self,
388 guard: &NamespaceGuard,
389 caller: &CallerIdentity,
390 registration: &ProtoRegisterWorker,
391 sender: WorkerTaskSender,
392 ) -> Result<WorkerRegistration, ServerError> {
393 guard
399 .scope(caller, &NamespaceOperation::register_worker(registration))
400 .await?;
401 let namespaces = guard.scope_worker_namespaces(caller, ®istration.namespaces)?;
402 self.mint_or_gate_namespaces(&namespaces).await?;
410 let node = optional_node(®istration.node);
411 self.enforce_pinned_placement(&namespaces, node.as_deref())
421 .await?;
422 self.register_namespaces(
423 namespaces,
424 registration.task_queue.clone(),
425 node,
426 registration.activity_types.iter(),
427 sender,
428 )
429 }
430
431 async fn mint_or_gate_namespaces(&self, namespaces: &[String]) -> Result<(), ServerError> {
448 let Some(minter) = &self.minter else {
449 return Ok(());
450 };
451 minter
452 .mint_or_gate(namespaces, NamespaceOrigin::WorkerMint)
453 .await
454 }
455
456 async fn enforce_pinned_placement(
479 &self,
480 namespaces: &[String],
481 node: Option<&str>,
482 ) -> Result<(), ServerError> {
483 let Some(minter) = &self.minter else {
484 return Ok(());
485 };
486 for namespace in namespaces {
487 let NamespacePlacement::Pinned { nodes } = minter.placement_of(namespace).await? else {
488 continue;
489 };
490 let admitted = node.is_some_and(|n| nodes.contains(n));
491 if !admitted {
492 return Err(ServerError::placement_admission_denied(
493 namespace, node, &nodes,
494 ));
495 }
496 }
497 Ok(())
498 }
499
500 pub fn register<'a>(
511 &self,
512 namespace: impl Into<String>,
513 activity_types: impl IntoIterator<Item = &'a String>,
514 sender: WorkerTaskSender,
515 ) -> Result<WorkerRegistration, ServerError> {
516 self.register_namespaces(
517 [namespace.into()],
518 String::from(DEFAULT_TASK_QUEUE),
519 None,
520 activity_types,
521 sender,
522 )
523 }
524
525 pub fn register_pool<'a>(
532 &self,
533 pool: PoolAddress,
534 activity_types: impl IntoIterator<Item = &'a String>,
535 sender: WorkerTaskSender,
536 ) -> Result<WorkerRegistration, ServerError> {
537 let PoolAddress {
538 namespace,
539 task_queue,
540 } = pool;
541 self.register_namespaces([namespace], task_queue, None, activity_types, sender)
542 }
543
544 pub fn register_namespaces<'a>(
556 &self,
557 namespaces: impl IntoIterator<Item = String>,
558 task_queue: impl Into<String>,
559 node: Option<String>,
560 activity_types: impl IntoIterator<Item = &'a String>,
561 sender: WorkerTaskSender,
562 ) -> Result<WorkerRegistration, ServerError> {
563 self.register_delivery(
564 namespaces,
565 task_queue,
566 node,
567 activity_types,
568 WorkerDelivery::Grpc(sender),
569 )
570 }
571
572 pub fn register_delivery<'a>(
585 &self,
586 namespaces: impl IntoIterator<Item = String>,
587 task_queue: impl Into<String>,
588 node: Option<String>,
589 activity_types: impl IntoIterator<Item = &'a String>,
590 delivery: WorkerDelivery,
591 ) -> Result<WorkerRegistration, ServerError> {
592 self.register_delivery_with_capabilities(
593 namespaces,
594 task_queue,
595 node,
596 activity_types,
597 delivery,
598 InterventionCapabilities::none(),
599 )
600 }
601
602 pub fn register_delivery_with_capabilities<'a>(
616 &self,
617 namespaces: impl IntoIterator<Item = String>,
618 task_queue: impl Into<String>,
619 node: Option<String>,
620 activity_types: impl IntoIterator<Item = &'a String>,
621 delivery: WorkerDelivery,
622 intervention_capabilities: InterventionCapabilities,
623 ) -> Result<WorkerRegistration, ServerError> {
624 let namespaces = namespaces.into_iter().collect::<BTreeSet<_>>();
625 let task_queue = task_queue.into();
626 let activity_types = activity_types.into_iter().cloned().collect::<BTreeSet<_>>();
627 let mut state = self.state()?;
628 let worker_id = WorkerId(state.next_worker_id);
629 state.next_worker_id = state.next_worker_id.saturating_add(1);
630
631 let node_for_event = node.clone();
634 let handle = WorkerHandle {
635 id: worker_id,
636 namespaces: namespaces.clone(),
637 task_queue: task_queue.clone(),
638 node,
639 activity_types: activity_types.clone(),
640 delivery,
641 intervention_capabilities,
642 };
643
644 for namespace in &namespaces {
645 let pool = PoolAddress::new(namespace.clone(), task_queue.clone());
646 for activity_type in &activity_types {
647 let key = ActivityKey::new(pool.clone(), activity_type.clone());
648 if let Some(by_node) = state.last_departure.get_mut(&key) {
654 by_node.remove(&handle.node);
655 if by_node.is_empty() {
656 state.last_departure.remove(&key);
657 }
658 }
659 state
660 .by_activity
661 .entry(key)
662 .or_default()
663 .insert(worker_id, handle.clone());
664 }
665 }
666 let transport = transport_of(&handle.delivery);
667 state.workers.insert(worker_id, handle);
668 drop(state);
669
670 if let Some(metrics) = &self.metrics {
671 for namespace in &namespaces {
672 metrics.worker_connected(namespace);
673 }
674 }
675
676 if let Some(publisher) = &self.cluster_publisher {
680 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
681 let task_queue_owned = task_queue.clone();
682 drop(publisher.emit(|meta| ClusterEvent::WorkerConnected {
683 meta,
684 worker_id: worker_id.value().to_string(),
685 namespaces: namespaces_vec,
686 task_queue: task_queue_owned,
687 transport,
688 node: node_for_event,
689 }));
690 }
691
692 self.worker_arrived.notify_waiters();
693
694 Ok(WorkerRegistration {
695 registry: self.clone(),
696 parts: Some(WorkerRegistrationParts {
697 worker_id,
698 namespaces,
699 task_queue,
700 activity_types,
701 }),
702 })
703 }
704
705 pub async fn wait_for_worker(&self) {
711 self.worker_arrived.notified().await;
712 }
713
714 pub fn workers_for(
736 &self,
737 namespace: &str,
738 task_queue: &str,
739 activity_type: &str,
740 node: Option<&str>,
741 ) -> Result<Vec<WorkerHandle>, ServerError> {
742 let mut state = self.state()?;
743 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
744 let mut workers: Vec<WorkerHandle> = state
745 .by_activity
746 .get(&key)
747 .map(|workers| {
748 workers
749 .values()
750 .filter(|worker| worker_matches_node(worker, node))
751 .cloned()
752 .collect()
753 })
754 .unwrap_or_default();
755 if workers.is_empty() {
756 return Ok(workers);
757 }
758 workers.sort_by_key(WorkerHandle::id);
759 let idx = state.rotation.entry(key).or_insert(0);
760 let start = *idx % workers.len();
761 *idx = idx.wrapping_add(1);
762 let mut rotated = Vec::with_capacity(workers.len());
763 rotated.extend_from_slice(&workers[start..]);
764 rotated.extend_from_slice(&workers[..start]);
765 Ok(rotated)
766 }
767
768 pub fn all_workers(&self) -> Result<Vec<WorkerHandle>, ServerError> {
774 let state = self.state()?;
775 Ok(state.workers.values().cloned().collect())
776 }
777
778 pub fn worker_by_id(&self, worker_id: WorkerId) -> Result<Option<WorkerHandle>, ServerError> {
790 Ok(self.state()?.workers.get(&worker_id).cloned())
791 }
792
793 pub fn set_intervention_capabilities(
804 &self,
805 worker_id: WorkerId,
806 capabilities: &InterventionCapabilities,
807 ) -> Result<bool, ServerError> {
808 let mut state = self.state()?;
809 if !state.workers.contains_key(&worker_id) {
810 return Ok(false);
811 }
812 if let Some(handle) = state.workers.get_mut(&worker_id) {
813 handle.intervention_capabilities = capabilities.clone();
814 }
815 for workers in state.by_activity.values_mut() {
818 if let Some(handle) = workers.get_mut(&worker_id) {
819 handle.intervention_capabilities = capabilities.clone();
820 }
821 }
822 Ok(true)
823 }
824
825 pub fn broadcast_drain(&self) -> Result<usize, ServerError> {
835 let workers = self.all_workers()?;
836 let mut delivered = 0usize;
837 for worker in workers {
838 if self.drain_worker(worker.id())? {
839 delivered = delivered.saturating_add(1);
840 }
841 }
842 Ok(delivered)
843 }
844
845 pub fn drain_worker(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
855 let worker = {
856 let mut state = self.state()?;
857 let Some(worker) = state.workers.get(&worker_id).cloned() else {
858 return Ok(false);
859 };
860 Self::remove_worker_from_service(&mut state, &worker);
861 worker
862 };
863 match worker.delivery() {
864 WorkerDelivery::Grpc(sender) => {
865 if sender.try_send(WorkerMessage::DrainRequest).is_ok() {
866 tracing::info!(worker_id = worker_id.value(), "worker drain requested");
867 Ok(true)
868 } else {
869 tracing::error!(
870 worker_id = worker_id.value(),
871 "worker drain signal failed; force-deregistering closed transport"
872 );
873 self.deregister(worker_id)?;
874 Ok(false)
875 }
876 }
877 #[cfg(feature = "liminal-transport")]
878 WorkerDelivery::Liminal(delivery) => {
879 tracing::error!(
880 worker_id = worker_id.value(),
881 connection_pid = delivery.pid(),
882 "liminal transport has no drain control channel; worker fenced and \
883 deregistered at drain start"
884 );
885 self.deregister(worker_id)?;
886 Ok(false)
887 }
888 }
889 }
890
891 pub fn select_worker(
901 &self,
902 namespace: &str,
903 task_queue: &str,
904 activity_type: &str,
905 node: Option<&str>,
906 ) -> Result<Option<WorkerHandle>, ServerError> {
907 let state = self.state()?;
908 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
909 Ok(state.by_activity.get(&key).and_then(|workers| {
910 workers
911 .values()
912 .filter(|worker| worker_matches_node(worker, node))
913 .filter(|worker| !state.dispatch_ineligible.contains(&worker.id))
920 .min_by_key(|worker| worker.id)
921 .cloned()
922 }))
923 }
924
925 pub fn set_dispatch_ineligible(
937 &self,
938 unreachable: BTreeSet<WorkerId>,
939 ) -> Result<(), ServerError> {
940 self.state()?.dispatch_ineligible = unreachable;
941 Ok(())
942 }
943
944 pub fn is_dispatch_ineligible(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
951 Ok(self.state()?.dispatch_ineligible.contains(&worker_id))
952 }
953
954 pub fn pool_census(
970 &self,
971 namespace: &str,
972 task_queue: &str,
973 activity_type: &str,
974 node: Option<&str>,
975 ) -> Result<super::queue_service::PoolCensus, ServerError> {
976 let state = self.state()?;
977 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
978 let workers_in_pool = state
979 .workers
980 .values()
981 .filter(|worker| {
982 worker.task_queue == task_queue && worker.namespaces.contains(namespace)
983 })
984 .count();
985 let serving = state.by_activity.get(&key);
986 let workers_serving_activity = serving.map_or(0, HashMap::len);
987 let compatible_workers = serving.map_or(0, |workers| {
988 workers
989 .values()
990 .filter(|worker| worker_matches_node(worker, node))
991 .count()
992 });
993 let last_compatible_poller_age = if compatible_workers > 0 {
994 Some(Duration::ZERO)
995 } else {
996 state
997 .last_departure
998 .get(&key)
999 .and_then(|by_node| {
1000 by_node
1001 .iter()
1002 .filter(|(departed_node, _)| match node {
1003 None => true,
1004 Some(node) => departed_node.as_deref() == Some(node),
1005 })
1006 .map(|(_, departed_at)| *departed_at)
1007 .max()
1008 })
1009 .map(|departed_at| departed_at.elapsed())
1010 };
1011 Ok(super::queue_service::PoolCensus {
1012 workers_in_pool,
1013 workers_serving_activity,
1014 compatible_workers,
1015 last_compatible_poller_age,
1016 })
1017 }
1018
1019 pub fn is_registered(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1030 Ok(self.state()?.workers.contains_key(&worker_id))
1031 }
1032
1033 pub fn deregister(&self, worker_id: WorkerId) -> Result<(), ServerError> {
1045 self.deregister_with_reason(worker_id, WorkerDeathReason::Disconnect)
1046 }
1047
1048 pub fn deregister_with_reason(
1056 &self,
1057 worker_id: WorkerId,
1058 reason: WorkerDeathReason,
1059 ) -> Result<(), ServerError> {
1060 let mut state = self.state()?;
1061 let removed_namespaces = Self::remove_worker(&mut state, worker_id);
1062 drop(state);
1063
1064 let Some(namespaces) = removed_namespaces else {
1065 return Ok(());
1067 };
1068
1069 if let Some(metrics) = &self.metrics {
1070 for namespace in &namespaces {
1071 metrics.worker_disconnected(namespace);
1072 }
1073 }
1074 self.emit_worker_disconnected(worker_id, &namespaces, reason);
1075
1076 Ok(())
1077 }
1078
1079 fn emit_worker_disconnected(
1081 &self,
1082 worker_id: WorkerId,
1083 namespaces: &BTreeSet<String>,
1084 reason: WorkerDeathReason,
1085 ) {
1086 if let Some(publisher) = &self.cluster_publisher {
1087 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
1088 drop(publisher.emit(|meta| ClusterEvent::WorkerDisconnected {
1089 meta,
1090 worker_id: worker_id.value().to_string(),
1091 namespaces: namespaces_vec,
1092 reason,
1093 }));
1094 }
1095 }
1096
1097 fn remove_worker(state: &mut RegistryState, worker_id: WorkerId) -> Option<BTreeSet<String>> {
1101 let handle = state.workers.remove(&worker_id)?;
1102 Self::remove_worker_from_service(state, &handle);
1103 Some(handle.namespaces)
1104 }
1105
1106 fn remove_worker_from_service(state: &mut RegistryState, handle: &WorkerHandle) {
1109 let departed_at = Instant::now();
1110 for namespace in &handle.namespaces {
1111 let pool = PoolAddress::new(namespace.clone(), handle.task_queue.clone());
1112 for activity_type in &handle.activity_types {
1113 let key = ActivityKey::new(pool.clone(), activity_type.clone());
1114 state
1118 .last_departure
1119 .entry(key.clone())
1120 .or_default()
1121 .insert(handle.node.clone(), departed_at);
1122 if let Some(workers) = state.by_activity.get_mut(&key) {
1123 workers.remove(&handle.id);
1124 if workers.is_empty() {
1125 state.by_activity.remove(&key);
1126 state.rotation.remove(&key);
1133 }
1134 }
1135 }
1136 }
1137 }
1138
1139 fn state(&self) -> Result<MutexGuard<'_, RegistryState>, ServerError> {
1140 self.inner
1141 .lock()
1142 .map_err(|_| ServerError::lock_poisoned("connected worker registry"))
1143 }
1144}
1145
1146const fn transport_of(delivery: &WorkerDelivery) -> WorkerTransport {
1149 match delivery {
1150 WorkerDelivery::Grpc(_) => WorkerTransport::Grpc,
1151 #[cfg(feature = "liminal-transport")]
1152 WorkerDelivery::Liminal(_) => WorkerTransport::Liminal,
1153 }
1154}
1155
1156pub(crate) fn optional_node(node: &str) -> Option<String> {
1166 if node.is_empty() {
1167 None
1168 } else {
1169 Some(node.to_owned())
1170 }
1171}
1172
1173fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
1178 match node {
1179 None => true,
1180 Some(node) => worker.node() == Some(node),
1181 }
1182}
1183
1184#[derive(Clone, Debug)]
1185struct WorkerRegistrationParts {
1186 worker_id: WorkerId,
1187 namespaces: BTreeSet<String>,
1188 task_queue: String,
1189 activity_types: BTreeSet<String>,
1190}
1191
1192#[derive(Debug)]
1197pub struct WorkerRegistration {
1198 registry: ConnectedWorkerRegistry,
1199 parts: Option<WorkerRegistrationParts>,
1200}
1201
1202impl WorkerRegistration {
1203 #[must_use]
1205 pub fn worker_id(&self) -> Option<WorkerId> {
1206 self.parts.as_ref().map(|parts| parts.worker_id)
1207 }
1208
1209 #[must_use]
1211 pub fn namespaces(&self) -> Option<&BTreeSet<String>> {
1212 self.parts.as_ref().map(|parts| &parts.namespaces)
1213 }
1214
1215 #[must_use]
1217 pub fn task_queue(&self) -> Option<&str> {
1218 self.parts.as_ref().map(|parts| parts.task_queue.as_str())
1219 }
1220
1221 #[must_use]
1223 pub fn activity_types(&self) -> Option<&BTreeSet<String>> {
1224 self.parts.as_ref().map(|parts| &parts.activity_types)
1225 }
1226
1227 pub fn deregister(mut self) -> Result<(), ServerError> {
1233 let Some(parts) = self.parts.take() else {
1234 return Ok(());
1235 };
1236 self.registry.deregister(parts.worker_id)
1237 }
1238}
1239
1240impl Drop for WorkerRegistration {
1241 fn drop(&mut self) {
1242 let Some(parts) = self.parts.take() else {
1243 return;
1244 };
1245 let removed_namespaces = self.registry.inner.lock().ok().and_then(|mut state| {
1246 ConnectedWorkerRegistry::remove_worker(&mut state, parts.worker_id)
1247 });
1248 if let Some(namespaces) = removed_namespaces {
1249 if let Some(metrics) = &self.registry.metrics {
1250 for namespace in &namespaces {
1251 metrics.worker_disconnected(namespace);
1252 }
1253 }
1254 self.registry.emit_worker_disconnected(
1258 parts.worker_id,
1259 &namespaces,
1260 WorkerDeathReason::Disconnect,
1261 );
1262 }
1263 }
1264}
1265
1266#[cfg(test)]
1267mod tests {
1268 use crate::config::NamespaceMode;
1269 use crate::namespace::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
1270 use crate::worker::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
1271
1272 use super::*;
1273
1274 fn guard() -> NamespaceGuard {
1275 NamespaceGuard::new(NamespaceResolver::authorization_only(
1276 NamespaceMode::SharedEngine,
1277 StaticWorkflowNamespaces::default(),
1278 StaticScheduleNamespaces::default(),
1279 ))
1280 }
1281
1282 fn caller(namespace: &str) -> CallerIdentity {
1283 CallerIdentity::new("worker", [namespace.to_owned()])
1284 }
1285
1286 fn test_failure(message: &str) -> ServerError {
1289 ServerError::worker_dispatch("default".to_owned(), "test".to_owned(), message.to_owned())
1290 }
1291
1292 fn registration(namespace: &str, activity_types: &[&str]) -> ProtoRegisterWorker {
1293 registration_with_queue(namespace, "", activity_types)
1294 }
1295
1296 fn registration_with_queue(
1297 namespace: &str,
1298 task_queue: &str,
1299 activity_types: &[&str],
1300 ) -> ProtoRegisterWorker {
1301 registration_full(&[namespace], task_queue, "", activity_types)
1302 }
1303
1304 fn registration_full(
1305 namespaces: &[&str],
1306 task_queue: &str,
1307 node: &str,
1308 activity_types: &[&str],
1309 ) -> ProtoRegisterWorker {
1310 ProtoRegisterWorker {
1311 namespaces: namespaces.iter().map(|value| (*value).to_owned()).collect(),
1312 activity_types: activity_types
1313 .iter()
1314 .map(|value| (*value).to_owned())
1315 .collect(),
1316 task_queue: task_queue.to_owned(),
1317 node: node.to_owned(),
1318 activities: Vec::new(),
1319 identity: String::new(),
1320 }
1321 }
1322
1323 fn multi_caller(namespaces: &[&str]) -> CallerIdentity {
1324 CallerIdentity::new("worker", namespaces.iter().map(|value| (*value).to_owned()))
1325 }
1326
1327 #[tokio::test]
1333 async fn set_intervention_capabilities_updates_live_worker() -> Result<(), ServerError> {
1334 let registry = ConnectedWorkerRegistry::default();
1335 let (sender, _receiver) = mpsc::channel(1);
1336 let types = ["scout".to_owned()];
1337 let guard = registry.register_delivery_with_capabilities(
1338 ["default".to_owned()],
1339 "default",
1340 None,
1341 types.iter(),
1342 WorkerDelivery::Grpc(sender),
1343 InterventionCapabilities::none(),
1344 )?;
1345 let Some(worker_id) = guard.worker_id() else {
1346 return Err(test_failure("registration carries an id"));
1347 };
1348
1349 let announced = InterventionCapabilities {
1350 supported: vec![aion_core::InterventionPrimitive::InjectMessage],
1351 };
1352 assert!(
1353 registry.set_intervention_capabilities(worker_id, &announced)?,
1354 "a live worker's capabilities must be updatable"
1355 );
1356 let Some(handle) = registry.worker_by_id(worker_id)? else {
1357 return Err(test_failure("worker stays registered"));
1358 };
1359 assert_eq!(handle.intervention_capabilities(), &announced);
1360
1361 assert!(
1362 !registry.set_intervention_capabilities(WorkerId(u64::MAX), &announced)?,
1363 "an unknown worker reports false, never an error"
1364 );
1365 Ok(())
1366 }
1367
1368 #[tokio::test]
1369 async fn register_and_deregister_are_namespace_isolated() -> Result<(), ServerError> {
1370 let registry = ConnectedWorkerRegistry::default();
1371 let (tenant_a_tx, _tenant_a_rx) = mpsc::channel(1);
1372 let (tenant_b_tx, _tenant_b_rx) = mpsc::channel(1);
1373
1374 let tenant_a = registry
1375 .accept_registration(
1376 &guard(),
1377 &caller("tenant-a"),
1378 ®istration("tenant-a", &["charge", "charge"]),
1379 tenant_a_tx,
1380 )
1381 .await?;
1382 let tenant_b = registry
1383 .accept_registration(
1384 &guard(),
1385 &caller("tenant-b"),
1386 ®istration("tenant-b", &["charge"]),
1387 tenant_b_tx,
1388 )
1389 .await?;
1390
1391 let tq = DEFAULT_TASK_QUEUE;
1392 assert_eq!(
1393 registry.workers_for("tenant-a", tq, "charge", None)?.len(),
1394 1
1395 );
1396 assert_eq!(
1397 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1398 1
1399 );
1400 assert!(
1401 registry
1402 .workers_for("tenant-a", tq, "missing", None)?
1403 .is_empty()
1404 );
1405
1406 let tenant_a_id = tenant_a.worker_id();
1407 tenant_a.deregister()?;
1408
1409 assert!(
1410 registry
1411 .workers_for("tenant-a", tq, "charge", None)?
1412 .is_empty()
1413 );
1414 assert_eq!(
1415 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1416 1
1417 );
1418 assert_ne!(tenant_a_id, tenant_b.worker_id());
1419
1420 tenant_b.deregister()?;
1421 assert!(
1422 registry
1423 .workers_for("tenant-b", tq, "charge", None)?
1424 .is_empty()
1425 );
1426 Ok(())
1427 }
1428
1429 #[tokio::test]
1430 async fn denied_namespace_is_not_registered() -> Result<(), ServerError> {
1431 let registry = ConnectedWorkerRegistry::default();
1432 let (tx, _rx) = mpsc::channel(1);
1433 let denied = registry
1434 .accept_registration(
1435 &guard(),
1436 &caller("tenant-a"),
1437 ®istration("tenant-b", &["charge"]),
1438 tx,
1439 )
1440 .await;
1441
1442 assert!(denied.is_err());
1443 assert!(
1444 registry
1445 .workers_for("tenant-b", DEFAULT_TASK_QUEUE, "charge", None)?
1446 .is_empty()
1447 );
1448 Ok(())
1449 }
1450
1451 #[tokio::test]
1452 async fn task_queues_partition_disjoint_pools_within_one_namespace() -> Result<(), ServerError>
1453 {
1454 let registry = ConnectedWorkerRegistry::default();
1458 let (norn_tx, _norn_rx) = mpsc::channel(1);
1459 let (claude_a_tx, _claude_a_rx) = mpsc::channel(1);
1460 let (claude_b_tx, _claude_b_rx) = mpsc::channel(1);
1461
1462 let norn = registry
1463 .accept_registration(
1464 &guard(),
1465 &caller("local"),
1466 ®istration_with_queue("local", "norn", &["dev"]),
1467 norn_tx,
1468 )
1469 .await?;
1470 let claude_a = registry
1472 .accept_registration(
1473 &guard(),
1474 &caller("local"),
1475 ®istration_with_queue("local", "claude", &["dev"]),
1476 claude_a_tx,
1477 )
1478 .await?;
1479 let claude_b = registry
1480 .accept_registration(
1481 &guard(),
1482 &caller("local"),
1483 ®istration_with_queue("local", "claude", &["dev"]),
1484 claude_b_tx,
1485 )
1486 .await?;
1487
1488 let norn_pool = registry.workers_for("local", "norn", "dev", None)?;
1489 assert_eq!(norn_pool.len(), 1, "norn pool has exactly its one worker");
1490 let norn_id = norn.worker_id().ok_or_else(missing_id)?;
1491 assert_eq!(norn_pool[0].id(), norn_id);
1492
1493 let claude_pool = registry.workers_for("local", "claude", "dev", None)?;
1494 assert_eq!(
1495 claude_pool.len(),
1496 2,
1497 "claude pool sees only its two workers"
1498 );
1499 let claude_ids: BTreeSet<WorkerId> = claude_pool.iter().map(WorkerHandle::id).collect();
1500 assert!(
1501 !claude_ids.contains(&norn_id),
1502 "the norn worker must never appear in the claude pool"
1503 );
1504
1505 assert!(
1508 !registry
1509 .workers_for("local", "norn", "dev", None)?
1510 .iter()
1511 .any(|worker| claude_ids.contains(&worker.id()))
1512 );
1513
1514 let first = registry.workers_for("local", "claude", "dev", None)?[0].id();
1518 let second = registry.workers_for("local", "claude", "dev", None)?[0].id();
1519 assert_ne!(
1520 first, second,
1521 "claude pool round-robins across both workers"
1522 );
1523 assert_eq!(
1524 registry.workers_for("local", "norn", "dev", None)?[0].id(),
1525 norn_id,
1526 "the norn pool rotation is unaffected by claude traffic"
1527 );
1528
1529 norn.deregister()?;
1530 claude_a.deregister()?;
1531 claude_b.deregister()?;
1532 Ok(())
1533 }
1534
1535 #[tokio::test]
1536 async fn same_task_queue_in_different_namespaces_is_isolated() -> Result<(), ServerError> {
1537 let registry = ConnectedWorkerRegistry::default();
1540 let (local_tx, _local_rx) = mpsc::channel(1);
1541 let (remote_tx, _remote_rx) = mpsc::channel(1);
1542
1543 let local = registry
1544 .accept_registration(
1545 &guard(),
1546 &caller("local"),
1547 ®istration_with_queue("local", "gpu", &["render"]),
1548 local_tx,
1549 )
1550 .await?;
1551 let remote = registry
1552 .accept_registration(
1553 &guard(),
1554 &caller("remote"),
1555 ®istration_with_queue("remote", "gpu", &["render"]),
1556 remote_tx,
1557 )
1558 .await?;
1559
1560 let local_pool = registry.workers_for("local", "gpu", "render", None)?;
1561 let remote_pool = registry.workers_for("remote", "gpu", "render", None)?;
1562 assert_eq!(local_pool.len(), 1);
1563 assert_eq!(remote_pool.len(), 1);
1564 assert_ne!(
1565 local_pool[0].id(),
1566 remote_pool[0].id(),
1567 "a shared task_queue string does not merge two namespaces"
1568 );
1569
1570 local.deregister()?;
1571 assert!(
1572 registry
1573 .workers_for("local", "gpu", "render", None)?
1574 .is_empty(),
1575 "deregistering the local worker leaves the remote namespace untouched"
1576 );
1577 assert_eq!(
1578 registry.workers_for("remote", "gpu", "render", None)?.len(),
1579 1
1580 );
1581
1582 remote.deregister()?;
1583 Ok(())
1584 }
1585
1586 #[tokio::test]
1587 async fn worker_serving_a_namespace_set_is_reachable_in_each() -> Result<(), ServerError> {
1588 let registry = ConnectedWorkerRegistry::default();
1591 let (ab_tx, _ab_rx) = mpsc::channel(1);
1592 let (a_tx, _a_rx) = mpsc::channel(1);
1593
1594 let worker_ab = registry
1595 .accept_registration(
1596 &guard(),
1597 &multi_caller(&["a", "b"]),
1598 ®istration_full(&["a", "b"], "default", "", &["dev"]),
1599 ab_tx,
1600 )
1601 .await?;
1602 let worker_a = registry
1603 .accept_registration(
1604 &guard(),
1605 &caller("a"),
1606 ®istration_full(&["a"], "default", "", &["dev"]),
1607 a_tx,
1608 )
1609 .await?;
1610
1611 let in_a = registry.workers_for("a", "default", "dev", None)?;
1612 let in_b = registry.workers_for("b", "default", "dev", None)?;
1613 let both_id = worker_ab.worker_id().ok_or_else(missing_id)?;
1614 let only_a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1615
1616 let a_ids: BTreeSet<WorkerId> = in_a.iter().map(WorkerHandle::id).collect();
1618 assert_eq!(a_ids, BTreeSet::from([both_id, only_a_id]));
1619 assert_eq!(in_b.len(), 1, "only the {{a, b}} worker is reachable in b");
1620 assert_eq!(in_b[0].id(), both_id);
1621 assert!(
1622 !in_b.iter().any(|worker| worker.id() == only_a_id),
1623 "the {{a}}-only worker must not be reachable in b"
1624 );
1625
1626 worker_ab.deregister()?;
1628 assert!(
1629 registry
1630 .workers_for("b", "default", "dev", None)?
1631 .is_empty()
1632 );
1633 assert_eq!(registry.workers_for("a", "default", "dev", None)?.len(), 1);
1634
1635 worker_a.deregister()?;
1636 Ok(())
1637 }
1638
1639 #[tokio::test]
1640 async fn node_pin_filters_within_pool() -> Result<(), ServerError> {
1641 let registry = ConnectedWorkerRegistry::default();
1645 let (n1_tx, _n1_rx) = mpsc::channel(1);
1646 let (n2_tx, _n2_rx) = mpsc::channel(1);
1647
1648 let on_n1 = registry
1649 .accept_registration(
1650 &guard(),
1651 &caller("ns"),
1652 ®istration_full(&["ns"], "tq", "n1", &["dev"]),
1653 n1_tx,
1654 )
1655 .await?;
1656 let on_n2 = registry
1657 .accept_registration(
1658 &guard(),
1659 &caller("ns"),
1660 ®istration_full(&["ns"], "tq", "n2", &["dev"]),
1661 n2_tx,
1662 )
1663 .await?;
1664 let n1_id = on_n1.worker_id().ok_or_else(missing_id)?;
1665 let n2_id = on_n2.worker_id().ok_or_else(missing_id)?;
1666
1667 let unpinned = registry.workers_for("ns", "tq", "dev", None)?;
1669 assert_eq!(unpinned.len(), 2, "unpinned reaches the whole pool");
1670 let first = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1671 let second = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1672 assert_ne!(first, second, "unpinned round-robins across both nodes");
1673
1674 let pinned_n1 = registry.workers_for("ns", "tq", "dev", Some("n1"))?;
1676 assert_eq!(pinned_n1.len(), 1);
1677 assert_eq!(pinned_n1[0].id(), n1_id);
1678 let pinned_n2 = registry.workers_for("ns", "tq", "dev", Some("n2"))?;
1679 assert_eq!(pinned_n2.len(), 1);
1680 assert_eq!(pinned_n2[0].id(), n2_id);
1681
1682 assert_eq!(
1684 registry
1685 .select_worker("ns", "tq", "dev", Some("n1"))?
1686 .map(|worker| worker.id()),
1687 Some(n1_id)
1688 );
1689
1690 assert!(
1693 registry
1694 .workers_for("ns", "tq", "dev", Some("absent"))?
1695 .is_empty(),
1696 "a pin to a node with no worker yields no candidate"
1697 );
1698 assert!(
1699 registry
1700 .select_worker("ns", "tq", "dev", Some("absent"))?
1701 .is_none()
1702 );
1703
1704 on_n1.deregister()?;
1705 on_n2.deregister()?;
1706 Ok(())
1707 }
1708
1709 #[tokio::test]
1710 async fn shared_node_id_round_robins_across_workers() -> Result<(), ServerError> {
1711 let registry = ConnectedWorkerRegistry::default();
1714 let (a_tx, _a_rx) = mpsc::channel(1);
1715 let (b_tx, _b_rx) = mpsc::channel(1);
1716
1717 let worker_a = registry
1718 .accept_registration(
1719 &guard(),
1720 &caller("ns"),
1721 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1722 a_tx,
1723 )
1724 .await?;
1725 let worker_b = registry
1726 .accept_registration(
1727 &guard(),
1728 &caller("ns"),
1729 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1730 b_tx,
1731 )
1732 .await?;
1733 let a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1734 let b_id = worker_b.worker_id().ok_or_else(missing_id)?;
1735
1736 let pinned = registry.workers_for("ns", "tq", "dev", Some("shared"))?;
1737 assert_eq!(
1738 pinned.len(),
1739 2,
1740 "both workers on the shared node are candidates"
1741 );
1742 let pinned_ids: BTreeSet<WorkerId> = pinned.iter().map(WorkerHandle::id).collect();
1743 assert_eq!(pinned_ids, BTreeSet::from([a_id, b_id]));
1744
1745 let first = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1746 let second = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1747 assert_ne!(
1748 first, second,
1749 "a pin to a shared node round-robins across both workers on it"
1750 );
1751
1752 worker_a.deregister()?;
1753 worker_b.deregister()?;
1754 Ok(())
1755 }
1756
1757 #[tokio::test]
1758 async fn rotation_cursor_is_pruned_when_last_worker_leaves() -> Result<(), ServerError> {
1759 let registry = ConnectedWorkerRegistry::default();
1764 let (tx, _rx) = mpsc::channel(1);
1765 let worker = registry
1766 .accept_registration(
1767 &guard(),
1768 &caller("ns"),
1769 ®istration_full(&["ns"], "tq", "", &["dev"]),
1770 tx,
1771 )
1772 .await?;
1773
1774 let _ = registry.workers_for("ns", "tq", "dev", None)?;
1776 let key = ActivityKey::new(PoolAddress::new("ns", "tq"), "dev");
1777 assert!(
1778 registry.state()?.rotation.contains_key(&key),
1779 "a lookup must have created the rotation cursor"
1780 );
1781
1782 worker.deregister()?;
1783 let state = registry.state()?;
1784 assert!(
1785 !state.rotation.contains_key(&key),
1786 "the rotation cursor must be pruned once the last worker leaves"
1787 );
1788 assert!(
1789 !state.by_activity.contains_key(&key),
1790 "the activity bucket must also be gone"
1791 );
1792 Ok(())
1793 }
1794
1795 fn missing_id() -> ServerError {
1796 ServerError::lock_poisoned("registration unexpectedly missing a worker id")
1797 }
1798
1799 fn namespace_store() -> Arc<dyn NamespaceStore> {
1802 Arc::new(aion_store::InMemoryStore::default())
1803 }
1804
1805 fn minting_registry(
1806 store: &Arc<dyn NamespaceStore>,
1807 policy: AutoCreate,
1808 ) -> ConnectedWorkerRegistry {
1809 ConnectedWorkerRegistry::default().with_namespace_minting(Arc::clone(store), policy)
1810 }
1811
1812 #[tokio::test]
1813 async fn open_register_mints_durable_record_and_is_idempotent() -> Result<(), ServerError> {
1814 let store = namespace_store();
1815 let registry = minting_registry(&store, AutoCreate::Open);
1816
1817 let (tx_one, _rx_one) = mpsc::channel(1);
1819 let first = registry
1820 .accept_registration(
1821 &guard(),
1822 &caller("orders"),
1823 ®istration("orders", &["charge"]),
1824 tx_one,
1825 )
1826 .await?;
1827 let record = store
1828 .get_namespace("orders")
1829 .await?
1830 .ok_or_else(|| ServerError::namespace_denied("expected a minted record"))?;
1831 assert_eq!(record.name, "orders");
1832 assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
1833
1834 let (tx_two, _rx_two) = mpsc::channel(1);
1837 registry
1838 .accept_registration(
1839 &guard(),
1840 &caller("orders"),
1841 ®istration("orders", &["refund"]),
1842 tx_two,
1843 )
1844 .await?;
1845 let all = store.list_namespaces().await?;
1846 assert_eq!(
1847 all.iter().filter(|r| r.name == "orders").count(),
1848 1,
1849 "re-register must not create a duplicate namespace row"
1850 );
1851 drop(first);
1852 Ok(())
1853 }
1854
1855 #[tokio::test]
1856 async fn open_register_mints_each_namespace_in_a_multi_namespace_worker()
1857 -> Result<(), ServerError> {
1858 let store = namespace_store();
1859 let registry = minting_registry(&store, AutoCreate::Open);
1860 let (tx, _rx) = mpsc::channel(1);
1861
1862 registry
1863 .accept_registration(
1864 &guard(),
1865 &multi_caller(&["alpha", "beta"]),
1866 ®istration_full(&["alpha", "beta"], "", "", &["charge"]),
1867 tx,
1868 )
1869 .await?;
1870
1871 assert!(store.get_namespace("alpha").await?.is_some());
1872 assert!(store.get_namespace("beta").await?.is_some());
1873 Ok(())
1874 }
1875
1876 async fn pinned_registry(
1882 store: &Arc<dyn NamespaceStore>,
1883 namespace: &str,
1884 nodes: &[&str],
1885 ) -> Result<ConnectedWorkerRegistry, ServerError> {
1886 store
1887 .register_namespace(namespace, NamespaceOrigin::Explicit)
1888 .await?;
1889 store
1890 .set_namespace_placement(
1891 namespace,
1892 NamespacePlacement::Pinned {
1893 nodes: nodes.iter().map(|n| (*n).to_owned()).collect(),
1894 },
1895 )
1896 .await?;
1897 Ok(minting_registry(store, AutoCreate::Open))
1898 }
1899
1900 #[tokio::test]
1903 async fn pinned_admits_a_worker_on_a_required_node() -> Result<(), ServerError> {
1904 let store = namespace_store();
1905 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1906 let (tx, _rx) = mpsc::channel(1);
1907
1908 let _registration = registry
1909 .accept_registration(
1910 &guard(),
1911 &caller("iso"),
1912 ®istration_full(&["iso"], "", "n1", &["charge"]),
1913 tx,
1914 )
1915 .await?;
1916
1917 assert_eq!(
1918 registry
1919 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", Some("n1"))?
1920 .len(),
1921 1,
1922 "an n1 worker must be admitted into the Pinned{{n1}} namespace's pool"
1923 );
1924 Ok(())
1925 }
1926
1927 #[tokio::test]
1931 async fn pinned_rejects_a_wrong_node_worker() -> Result<(), ServerError> {
1932 let store = namespace_store();
1933 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1934 let (tx, _rx) = mpsc::channel(1);
1935
1936 let denied = registry
1937 .accept_registration(
1938 &guard(),
1939 &caller("iso"),
1940 ®istration_full(&["iso"], "", "n2", &["charge"]),
1941 tx,
1942 )
1943 .await;
1944 assert!(
1945 matches!(denied, Err(ServerError::Namespace { .. })),
1946 "a wrong-node (n2) worker must be rejected from a Pinned{{n1}} namespace"
1947 );
1948 assert!(
1949 registry
1950 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
1951 .is_empty(),
1952 "a rejected registration must not insert a worker on any node"
1953 );
1954 Ok(())
1955 }
1956
1957 #[tokio::test]
1960 async fn pinned_rejects_a_node_less_worker() -> Result<(), ServerError> {
1961 let store = namespace_store();
1962 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1963 let (tx, _rx) = mpsc::channel(1);
1964
1965 let denied = registry
1966 .accept_registration(
1967 &guard(),
1968 &caller("iso"),
1969 ®istration_full(&["iso"], "", "", &["charge"]),
1970 tx,
1971 )
1972 .await;
1973 assert!(
1974 matches!(denied, Err(ServerError::Namespace { .. })),
1975 "a node-less worker must be rejected from a Pinned{{n1}} namespace"
1976 );
1977 assert!(
1978 registry
1979 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
1980 .is_empty(),
1981 "a rejected node-less registration must not insert a worker"
1982 );
1983 Ok(())
1984 }
1985
1986 #[tokio::test]
1990 async fn pinned_violation_rejects_the_whole_multi_namespace_registration()
1991 -> Result<(), ServerError> {
1992 let store = namespace_store();
1993 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
1994 let (tx, _rx) = mpsc::channel(1);
1995
1996 let denied = registry
1997 .accept_registration(
1998 &guard(),
1999 &multi_caller(&["free", "iso"]),
2000 ®istration_full(&["free", "iso"], "", "n2", &["charge"]),
2001 tx,
2002 )
2003 .await;
2004 assert!(
2005 matches!(denied, Err(ServerError::Namespace { .. })),
2006 "a wrong-node worker serving a Pinned namespace fails the WHOLE registration"
2007 );
2008 assert!(
2009 registry
2010 .workers_for("free", DEFAULT_TASK_QUEUE, "charge", None)?
2011 .is_empty(),
2012 "the compliant namespace must NOT be partially admitted"
2013 );
2014 Ok(())
2015 }
2016
2017 #[tokio::test]
2020 async fn unplaced_and_prefer_admission_is_unaffected_by_the_pinned_gate()
2021 -> Result<(), ServerError> {
2022 let store = namespace_store();
2023 store
2026 .register_namespace("pref", NamespaceOrigin::Explicit)
2027 .await?;
2028 store
2029 .set_namespace_placement(
2030 "pref",
2031 NamespacePlacement::Prefer {
2032 nodes: ["n1".to_owned()].into_iter().collect(),
2033 },
2034 )
2035 .await?;
2036 let registry = minting_registry(&store, AutoCreate::Open);
2037
2038 let (tx_a, _rx_a) = mpsc::channel(1);
2039 let _reg_a = registry
2040 .accept_registration(
2041 &guard(),
2042 &caller("unpl"),
2043 ®istration_full(&["unpl"], "", "", &["charge"]),
2044 tx_a,
2045 )
2046 .await?;
2047 let (tx_b, _rx_b) = mpsc::channel(1);
2048 let _reg_b = registry
2049 .accept_registration(
2050 &guard(),
2051 &caller("pref"),
2052 ®istration_full(&["pref"], "", "", &["charge"]),
2053 tx_b,
2054 )
2055 .await?;
2056
2057 assert_eq!(
2058 registry
2059 .workers_for("unpl", DEFAULT_TASK_QUEUE, "charge", None)?
2060 .len(),
2061 1,
2062 "an Unplaced namespace admits a node-less worker unchanged"
2063 );
2064 assert_eq!(
2065 registry
2066 .workers_for("pref", DEFAULT_TASK_QUEUE, "charge", None)?
2067 .len(),
2068 1,
2069 "a Prefer namespace admits a node-less worker unchanged (only Pinned gates)"
2070 );
2071 Ok(())
2072 }
2073
2074 #[tokio::test]
2078 async fn no_minter_registry_skips_the_placement_gate() -> Result<(), ServerError> {
2079 let registry = ConnectedWorkerRegistry::default();
2080 let (tx, _rx) = mpsc::channel(1);
2081 let _registration = registry
2082 .accept_registration(
2083 &guard(),
2084 &caller("plain"),
2085 ®istration_full(&["plain"], "", "", &["charge"]),
2086 tx,
2087 )
2088 .await?;
2089 assert_eq!(
2090 registry
2091 .workers_for("plain", DEFAULT_TASK_QUEUE, "charge", None)?
2092 .len(),
2093 1,
2094 "with no minter the placement gate is a no-op — registration is unchanged"
2095 );
2096 Ok(())
2097 }
2098
2099 #[tokio::test]
2100 async fn concurrent_registrations_for_a_new_namespace_create_exactly_one_record()
2101 -> Result<(), ServerError> {
2102 let store = namespace_store();
2103 let registry = minting_registry(&store, AutoCreate::Open);
2104
2105 let mut handles = Vec::new();
2106 for _ in 0..8 {
2107 let registry = registry.clone();
2108 handles.push(tokio::spawn(async move {
2109 let (tx, rx) = mpsc::channel(1);
2110 let outcome = registry
2111 .accept_registration(
2112 &guard(),
2113 &caller("rush"),
2114 ®istration("rush", &["charge"]),
2115 tx,
2116 )
2117 .await;
2118 drop(rx);
2120 outcome.map(|registration| registration.worker_id())
2121 }));
2122 }
2123 for handle in handles {
2124 handle
2125 .await
2126 .map_err(|_| ServerError::lock_poisoned("registration task panicked"))??;
2127 }
2128
2129 let all = store.list_namespaces().await?;
2130 assert_eq!(
2131 all.iter().filter(|r| r.name == "rush").count(),
2132 1,
2133 "racing registrations must converge on exactly one durable record"
2134 );
2135 Ok(())
2136 }
2137
2138 #[tokio::test]
2139 async fn closed_rejects_unknown_namespace_and_does_not_create_it() -> Result<(), ServerError> {
2140 let store = namespace_store();
2141 let registry = minting_registry(&store, AutoCreate::Closed);
2142 let (tx, _rx) = mpsc::channel(1);
2143
2144 let denied = registry
2145 .accept_registration(
2146 &guard(),
2147 &caller("ghost"),
2148 ®istration("ghost", &["charge"]),
2149 tx,
2150 )
2151 .await;
2152 assert!(
2153 matches!(denied, Err(ServerError::Namespace { .. })),
2154 "closed policy must reject an unknown namespace"
2155 );
2156 assert!(
2157 store.get_namespace("ghost").await?.is_none(),
2158 "closed policy must NOT create the namespace it rejected"
2159 );
2160 let tq = DEFAULT_TASK_QUEUE;
2161 assert!(
2162 registry
2163 .workers_for("ghost", tq, "charge", None)?
2164 .is_empty(),
2165 "a rejected registration must not insert a worker"
2166 );
2167 Ok(())
2168 }
2169
2170 #[tokio::test]
2171 async fn closed_admits_a_known_namespace() -> Result<(), ServerError> {
2172 let store = namespace_store();
2173 store
2175 .register_namespace("known", NamespaceOrigin::Explicit)
2176 .await?;
2177 let registry = minting_registry(&store, AutoCreate::Closed);
2178 let (tx, _rx) = mpsc::channel(1);
2179
2180 let _registration = registry
2182 .accept_registration(
2183 &guard(),
2184 &caller("known"),
2185 ®istration("known", &["charge"]),
2186 tx,
2187 )
2188 .await?;
2189 let tq = DEFAULT_TASK_QUEUE;
2190 assert_eq!(
2191 registry.workers_for("known", tq, "charge", None)?.len(),
2192 1,
2193 "a known namespace must register under closed policy"
2194 );
2195 Ok(())
2196 }
2197
2198 #[tokio::test]
2199 async fn no_minter_leaves_registration_untouched() -> Result<(), ServerError> {
2200 let registry = ConnectedWorkerRegistry::default();
2203 let (tx, _rx) = mpsc::channel(1);
2204 let _registration = registry
2205 .accept_registration(
2206 &guard(),
2207 &caller("orders"),
2208 ®istration("orders", &["charge"]),
2209 tx,
2210 )
2211 .await?;
2212 let tq = DEFAULT_TASK_QUEUE;
2213 assert_eq!(registry.workers_for("orders", tq, "charge", None)?.len(), 1);
2214 Ok(())
2215 }
2216
2217 #[test]
2220 fn a_never_served_address_censuses_empty_with_no_poller_age() -> Result<(), ServerError> {
2221 let registry = ConnectedWorkerRegistry::default();
2222 let census = registry.pool_census("default", "general", "greet", None)?;
2223 assert_eq!(census.workers_in_pool, 0);
2224 assert_eq!(census.workers_serving_activity, 0);
2225 assert_eq!(census.compatible_workers, 0);
2226 assert_eq!(census.last_compatible_poller_age, None);
2227 assert!(!census.is_served());
2228 Ok(())
2229 }
2230
2231 #[test]
2235 fn the_census_separates_pool_activity_and_node_coverage() -> Result<(), ServerError> {
2236 let registry = ConnectedWorkerRegistry::default();
2237 let (tx, _rx) = mpsc::channel(1);
2238 let _worker = registry.register_namespaces(
2239 [String::from("default")],
2240 "general",
2241 Some(String::from("n1")),
2242 [String::from("greet")].iter(),
2243 tx,
2244 )?;
2245
2246 let unpinned = registry.pool_census("default", "general", "greet", None)?;
2247 assert_eq!(unpinned.workers_in_pool, 1);
2248 assert_eq!(unpinned.workers_serving_activity, 1);
2249 assert_eq!(unpinned.compatible_workers, 1);
2250 assert_eq!(unpinned.last_compatible_poller_age, Some(Duration::ZERO));
2251
2252 let other_activity = registry.pool_census("default", "general", "settle", None)?;
2255 assert_eq!(other_activity.workers_in_pool, 1);
2256 assert_eq!(other_activity.workers_serving_activity, 0);
2257 assert_eq!(other_activity.compatible_workers, 0);
2258
2259 let wrong_node = registry.pool_census("default", "general", "greet", Some("n2"))?;
2262 assert_eq!(wrong_node.workers_serving_activity, 1);
2263 assert_eq!(wrong_node.compatible_workers, 0);
2264 assert_eq!(wrong_node.last_compatible_poller_age, None);
2265 Ok(())
2266 }
2267
2268 #[test]
2289 fn a_pump_alive_worker_the_server_cannot_reach_is_not_selected() -> Result<(), ServerError> {
2290 const WINDOW: Duration = Duration::from_secs(30);
2291
2292 let registry = ConnectedWorkerRegistry::default();
2293 let tracker = HeartbeatTracker::new(WINDOW);
2294 let (tx, _rx) = mpsc::channel(1);
2295 let worker = registry.register_namespaces(
2296 [String::from("default")],
2297 "general",
2298 None,
2299 [String::from("greet")].iter(),
2300 tx,
2301 )?;
2302 let Some(worker_id) = worker.worker_id() else {
2303 return Err(test_failure("registration carries an id"));
2304 };
2305 let start = Instant::now();
2306 tracker.register_connection(worker_id, start)?;
2307
2308 for _ in 0..DISPATCH_PROBATION_PINGS {
2313 assert!(
2314 tracker.record_dispatch_reachability(worker_id, start)?,
2315 "the worker is tracked while it serves its probation"
2316 );
2317 }
2318 registry
2319 .set_dispatch_ineligible(tracker.unreachable_workers(start)?.into_iter().collect())?;
2320 assert!(
2321 registry
2322 .select_worker("default", "general", "greet", None)?
2323 .is_some(),
2324 "a worker that has answered a full run of pings must be selectable: this is the \
2325 control, and without it a registry selecting NOBODY would satisfy every assertion \
2326 below"
2327 );
2328
2329 let much_later = start + WINDOW * 4;
2335 assert!(
2336 tracker.record_connection_activity(worker_id, much_later)?,
2337 "the worker is still tracked; its process is plainly alive"
2338 );
2339 assert!(
2340 tracker.record_dispatch_unreachable(worker_id)?,
2341 "the probe fired and went unanswered"
2342 );
2343 registry.set_dispatch_ineligible(
2344 tracker
2345 .unreachable_workers(much_later)?
2346 .into_iter()
2347 .collect(),
2348 )?;
2349
2350 assert!(
2351 registry.is_dispatch_ineligible(worker_id)?,
2352 "a worker the server cannot reach must be marked ineligible however alive its \
2353 process looks"
2354 );
2355 assert!(
2356 registry
2357 .select_worker("default", "general", "greet", None)?
2358 .is_none(),
2359 "an ineligible worker must not be SELECTED: selecting one produces a dispatch that \
2360 can only fail, and on the liminal transport it fails by consuming connection \
2361 capacity — making the unreachability worse"
2362 );
2363
2364 assert!(
2369 tracker.record_dispatch_reachability(worker_id, much_later)?,
2370 "the worker is still tracked"
2371 );
2372 registry.set_dispatch_ineligible(
2373 tracker
2374 .unreachable_workers(much_later)?
2375 .into_iter()
2376 .collect(),
2377 )?;
2378 assert!(
2379 registry.is_dispatch_ineligible(worker_id)?,
2380 "one answer part-way through a fresh probation must NOT restore eligibility: a link \
2381 answering one probe in three would otherwise flap in and out of selection"
2382 );
2383
2384 for _ in 1..DISPATCH_PROBATION_PINGS {
2385 assert!(
2386 tracker.record_dispatch_reachability(worker_id, much_later)?,
2387 "the worker is still tracked"
2388 );
2389 }
2390 registry.set_dispatch_ineligible(
2391 tracker
2392 .unreachable_workers(much_later)?
2393 .into_iter()
2394 .collect(),
2395 )?;
2396 assert!(
2397 !registry.is_dispatch_ineligible(worker_id)?,
2398 "an answered ping must clear the exclusion"
2399 );
2400 assert!(
2401 registry
2402 .select_worker("default", "general", "greet", None)?
2403 .is_some(),
2404 "and the worker must be selectable again"
2405 );
2406 Ok(())
2407 }
2408
2409 #[test]
2413 fn a_departed_worker_leaves_a_last_compatible_poller_age() -> Result<(), ServerError> {
2414 let registry = ConnectedWorkerRegistry::default();
2415 let (tx, _rx) = mpsc::channel(1);
2416 let worker = registry.register_namespaces(
2417 [String::from("default")],
2418 "general",
2419 None,
2420 [String::from("greet")].iter(),
2421 tx,
2422 )?;
2423 worker.deregister()?;
2424
2425 let census = registry.pool_census("default", "general", "greet", None)?;
2426 assert_eq!(census.workers_in_pool, 0);
2427 assert_eq!(census.compatible_workers, 0);
2428 let age = census
2429 .last_compatible_poller_age
2430 .ok_or_else(|| test_failure("a departed worker must leave an age behind"))?;
2431 assert!(
2432 age < Duration::from_secs(60),
2433 "the recorded departure is implausibly old: {age:?}"
2434 );
2435
2436 let (tx, _rx) = mpsc::channel(1);
2439 let _back = registry.register_namespaces(
2440 [String::from("default")],
2441 "general",
2442 None,
2443 [String::from("greet")].iter(),
2444 tx,
2445 )?;
2446 let served = registry.pool_census("default", "general", "greet", None)?;
2447 assert_eq!(served.last_compatible_poller_age, Some(Duration::ZERO));
2448 assert!(served.is_served());
2449 Ok(())
2450 }
2451}