1use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::sync::{Arc, Mutex, MutexGuard};
5use std::time::{Duration, Instant};
6
7use aion_core::{
8 ClusterEvent, DeploymentAssociation, InterventionCapabilities, WorkerDeathReason,
9 WorkerTransport,
10};
11use aion_proto::{ProtoActivityTask, ProtoRegisterWorker};
12use aion_store::{NamespaceOrigin, NamespacePlacement, NamespaceStore, WorkerDeploymentStore};
13use tokio::sync::{Notify, mpsc};
14
15use crate::cluster_publisher::ClusterEventPublisher;
16use crate::config::AutoCreate;
17use crate::error::ServerError;
18use crate::namespace::{CallerIdentity, NamespaceGuard, NamespaceMinter, NamespaceOperation};
19use crate::observability::Metrics;
20use crate::worker::admission_audit::AdmissionAudit;
21
22pub use aion_core::DEFAULT_TASK_QUEUE;
33
34pub type WorkerTaskSender = mpsc::Sender<WorkerMessage>;
36
37#[derive(Clone, Debug)]
49pub enum WorkerDelivery {
50 Grpc(WorkerTaskSender),
53 #[cfg(feature = "liminal-transport")]
58 Liminal(crate::worker::liminal_transport::LiminalWorkerDelivery),
59}
60
61#[derive(Clone, Debug, Eq, PartialEq)]
63pub enum WorkerMessage {
64 ActivityTask(Box<ProtoActivityTask>),
66 DrainRequest,
68}
69
70#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
82pub struct PoolAddress {
83 namespace: String,
84 task_queue: String,
85}
86
87impl PoolAddress {
88 #[must_use]
92 pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
93 let task_queue = task_queue.into();
94 let task_queue = if task_queue.is_empty() {
95 String::from(DEFAULT_TASK_QUEUE)
96 } else {
97 task_queue
98 };
99 Self {
100 namespace: namespace.into(),
101 task_queue,
102 }
103 }
104
105 #[must_use]
107 pub fn namespace(&self) -> &str {
108 &self.namespace
109 }
110
111 #[must_use]
113 pub fn task_queue(&self) -> &str {
114 &self.task_queue
115 }
116}
117
118#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
122struct ActivityKey {
123 pool: PoolAddress,
124 activity_type: String,
125}
126
127impl ActivityKey {
128 fn new(pool: PoolAddress, activity_type: impl Into<String>) -> Self {
129 Self {
130 pool,
131 activity_type: activity_type.into(),
132 }
133 }
134}
135
136type WorkerMap = HashMap<WorkerId, WorkerHandle>;
137type RegistryMap = HashMap<ActivityKey, WorkerMap>;
138
139#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
141pub struct WorkerId(u64);
142
143impl WorkerId {
144 #[must_use]
147 pub const fn from_value(value: u64) -> Self {
148 Self(value)
149 }
150
151 #[must_use]
154 pub const fn value(self) -> u64 {
155 self.0
156 }
157}
158
159#[derive(Clone, Debug, PartialEq, Eq)]
161pub struct WorkerInstanceIdentity {
162 pub deployment: String,
164 pub instance_id: String,
166 pub association: DeploymentAssociation,
168}
169
170struct RegistrationOptions {
171 instance: Option<WorkerInstanceIdentity>,
172 intervention_capabilities: InterventionCapabilities,
173}
174
175#[derive(Clone, Debug)]
183pub struct WorkerHandle {
184 id: WorkerId,
185 namespaces: BTreeSet<String>,
186 task_queue: String,
187 node: Option<String>,
188 activity_types: BTreeSet<String>,
189 instance: Option<WorkerInstanceIdentity>,
190 delivery: WorkerDelivery,
191 intervention_capabilities: InterventionCapabilities,
197}
198
199impl WorkerHandle {
200 #[must_use]
202 pub const fn id(&self) -> WorkerId {
203 self.id
204 }
205
206 #[must_use]
209 pub const fn namespaces(&self) -> &BTreeSet<String> {
210 &self.namespaces
211 }
212
213 #[must_use]
215 pub fn task_queue(&self) -> &str {
216 &self.task_queue
217 }
218
219 #[must_use]
222 pub fn node(&self) -> Option<&str> {
223 self.node.as_deref()
224 }
225
226 #[must_use]
228 pub fn activity_types(&self) -> &BTreeSet<String> {
229 &self.activity_types
230 }
231
232 #[must_use]
234 pub const fn instance(&self) -> Option<&WorkerInstanceIdentity> {
235 self.instance.as_ref()
236 }
237
238 #[must_use]
240 pub const fn delivery(&self) -> &WorkerDelivery {
241 &self.delivery
242 }
243
244 #[must_use]
250 pub const fn intervention_capabilities(&self) -> &InterventionCapabilities {
251 &self.intervention_capabilities
252 }
253
254 #[must_use]
261 pub fn sender(&self) -> Option<&WorkerTaskSender> {
262 match &self.delivery {
263 WorkerDelivery::Grpc(sender) => Some(sender),
264 #[cfg(feature = "liminal-transport")]
265 WorkerDelivery::Liminal(_) => None,
266 }
267 }
268}
269
270#[derive(Debug)]
271struct RegistryState {
272 next_worker_id: u64,
273 workers: BTreeMap<WorkerId, WorkerHandle>,
274 by_activity: RegistryMap,
275 rotation: HashMap<ActivityKey, usize>,
278 last_departure: HashMap<ActivityKey, BTreeMap<Option<String>, Instant>>,
285 dispatch_ineligible: BTreeSet<WorkerId>,
298}
299
300impl Default for RegistryState {
301 fn default() -> Self {
302 Self {
303 next_worker_id: 1,
304 workers: BTreeMap::new(),
305 by_activity: HashMap::new(),
306 rotation: HashMap::new(),
307 last_departure: HashMap::new(),
308 dispatch_ineligible: BTreeSet::new(),
309 }
310 }
311}
312
313#[derive(Clone)]
315pub struct ConnectedWorkerRegistry {
316 inner: Arc<Mutex<RegistryState>>,
317 metrics: Option<Metrics>,
318 cluster_publisher: Option<ClusterEventPublisher>,
322 minter: Option<NamespaceMinter>,
326 deployment_store: Option<Arc<dyn WorkerDeploymentStore>>,
328 worker_arrived: Arc<Notify>,
329 audit: Arc<AdmissionAudit>,
337}
338
339impl std::fmt::Debug for ConnectedWorkerRegistry {
340 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
341 formatter
342 .debug_struct("ConnectedWorkerRegistry")
343 .field(
344 "deployment_store_attached",
345 &self.deployment_store.is_some(),
346 )
347 .finish_non_exhaustive()
348 }
349}
350
351impl Default for ConnectedWorkerRegistry {
352 fn default() -> Self {
353 Self {
354 inner: Arc::new(Mutex::new(RegistryState::default())),
355 metrics: None,
356 cluster_publisher: None,
357 minter: None,
358 deployment_store: None,
359 worker_arrived: Arc::new(Notify::new()),
360 audit: Arc::new(AdmissionAudit::new()),
361 }
362 }
363}
364
365impl ConnectedWorkerRegistry {
366 #[must_use]
368 pub fn with_metrics(metrics: Metrics) -> Self {
369 Self {
370 inner: Arc::new(Mutex::new(RegistryState::default())),
371 metrics: Some(metrics),
372 cluster_publisher: None,
373 minter: None,
374 deployment_store: None,
375 worker_arrived: Arc::new(Notify::new()),
376 audit: Arc::new(AdmissionAudit::new()),
377 }
378 }
379
380 #[must_use]
384 pub fn admission_audit(&self) -> &AdmissionAudit {
385 &self.audit
386 }
387
388 #[must_use]
391 pub fn with_cluster_publisher(mut self, publisher: ClusterEventPublisher) -> Self {
392 self.cluster_publisher = Some(publisher);
393 self
394 }
395
396 #[must_use]
398 pub fn with_worker_deployment_store(mut self, store: Arc<dyn WorkerDeploymentStore>) -> Self {
399 self.deployment_store = Some(store);
400 self
401 }
402
403 #[must_use]
418 pub fn with_namespace_minting(
419 mut self,
420 store: Arc<dyn NamespaceStore>,
421 policy: AutoCreate,
422 ) -> Self {
423 let minter = NamespaceMinter::new(store, policy);
424 let minter = match &self.cluster_publisher {
425 Some(publisher) => minter.with_cluster_publisher(publisher.clone()),
426 None => minter,
427 };
428 self.minter = Some(minter);
429 self
430 }
431
432 #[must_use]
444 pub fn with_namespace_routing(mut self, routing: crate::namespace::NamespaceRouting) -> Self {
445 self.minter = self.minter.map(|minter| minter.with_routing(routing));
446 self
447 }
448
449 pub async fn accept_registration(
455 &self,
456 guard: &NamespaceGuard,
457 caller: &CallerIdentity,
458 registration: &ProtoRegisterWorker,
459 sender: WorkerTaskSender,
460 ) -> Result<WorkerRegistration, ServerError> {
461 guard
467 .scope(caller, &NamespaceOperation::register_worker(registration))
468 .await?;
469 let namespaces = guard.scope_worker_namespaces(caller, ®istration.namespaces)?;
470 self.mint_or_gate_namespaces(&namespaces).await?;
478 let node = optional_node(®istration.node);
479 self.enforce_pinned_placement(&namespaces, node.as_deref())
489 .await?;
490 let instance = self
491 .resolve_instance_identity(registration.instance.as_ref())
492 .await?;
493 self.register_delivery(
494 namespaces,
495 registration.task_queue.clone(),
496 node,
497 instance,
498 registration.activity_types.iter(),
499 WorkerDelivery::Grpc(sender),
500 )
501 }
502
503 pub async fn resolve_instance_identity(
520 &self,
521 instance: Option<&aion_proto::ProtoWorkerInstanceIdentity>,
522 ) -> Result<Option<WorkerInstanceIdentity>, ServerError> {
523 let Some(instance) = instance else {
524 return Ok(None);
525 };
526 let association = match &self.deployment_store {
527 Some(store) => {
528 if store
529 .get_worker_deployment(&instance.deployment)
530 .await
531 .map_err(ServerError::from)?
532 .is_some()
533 {
534 DeploymentAssociation::Known
535 } else {
536 DeploymentAssociation::Absent
537 }
538 }
539 None => DeploymentAssociation::Unchecked,
540 };
541 Ok(Some(WorkerInstanceIdentity {
542 deployment: instance.deployment.clone(),
543 instance_id: instance.instance_id.clone(),
544 association,
545 }))
546 }
547
548 async fn mint_or_gate_namespaces(&self, namespaces: &[String]) -> Result<(), ServerError> {
565 let Some(minter) = &self.minter else {
566 return Ok(());
567 };
568 minter
569 .mint_or_gate(namespaces, NamespaceOrigin::WorkerMint)
570 .await
571 }
572
573 async fn enforce_pinned_placement(
596 &self,
597 namespaces: &[String],
598 node: Option<&str>,
599 ) -> Result<(), ServerError> {
600 let Some(minter) = &self.minter else {
601 return Ok(());
602 };
603 for namespace in namespaces {
604 let NamespacePlacement::Pinned { nodes } = minter.placement_of(namespace).await? else {
605 continue;
606 };
607 let admitted = node.is_some_and(|n| nodes.contains(n));
608 if !admitted {
609 return Err(ServerError::placement_admission_denied(
610 namespace, node, &nodes,
611 ));
612 }
613 }
614 Ok(())
615 }
616
617 pub fn register<'a>(
628 &self,
629 namespace: impl Into<String>,
630 activity_types: impl IntoIterator<Item = &'a String>,
631 sender: WorkerTaskSender,
632 ) -> Result<WorkerRegistration, ServerError> {
633 self.register_namespaces(
634 [namespace.into()],
635 String::from(DEFAULT_TASK_QUEUE),
636 None,
637 activity_types,
638 sender,
639 )
640 }
641
642 pub fn register_pool<'a>(
649 &self,
650 pool: PoolAddress,
651 activity_types: impl IntoIterator<Item = &'a String>,
652 sender: WorkerTaskSender,
653 ) -> Result<WorkerRegistration, ServerError> {
654 let PoolAddress {
655 namespace,
656 task_queue,
657 } = pool;
658 self.register_namespaces([namespace], task_queue, None, activity_types, sender)
659 }
660
661 pub fn register_namespaces<'a>(
673 &self,
674 namespaces: impl IntoIterator<Item = String>,
675 task_queue: impl Into<String>,
676 node: Option<String>,
677 activity_types: impl IntoIterator<Item = &'a String>,
678 sender: WorkerTaskSender,
679 ) -> Result<WorkerRegistration, ServerError> {
680 self.register_delivery(
681 namespaces,
682 task_queue,
683 node,
684 None,
685 activity_types,
686 WorkerDelivery::Grpc(sender),
687 )
688 }
689
690 pub fn register_delivery<'a>(
703 &self,
704 namespaces: impl IntoIterator<Item = String>,
705 task_queue: impl Into<String>,
706 node: Option<String>,
707 instance: Option<WorkerInstanceIdentity>,
708 activity_types: impl IntoIterator<Item = &'a String>,
709 delivery: WorkerDelivery,
710 ) -> Result<WorkerRegistration, ServerError> {
711 self.register_delivery_inner(
712 namespaces,
713 task_queue,
714 node,
715 activity_types,
716 delivery,
717 RegistrationOptions {
718 instance,
719 intervention_capabilities: InterventionCapabilities::none(),
720 },
721 )
722 }
723
724 pub fn register_delivery_with_capabilities<'a>(
738 &self,
739 namespaces: impl IntoIterator<Item = String>,
740 task_queue: impl Into<String>,
741 node: Option<String>,
742 activity_types: impl IntoIterator<Item = &'a String>,
743 delivery: WorkerDelivery,
744 intervention_capabilities: InterventionCapabilities,
745 ) -> Result<WorkerRegistration, ServerError> {
746 self.register_delivery_inner(
747 namespaces,
748 task_queue,
749 node,
750 activity_types,
751 delivery,
752 RegistrationOptions {
753 instance: None,
754 intervention_capabilities,
755 },
756 )
757 }
758
759 fn register_delivery_inner<'a>(
760 &self,
761 namespaces: impl IntoIterator<Item = String>,
762 task_queue: impl Into<String>,
763 node: Option<String>,
764 activity_types: impl IntoIterator<Item = &'a String>,
765 delivery: WorkerDelivery,
766 options: RegistrationOptions,
767 ) -> Result<WorkerRegistration, ServerError> {
768 let namespaces = namespaces.into_iter().collect::<BTreeSet<_>>();
769 let task_queue = task_queue.into();
770 let activity_types = activity_types.into_iter().cloned().collect::<BTreeSet<_>>();
771 let mut state = self.state()?;
772 let worker_id = WorkerId(state.next_worker_id);
773 state.next_worker_id = state.next_worker_id.saturating_add(1);
774
775 let node_for_event = node.clone();
778 let instance_for_event = options.instance.clone();
779 let handle = WorkerHandle {
780 id: worker_id,
781 namespaces: namespaces.clone(),
782 task_queue: task_queue.clone(),
783 node,
784 activity_types: activity_types.clone(),
785 instance: options.instance,
786 delivery,
787 intervention_capabilities: options.intervention_capabilities,
788 };
789
790 for namespace in &namespaces {
791 let pool = PoolAddress::new(namespace.clone(), task_queue.clone());
792 for activity_type in &activity_types {
793 let key = ActivityKey::new(pool.clone(), activity_type.clone());
794 if let Some(by_node) = state.last_departure.get_mut(&key) {
800 by_node.remove(&handle.node);
801 if by_node.is_empty() {
802 state.last_departure.remove(&key);
803 }
804 }
805 state
806 .by_activity
807 .entry(key)
808 .or_default()
809 .insert(worker_id, handle.clone());
810 }
811 }
812 let transport = transport_of(&handle.delivery);
813 state.workers.insert(worker_id, handle);
814 drop(state);
815
816 if let Some(metrics) = &self.metrics {
817 for namespace in &namespaces {
818 metrics.worker_connected(namespace);
819 }
820 }
821
822 if let Some(publisher) = &self.cluster_publisher {
826 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
827 let task_queue_owned = task_queue.clone();
828 drop(publisher.emit(|meta| {
829 ClusterEvent::WorkerConnected {
830 meta,
831 worker_id: worker_id.value().to_string(),
832 namespaces: namespaces_vec,
833 task_queue: task_queue_owned,
834 transport,
835 node: node_for_event,
836 deployment: instance_for_event
837 .as_ref()
838 .map(|identity| identity.deployment.clone()),
839 deployment_association: instance_for_event.map(|identity| identity.association),
840 }
841 }));
842 }
843
844 self.worker_arrived.notify_waiters();
845
846 Ok(WorkerRegistration {
847 registry: self.clone(),
848 parts: Some(WorkerRegistrationParts {
849 worker_id,
850 namespaces,
851 task_queue,
852 activity_types,
853 }),
854 })
855 }
856
857 pub async fn wait_for_worker(&self) {
863 self.worker_arrived.notified().await;
864 }
865
866 pub fn workers_for(
888 &self,
889 namespace: &str,
890 task_queue: &str,
891 activity_type: &str,
892 node: Option<&str>,
893 ) -> Result<Vec<WorkerHandle>, ServerError> {
894 let mut state = self.state()?;
895 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
896 let mut workers: Vec<WorkerHandle> = state
897 .by_activity
898 .get(&key)
899 .map(|workers| {
900 workers
901 .values()
902 .filter(|worker| worker_matches_node(worker, node))
903 .cloned()
904 .collect()
905 })
906 .unwrap_or_default();
907 if workers.is_empty() {
908 return Ok(workers);
909 }
910 workers.sort_by_key(WorkerHandle::id);
911 let idx = state.rotation.entry(key).or_insert(0);
912 let start = *idx % workers.len();
913 *idx = idx.wrapping_add(1);
914 let mut rotated = Vec::with_capacity(workers.len());
915 rotated.extend_from_slice(&workers[start..]);
916 rotated.extend_from_slice(&workers[..start]);
917 Ok(rotated)
918 }
919
920 pub fn all_workers(&self) -> Result<Vec<WorkerHandle>, ServerError> {
926 let state = self.state()?;
927 Ok(state.workers.values().cloned().collect())
928 }
929
930 pub fn worker_by_id(&self, worker_id: WorkerId) -> Result<Option<WorkerHandle>, ServerError> {
942 Ok(self.state()?.workers.get(&worker_id).cloned())
943 }
944
945 pub fn set_intervention_capabilities(
956 &self,
957 worker_id: WorkerId,
958 capabilities: &InterventionCapabilities,
959 ) -> Result<bool, ServerError> {
960 let mut state = self.state()?;
961 if !state.workers.contains_key(&worker_id) {
962 return Ok(false);
963 }
964 if let Some(handle) = state.workers.get_mut(&worker_id) {
965 handle.intervention_capabilities = capabilities.clone();
966 }
967 for workers in state.by_activity.values_mut() {
970 if let Some(handle) = workers.get_mut(&worker_id) {
971 handle.intervention_capabilities = capabilities.clone();
972 }
973 }
974 Ok(true)
975 }
976
977 pub fn broadcast_drain(&self) -> Result<usize, ServerError> {
987 let workers = self.all_workers()?;
988 let mut delivered = 0usize;
989 for worker in workers {
990 if self.drain_worker(worker.id())? {
991 delivered = delivered.saturating_add(1);
992 }
993 }
994 Ok(delivered)
995 }
996
997 pub fn drain_worker(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1007 let worker = {
1008 let mut state = self.state()?;
1009 let Some(worker) = state.workers.get(&worker_id).cloned() else {
1010 return Ok(false);
1011 };
1012 Self::remove_worker_from_service(&mut state, &worker);
1013 worker
1014 };
1015 match worker.delivery() {
1016 WorkerDelivery::Grpc(sender) => {
1017 if sender.try_send(WorkerMessage::DrainRequest).is_ok() {
1018 tracing::info!(worker_id = worker_id.value(), "worker drain requested");
1019 Ok(true)
1020 } else {
1021 tracing::error!(
1022 worker_id = worker_id.value(),
1023 "worker drain signal failed; force-deregistering closed transport"
1024 );
1025 self.deregister(worker_id)?;
1026 Ok(false)
1027 }
1028 }
1029 #[cfg(feature = "liminal-transport")]
1030 WorkerDelivery::Liminal(delivery) => {
1031 tracing::error!(
1032 worker_id = worker_id.value(),
1033 connection_pid = delivery.pid(),
1034 "liminal transport has no drain control channel; worker fenced and \
1035 deregistered at drain start"
1036 );
1037 self.deregister(worker_id)?;
1038 Ok(false)
1039 }
1040 }
1041 }
1042
1043 pub fn select_worker(
1053 &self,
1054 namespace: &str,
1055 task_queue: &str,
1056 activity_type: &str,
1057 node: Option<&str>,
1058 ) -> Result<Option<WorkerHandle>, ServerError> {
1059 let state = self.state()?;
1060 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1061 Ok(state.by_activity.get(&key).and_then(|workers| {
1062 workers
1063 .values()
1064 .filter(|worker| worker_matches_node(worker, node))
1065 .filter(|worker| !state.dispatch_ineligible.contains(&worker.id))
1072 .min_by_key(|worker| worker.id)
1073 .cloned()
1074 }))
1075 }
1076
1077 pub fn set_dispatch_ineligible(
1089 &self,
1090 unreachable: BTreeSet<WorkerId>,
1091 ) -> Result<(), ServerError> {
1092 self.state()?.dispatch_ineligible = unreachable;
1093 Ok(())
1094 }
1095
1096 pub fn is_dispatch_ineligible(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1103 Ok(self.state()?.dispatch_ineligible.contains(&worker_id))
1104 }
1105
1106 pub fn dispatch_ineligible(&self) -> Result<BTreeSet<WorkerId>, ServerError> {
1118 Ok(self.state()?.dispatch_ineligible.clone())
1119 }
1120
1121 pub fn pool_census(
1137 &self,
1138 namespace: &str,
1139 task_queue: &str,
1140 activity_type: &str,
1141 node: Option<&str>,
1142 ) -> Result<super::queue_service::PoolCensus, ServerError> {
1143 let state = self.state()?;
1144 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1145 let workers_in_pool = state
1146 .workers
1147 .values()
1148 .filter(|worker| {
1149 worker.task_queue == task_queue && worker.namespaces.contains(namespace)
1150 })
1151 .count();
1152 let serving = state.by_activity.get(&key);
1153 let workers_serving_activity = serving.map_or(0, HashMap::len);
1154 let compatible_workers = serving.map_or(0, |workers| {
1155 workers
1156 .values()
1157 .filter(|worker| worker_matches_node(worker, node))
1158 .count()
1159 });
1160 let last_compatible_poller_age = if compatible_workers > 0 {
1161 Some(Duration::ZERO)
1162 } else {
1163 state
1164 .last_departure
1165 .get(&key)
1166 .and_then(|by_node| {
1167 by_node
1168 .iter()
1169 .filter(|(departed_node, _)| match node {
1170 None => true,
1171 Some(node) => departed_node.as_deref() == Some(node),
1172 })
1173 .map(|(_, departed_at)| *departed_at)
1174 .max()
1175 })
1176 .map(|departed_at| departed_at.elapsed())
1177 };
1178 Ok(super::queue_service::PoolCensus {
1179 workers_in_pool,
1180 workers_serving_activity,
1181 compatible_workers,
1182 last_compatible_poller_age,
1183 })
1184 }
1185
1186 pub fn is_registered(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1197 Ok(self.state()?.workers.contains_key(&worker_id))
1198 }
1199
1200 pub fn deregister(&self, worker_id: WorkerId) -> Result<(), ServerError> {
1212 self.deregister_with_reason(worker_id, WorkerDeathReason::Disconnect)
1213 }
1214
1215 pub fn deregister_with_reason(
1223 &self,
1224 worker_id: WorkerId,
1225 reason: WorkerDeathReason,
1226 ) -> Result<(), ServerError> {
1227 let mut state = self.state()?;
1228 let removed_namespaces = Self::remove_worker(&mut state, worker_id);
1229 drop(state);
1230
1231 let Some(namespaces) = removed_namespaces else {
1232 return Ok(());
1234 };
1235
1236 if let Some(metrics) = &self.metrics {
1237 for namespace in &namespaces {
1238 metrics.worker_disconnected(namespace);
1239 }
1240 }
1241 self.emit_worker_disconnected(worker_id, &namespaces, reason);
1242
1243 Ok(())
1244 }
1245
1246 fn emit_worker_disconnected(
1248 &self,
1249 worker_id: WorkerId,
1250 namespaces: &BTreeSet<String>,
1251 reason: WorkerDeathReason,
1252 ) {
1253 if let Some(publisher) = &self.cluster_publisher {
1254 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
1255 drop(publisher.emit(|meta| ClusterEvent::WorkerDisconnected {
1256 meta,
1257 worker_id: worker_id.value().to_string(),
1258 namespaces: namespaces_vec,
1259 reason,
1260 }));
1261 }
1262 }
1263
1264 fn remove_worker(state: &mut RegistryState, worker_id: WorkerId) -> Option<BTreeSet<String>> {
1268 let handle = state.workers.remove(&worker_id)?;
1269 Self::remove_worker_from_service(state, &handle);
1270 Some(handle.namespaces)
1271 }
1272
1273 fn remove_worker_from_service(state: &mut RegistryState, handle: &WorkerHandle) {
1276 let departed_at = Instant::now();
1277 for namespace in &handle.namespaces {
1278 let pool = PoolAddress::new(namespace.clone(), handle.task_queue.clone());
1279 for activity_type in &handle.activity_types {
1280 let key = ActivityKey::new(pool.clone(), activity_type.clone());
1281 state
1285 .last_departure
1286 .entry(key.clone())
1287 .or_default()
1288 .insert(handle.node.clone(), departed_at);
1289 if let Some(workers) = state.by_activity.get_mut(&key) {
1290 workers.remove(&handle.id);
1291 if workers.is_empty() {
1292 state.by_activity.remove(&key);
1293 state.rotation.remove(&key);
1300 }
1301 }
1302 }
1303 }
1304 }
1305
1306 fn state(&self) -> Result<MutexGuard<'_, RegistryState>, ServerError> {
1307 self.inner
1308 .lock()
1309 .map_err(|_| ServerError::lock_poisoned("connected worker registry"))
1310 }
1311}
1312
1313const fn transport_of(delivery: &WorkerDelivery) -> WorkerTransport {
1316 match delivery {
1317 WorkerDelivery::Grpc(_) => WorkerTransport::Grpc,
1318 #[cfg(feature = "liminal-transport")]
1319 WorkerDelivery::Liminal(_) => WorkerTransport::Liminal,
1320 }
1321}
1322
1323pub(crate) fn optional_node(node: &str) -> Option<String> {
1333 if node.is_empty() {
1334 None
1335 } else {
1336 Some(node.to_owned())
1337 }
1338}
1339
1340fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
1345 match node {
1346 None => true,
1347 Some(node) => worker.node() == Some(node),
1348 }
1349}
1350
1351#[derive(Clone, Debug)]
1352struct WorkerRegistrationParts {
1353 worker_id: WorkerId,
1354 namespaces: BTreeSet<String>,
1355 task_queue: String,
1356 activity_types: BTreeSet<String>,
1357}
1358
1359#[derive(Debug)]
1364pub struct WorkerRegistration {
1365 registry: ConnectedWorkerRegistry,
1366 parts: Option<WorkerRegistrationParts>,
1367}
1368
1369impl WorkerRegistration {
1370 #[must_use]
1372 pub fn worker_id(&self) -> Option<WorkerId> {
1373 self.parts.as_ref().map(|parts| parts.worker_id)
1374 }
1375
1376 #[must_use]
1378 pub fn namespaces(&self) -> Option<&BTreeSet<String>> {
1379 self.parts.as_ref().map(|parts| &parts.namespaces)
1380 }
1381
1382 #[must_use]
1384 pub fn task_queue(&self) -> Option<&str> {
1385 self.parts.as_ref().map(|parts| parts.task_queue.as_str())
1386 }
1387
1388 #[must_use]
1390 pub fn activity_types(&self) -> Option<&BTreeSet<String>> {
1391 self.parts.as_ref().map(|parts| &parts.activity_types)
1392 }
1393
1394 pub fn deregister(mut self) -> Result<(), ServerError> {
1400 let Some(parts) = self.parts.take() else {
1401 return Ok(());
1402 };
1403 self.registry.deregister(parts.worker_id)
1404 }
1405}
1406
1407impl Drop for WorkerRegistration {
1408 fn drop(&mut self) {
1409 let Some(parts) = self.parts.take() else {
1410 return;
1411 };
1412 let removed_namespaces = self.registry.inner.lock().ok().and_then(|mut state| {
1413 ConnectedWorkerRegistry::remove_worker(&mut state, parts.worker_id)
1414 });
1415 if let Some(namespaces) = removed_namespaces {
1416 if let Some(metrics) = &self.registry.metrics {
1417 for namespace in &namespaces {
1418 metrics.worker_disconnected(namespace);
1419 }
1420 }
1421 self.registry.emit_worker_disconnected(
1425 parts.worker_id,
1426 &namespaces,
1427 WorkerDeathReason::Disconnect,
1428 );
1429 }
1430 }
1431}
1432
1433#[cfg(test)]
1434mod tests {
1435 use crate::config::NamespaceMode;
1436 use crate::namespace::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
1437 use crate::worker::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
1438
1439 use super::*;
1440
1441 fn guard() -> NamespaceGuard {
1442 NamespaceGuard::new(NamespaceResolver::authorization_only(
1443 NamespaceMode::SharedEngine,
1444 StaticWorkflowNamespaces::default(),
1445 StaticScheduleNamespaces::default(),
1446 ))
1447 }
1448
1449 fn caller(namespace: &str) -> CallerIdentity {
1450 CallerIdentity::new("worker", [namespace.to_owned()])
1451 }
1452
1453 fn test_failure(message: &str) -> ServerError {
1456 ServerError::worker_dispatch("default".to_owned(), "test".to_owned(), message.to_owned())
1457 }
1458
1459 fn registration(namespace: &str, activity_types: &[&str]) -> ProtoRegisterWorker {
1460 registration_with_queue(namespace, "", activity_types)
1461 }
1462
1463 fn registration_with_queue(
1464 namespace: &str,
1465 task_queue: &str,
1466 activity_types: &[&str],
1467 ) -> ProtoRegisterWorker {
1468 registration_full(&[namespace], task_queue, "", activity_types)
1469 }
1470
1471 fn registration_full(
1472 namespaces: &[&str],
1473 task_queue: &str,
1474 node: &str,
1475 activity_types: &[&str],
1476 ) -> ProtoRegisterWorker {
1477 ProtoRegisterWorker {
1478 namespaces: namespaces.iter().map(|value| (*value).to_owned()).collect(),
1479 activity_types: activity_types
1480 .iter()
1481 .map(|value| (*value).to_owned())
1482 .collect(),
1483 task_queue: task_queue.to_owned(),
1484 node: node.to_owned(),
1485 activities: Vec::new(),
1486 identity: String::new(),
1487 instance: None,
1488 }
1489 }
1490
1491 fn multi_caller(namespaces: &[&str]) -> CallerIdentity {
1492 CallerIdentity::new("worker", namespaces.iter().map(|value| (*value).to_owned()))
1493 }
1494
1495 #[tokio::test]
1501 async fn set_intervention_capabilities_updates_live_worker() -> Result<(), ServerError> {
1502 let registry = ConnectedWorkerRegistry::default();
1503 let (sender, _receiver) = mpsc::channel(1);
1504 let types = ["scout".to_owned()];
1505 let guard = registry.register_delivery_with_capabilities(
1506 ["default".to_owned()],
1507 "default",
1508 None,
1509 types.iter(),
1510 WorkerDelivery::Grpc(sender),
1511 InterventionCapabilities::none(),
1512 )?;
1513 let Some(worker_id) = guard.worker_id() else {
1514 return Err(test_failure("registration carries an id"));
1515 };
1516
1517 let announced = InterventionCapabilities {
1518 supported: vec![aion_core::InterventionPrimitive::InjectMessage],
1519 };
1520 assert!(
1521 registry.set_intervention_capabilities(worker_id, &announced)?,
1522 "a live worker's capabilities must be updatable"
1523 );
1524 let Some(handle) = registry.worker_by_id(worker_id)? else {
1525 return Err(test_failure("worker stays registered"));
1526 };
1527 assert_eq!(handle.intervention_capabilities(), &announced);
1528
1529 assert!(
1530 !registry.set_intervention_capabilities(WorkerId(u64::MAX), &announced)?,
1531 "an unknown worker reports false, never an error"
1532 );
1533 Ok(())
1534 }
1535
1536 #[tokio::test]
1537 async fn register_and_deregister_are_namespace_isolated() -> Result<(), ServerError> {
1538 let registry = ConnectedWorkerRegistry::default();
1539 let (tenant_a_tx, _tenant_a_rx) = mpsc::channel(1);
1540 let (tenant_b_tx, _tenant_b_rx) = mpsc::channel(1);
1541
1542 let tenant_a = registry
1543 .accept_registration(
1544 &guard(),
1545 &caller("tenant-a"),
1546 ®istration("tenant-a", &["charge", "charge"]),
1547 tenant_a_tx,
1548 )
1549 .await?;
1550 let tenant_b = registry
1551 .accept_registration(
1552 &guard(),
1553 &caller("tenant-b"),
1554 ®istration("tenant-b", &["charge"]),
1555 tenant_b_tx,
1556 )
1557 .await?;
1558
1559 let tq = DEFAULT_TASK_QUEUE;
1560 assert_eq!(
1561 registry.workers_for("tenant-a", tq, "charge", None)?.len(),
1562 1
1563 );
1564 assert_eq!(
1565 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1566 1
1567 );
1568 assert!(
1569 registry
1570 .workers_for("tenant-a", tq, "missing", None)?
1571 .is_empty()
1572 );
1573
1574 let tenant_a_id = tenant_a.worker_id();
1575 tenant_a.deregister()?;
1576
1577 assert!(
1578 registry
1579 .workers_for("tenant-a", tq, "charge", None)?
1580 .is_empty()
1581 );
1582 assert_eq!(
1583 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1584 1
1585 );
1586 assert_ne!(tenant_a_id, tenant_b.worker_id());
1587
1588 tenant_b.deregister()?;
1589 assert!(
1590 registry
1591 .workers_for("tenant-b", tq, "charge", None)?
1592 .is_empty()
1593 );
1594 Ok(())
1595 }
1596
1597 #[tokio::test]
1598 async fn denied_namespace_is_not_registered() -> Result<(), ServerError> {
1599 let registry = ConnectedWorkerRegistry::default();
1600 let (tx, _rx) = mpsc::channel(1);
1601 let denied = registry
1602 .accept_registration(
1603 &guard(),
1604 &caller("tenant-a"),
1605 ®istration("tenant-b", &["charge"]),
1606 tx,
1607 )
1608 .await;
1609
1610 assert!(denied.is_err());
1611 assert!(
1612 registry
1613 .workers_for("tenant-b", DEFAULT_TASK_QUEUE, "charge", None)?
1614 .is_empty()
1615 );
1616 Ok(())
1617 }
1618
1619 #[tokio::test]
1620 async fn task_queues_partition_disjoint_pools_within_one_namespace() -> Result<(), ServerError>
1621 {
1622 let registry = ConnectedWorkerRegistry::default();
1626 let (norn_tx, _norn_rx) = mpsc::channel(1);
1627 let (claude_a_tx, _claude_a_rx) = mpsc::channel(1);
1628 let (claude_b_tx, _claude_b_rx) = mpsc::channel(1);
1629
1630 let norn = registry
1631 .accept_registration(
1632 &guard(),
1633 &caller("local"),
1634 ®istration_with_queue("local", "norn", &["dev"]),
1635 norn_tx,
1636 )
1637 .await?;
1638 let claude_a = registry
1640 .accept_registration(
1641 &guard(),
1642 &caller("local"),
1643 ®istration_with_queue("local", "claude", &["dev"]),
1644 claude_a_tx,
1645 )
1646 .await?;
1647 let claude_b = registry
1648 .accept_registration(
1649 &guard(),
1650 &caller("local"),
1651 ®istration_with_queue("local", "claude", &["dev"]),
1652 claude_b_tx,
1653 )
1654 .await?;
1655
1656 let norn_pool = registry.workers_for("local", "norn", "dev", None)?;
1657 assert_eq!(norn_pool.len(), 1, "norn pool has exactly its one worker");
1658 let norn_id = norn.worker_id().ok_or_else(missing_id)?;
1659 assert_eq!(norn_pool[0].id(), norn_id);
1660
1661 let claude_pool = registry.workers_for("local", "claude", "dev", None)?;
1662 assert_eq!(
1663 claude_pool.len(),
1664 2,
1665 "claude pool sees only its two workers"
1666 );
1667 let claude_ids: BTreeSet<WorkerId> = claude_pool.iter().map(WorkerHandle::id).collect();
1668 assert!(
1669 !claude_ids.contains(&norn_id),
1670 "the norn worker must never appear in the claude pool"
1671 );
1672
1673 assert!(
1676 !registry
1677 .workers_for("local", "norn", "dev", None)?
1678 .iter()
1679 .any(|worker| claude_ids.contains(&worker.id()))
1680 );
1681
1682 let first = registry.workers_for("local", "claude", "dev", None)?[0].id();
1686 let second = registry.workers_for("local", "claude", "dev", None)?[0].id();
1687 assert_ne!(
1688 first, second,
1689 "claude pool round-robins across both workers"
1690 );
1691 assert_eq!(
1692 registry.workers_for("local", "norn", "dev", None)?[0].id(),
1693 norn_id,
1694 "the norn pool rotation is unaffected by claude traffic"
1695 );
1696
1697 norn.deregister()?;
1698 claude_a.deregister()?;
1699 claude_b.deregister()?;
1700 Ok(())
1701 }
1702
1703 #[tokio::test]
1704 async fn same_task_queue_in_different_namespaces_is_isolated() -> Result<(), ServerError> {
1705 let registry = ConnectedWorkerRegistry::default();
1708 let (local_tx, _local_rx) = mpsc::channel(1);
1709 let (remote_tx, _remote_rx) = mpsc::channel(1);
1710
1711 let local = registry
1712 .accept_registration(
1713 &guard(),
1714 &caller("local"),
1715 ®istration_with_queue("local", "gpu", &["render"]),
1716 local_tx,
1717 )
1718 .await?;
1719 let remote = registry
1720 .accept_registration(
1721 &guard(),
1722 &caller("remote"),
1723 ®istration_with_queue("remote", "gpu", &["render"]),
1724 remote_tx,
1725 )
1726 .await?;
1727
1728 let local_pool = registry.workers_for("local", "gpu", "render", None)?;
1729 let remote_pool = registry.workers_for("remote", "gpu", "render", None)?;
1730 assert_eq!(local_pool.len(), 1);
1731 assert_eq!(remote_pool.len(), 1);
1732 assert_ne!(
1733 local_pool[0].id(),
1734 remote_pool[0].id(),
1735 "a shared task_queue string does not merge two namespaces"
1736 );
1737
1738 local.deregister()?;
1739 assert!(
1740 registry
1741 .workers_for("local", "gpu", "render", None)?
1742 .is_empty(),
1743 "deregistering the local worker leaves the remote namespace untouched"
1744 );
1745 assert_eq!(
1746 registry.workers_for("remote", "gpu", "render", None)?.len(),
1747 1
1748 );
1749
1750 remote.deregister()?;
1751 Ok(())
1752 }
1753
1754 #[tokio::test]
1755 async fn worker_serving_a_namespace_set_is_reachable_in_each() -> Result<(), ServerError> {
1756 let registry = ConnectedWorkerRegistry::default();
1759 let (ab_tx, _ab_rx) = mpsc::channel(1);
1760 let (a_tx, _a_rx) = mpsc::channel(1);
1761
1762 let worker_ab = registry
1763 .accept_registration(
1764 &guard(),
1765 &multi_caller(&["a", "b"]),
1766 ®istration_full(&["a", "b"], "default", "", &["dev"]),
1767 ab_tx,
1768 )
1769 .await?;
1770 let worker_a = registry
1771 .accept_registration(
1772 &guard(),
1773 &caller("a"),
1774 ®istration_full(&["a"], "default", "", &["dev"]),
1775 a_tx,
1776 )
1777 .await?;
1778
1779 let in_a = registry.workers_for("a", "default", "dev", None)?;
1780 let in_b = registry.workers_for("b", "default", "dev", None)?;
1781 let both_id = worker_ab.worker_id().ok_or_else(missing_id)?;
1782 let only_a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1783
1784 let a_ids: BTreeSet<WorkerId> = in_a.iter().map(WorkerHandle::id).collect();
1786 assert_eq!(a_ids, BTreeSet::from([both_id, only_a_id]));
1787 assert_eq!(in_b.len(), 1, "only the {{a, b}} worker is reachable in b");
1788 assert_eq!(in_b[0].id(), both_id);
1789 assert!(
1790 !in_b.iter().any(|worker| worker.id() == only_a_id),
1791 "the {{a}}-only worker must not be reachable in b"
1792 );
1793
1794 worker_ab.deregister()?;
1796 assert!(
1797 registry
1798 .workers_for("b", "default", "dev", None)?
1799 .is_empty()
1800 );
1801 assert_eq!(registry.workers_for("a", "default", "dev", None)?.len(), 1);
1802
1803 worker_a.deregister()?;
1804 Ok(())
1805 }
1806
1807 #[tokio::test]
1808 async fn node_pin_filters_within_pool() -> Result<(), ServerError> {
1809 let registry = ConnectedWorkerRegistry::default();
1813 let (n1_tx, _n1_rx) = mpsc::channel(1);
1814 let (n2_tx, _n2_rx) = mpsc::channel(1);
1815
1816 let on_n1 = registry
1817 .accept_registration(
1818 &guard(),
1819 &caller("ns"),
1820 ®istration_full(&["ns"], "tq", "n1", &["dev"]),
1821 n1_tx,
1822 )
1823 .await?;
1824 let on_n2 = registry
1825 .accept_registration(
1826 &guard(),
1827 &caller("ns"),
1828 ®istration_full(&["ns"], "tq", "n2", &["dev"]),
1829 n2_tx,
1830 )
1831 .await?;
1832 let n1_id = on_n1.worker_id().ok_or_else(missing_id)?;
1833 let n2_id = on_n2.worker_id().ok_or_else(missing_id)?;
1834
1835 let unpinned = registry.workers_for("ns", "tq", "dev", None)?;
1837 assert_eq!(unpinned.len(), 2, "unpinned reaches the whole pool");
1838 let first = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1839 let second = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1840 assert_ne!(first, second, "unpinned round-robins across both nodes");
1841
1842 let pinned_n1 = registry.workers_for("ns", "tq", "dev", Some("n1"))?;
1844 assert_eq!(pinned_n1.len(), 1);
1845 assert_eq!(pinned_n1[0].id(), n1_id);
1846 let pinned_n2 = registry.workers_for("ns", "tq", "dev", Some("n2"))?;
1847 assert_eq!(pinned_n2.len(), 1);
1848 assert_eq!(pinned_n2[0].id(), n2_id);
1849
1850 assert_eq!(
1852 registry
1853 .select_worker("ns", "tq", "dev", Some("n1"))?
1854 .map(|worker| worker.id()),
1855 Some(n1_id)
1856 );
1857
1858 assert!(
1861 registry
1862 .workers_for("ns", "tq", "dev", Some("absent"))?
1863 .is_empty(),
1864 "a pin to a node with no worker yields no candidate"
1865 );
1866 assert!(
1867 registry
1868 .select_worker("ns", "tq", "dev", Some("absent"))?
1869 .is_none()
1870 );
1871
1872 on_n1.deregister()?;
1873 on_n2.deregister()?;
1874 Ok(())
1875 }
1876
1877 #[tokio::test]
1878 async fn shared_node_id_round_robins_across_workers() -> Result<(), ServerError> {
1879 let registry = ConnectedWorkerRegistry::default();
1882 let (a_tx, _a_rx) = mpsc::channel(1);
1883 let (b_tx, _b_rx) = mpsc::channel(1);
1884
1885 let worker_a = registry
1886 .accept_registration(
1887 &guard(),
1888 &caller("ns"),
1889 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1890 a_tx,
1891 )
1892 .await?;
1893 let worker_b = registry
1894 .accept_registration(
1895 &guard(),
1896 &caller("ns"),
1897 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1898 b_tx,
1899 )
1900 .await?;
1901 let a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1902 let b_id = worker_b.worker_id().ok_or_else(missing_id)?;
1903
1904 let pinned = registry.workers_for("ns", "tq", "dev", Some("shared"))?;
1905 assert_eq!(
1906 pinned.len(),
1907 2,
1908 "both workers on the shared node are candidates"
1909 );
1910 let pinned_ids: BTreeSet<WorkerId> = pinned.iter().map(WorkerHandle::id).collect();
1911 assert_eq!(pinned_ids, BTreeSet::from([a_id, b_id]));
1912
1913 let first = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1914 let second = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
1915 assert_ne!(
1916 first, second,
1917 "a pin to a shared node round-robins across both workers on it"
1918 );
1919
1920 worker_a.deregister()?;
1921 worker_b.deregister()?;
1922 Ok(())
1923 }
1924
1925 #[tokio::test]
1926 async fn rotation_cursor_is_pruned_when_last_worker_leaves() -> Result<(), ServerError> {
1927 let registry = ConnectedWorkerRegistry::default();
1932 let (tx, _rx) = mpsc::channel(1);
1933 let worker = registry
1934 .accept_registration(
1935 &guard(),
1936 &caller("ns"),
1937 ®istration_full(&["ns"], "tq", "", &["dev"]),
1938 tx,
1939 )
1940 .await?;
1941
1942 let _ = registry.workers_for("ns", "tq", "dev", None)?;
1944 let key = ActivityKey::new(PoolAddress::new("ns", "tq"), "dev");
1945 assert!(
1946 registry.state()?.rotation.contains_key(&key),
1947 "a lookup must have created the rotation cursor"
1948 );
1949
1950 worker.deregister()?;
1951 let state = registry.state()?;
1952 assert!(
1953 !state.rotation.contains_key(&key),
1954 "the rotation cursor must be pruned once the last worker leaves"
1955 );
1956 assert!(
1957 !state.by_activity.contains_key(&key),
1958 "the activity bucket must also be gone"
1959 );
1960 Ok(())
1961 }
1962
1963 fn missing_id() -> ServerError {
1964 ServerError::lock_poisoned("registration unexpectedly missing a worker id")
1965 }
1966
1967 fn namespace_store() -> Arc<dyn NamespaceStore> {
1970 Arc::new(aion_store::InMemoryStore::default())
1971 }
1972
1973 fn minting_registry(
1974 store: &Arc<dyn NamespaceStore>,
1975 policy: AutoCreate,
1976 ) -> ConnectedWorkerRegistry {
1977 ConnectedWorkerRegistry::default().with_namespace_minting(Arc::clone(store), policy)
1978 }
1979
1980 #[tokio::test]
1981 async fn open_register_mints_durable_record_and_is_idempotent() -> Result<(), ServerError> {
1982 let store = namespace_store();
1983 let registry = minting_registry(&store, AutoCreate::Open);
1984
1985 let (tx_one, _rx_one) = mpsc::channel(1);
1987 let first = registry
1988 .accept_registration(
1989 &guard(),
1990 &caller("orders"),
1991 ®istration("orders", &["charge"]),
1992 tx_one,
1993 )
1994 .await?;
1995 let record = store
1996 .get_namespace("orders")
1997 .await?
1998 .ok_or_else(|| ServerError::namespace_denied("expected a minted record"))?;
1999 assert_eq!(record.name, "orders");
2000 assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
2001
2002 let (tx_two, _rx_two) = mpsc::channel(1);
2005 registry
2006 .accept_registration(
2007 &guard(),
2008 &caller("orders"),
2009 ®istration("orders", &["refund"]),
2010 tx_two,
2011 )
2012 .await?;
2013 let all = store.list_namespaces().await?;
2014 assert_eq!(
2015 all.iter().filter(|r| r.name == "orders").count(),
2016 1,
2017 "re-register must not create a duplicate namespace row"
2018 );
2019 drop(first);
2020 Ok(())
2021 }
2022
2023 #[tokio::test]
2024 async fn open_register_mints_each_namespace_in_a_multi_namespace_worker()
2025 -> Result<(), ServerError> {
2026 let store = namespace_store();
2027 let registry = minting_registry(&store, AutoCreate::Open);
2028 let (tx, _rx) = mpsc::channel(1);
2029
2030 registry
2031 .accept_registration(
2032 &guard(),
2033 &multi_caller(&["alpha", "beta"]),
2034 ®istration_full(&["alpha", "beta"], "", "", &["charge"]),
2035 tx,
2036 )
2037 .await?;
2038
2039 assert!(store.get_namespace("alpha").await?.is_some());
2040 assert!(store.get_namespace("beta").await?.is_some());
2041 Ok(())
2042 }
2043
2044 async fn pinned_registry(
2050 store: &Arc<dyn NamespaceStore>,
2051 namespace: &str,
2052 nodes: &[&str],
2053 ) -> Result<ConnectedWorkerRegistry, ServerError> {
2054 store
2055 .register_namespace(namespace, NamespaceOrigin::Explicit)
2056 .await?;
2057 store
2058 .set_namespace_placement(
2059 namespace,
2060 NamespacePlacement::Pinned {
2061 nodes: nodes.iter().map(|n| (*n).to_owned()).collect(),
2062 },
2063 )
2064 .await?;
2065 Ok(minting_registry(store, AutoCreate::Open))
2066 }
2067
2068 #[tokio::test]
2071 async fn pinned_admits_a_worker_on_a_required_node() -> Result<(), ServerError> {
2072 let store = namespace_store();
2073 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2074 let (tx, _rx) = mpsc::channel(1);
2075
2076 let _registration = registry
2077 .accept_registration(
2078 &guard(),
2079 &caller("iso"),
2080 ®istration_full(&["iso"], "", "n1", &["charge"]),
2081 tx,
2082 )
2083 .await?;
2084
2085 assert_eq!(
2086 registry
2087 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", Some("n1"))?
2088 .len(),
2089 1,
2090 "an n1 worker must be admitted into the Pinned{{n1}} namespace's pool"
2091 );
2092 Ok(())
2093 }
2094
2095 #[tokio::test]
2099 async fn pinned_rejects_a_wrong_node_worker() -> Result<(), ServerError> {
2100 let store = namespace_store();
2101 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2102 let (tx, _rx) = mpsc::channel(1);
2103
2104 let denied = registry
2105 .accept_registration(
2106 &guard(),
2107 &caller("iso"),
2108 ®istration_full(&["iso"], "", "n2", &["charge"]),
2109 tx,
2110 )
2111 .await;
2112 assert!(
2113 matches!(denied, Err(ServerError::Namespace { .. })),
2114 "a wrong-node (n2) worker must be rejected from a Pinned{{n1}} namespace"
2115 );
2116 assert!(
2117 registry
2118 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
2119 .is_empty(),
2120 "a rejected registration must not insert a worker on any node"
2121 );
2122 Ok(())
2123 }
2124
2125 #[tokio::test]
2128 async fn pinned_rejects_a_node_less_worker() -> Result<(), ServerError> {
2129 let store = namespace_store();
2130 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2131 let (tx, _rx) = mpsc::channel(1);
2132
2133 let denied = registry
2134 .accept_registration(
2135 &guard(),
2136 &caller("iso"),
2137 ®istration_full(&["iso"], "", "", &["charge"]),
2138 tx,
2139 )
2140 .await;
2141 assert!(
2142 matches!(denied, Err(ServerError::Namespace { .. })),
2143 "a node-less worker must be rejected from a Pinned{{n1}} namespace"
2144 );
2145 assert!(
2146 registry
2147 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
2148 .is_empty(),
2149 "a rejected node-less registration must not insert a worker"
2150 );
2151 Ok(())
2152 }
2153
2154 #[tokio::test]
2158 async fn pinned_violation_rejects_the_whole_multi_namespace_registration()
2159 -> Result<(), ServerError> {
2160 let store = namespace_store();
2161 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2162 let (tx, _rx) = mpsc::channel(1);
2163
2164 let denied = registry
2165 .accept_registration(
2166 &guard(),
2167 &multi_caller(&["free", "iso"]),
2168 ®istration_full(&["free", "iso"], "", "n2", &["charge"]),
2169 tx,
2170 )
2171 .await;
2172 assert!(
2173 matches!(denied, Err(ServerError::Namespace { .. })),
2174 "a wrong-node worker serving a Pinned namespace fails the WHOLE registration"
2175 );
2176 assert!(
2177 registry
2178 .workers_for("free", DEFAULT_TASK_QUEUE, "charge", None)?
2179 .is_empty(),
2180 "the compliant namespace must NOT be partially admitted"
2181 );
2182 Ok(())
2183 }
2184
2185 #[tokio::test]
2188 async fn unplaced_and_prefer_admission_is_unaffected_by_the_pinned_gate()
2189 -> Result<(), ServerError> {
2190 let store = namespace_store();
2191 store
2194 .register_namespace("pref", NamespaceOrigin::Explicit)
2195 .await?;
2196 store
2197 .set_namespace_placement(
2198 "pref",
2199 NamespacePlacement::Prefer {
2200 nodes: ["n1".to_owned()].into_iter().collect(),
2201 },
2202 )
2203 .await?;
2204 let registry = minting_registry(&store, AutoCreate::Open);
2205
2206 let (tx_a, _rx_a) = mpsc::channel(1);
2207 let _reg_a = registry
2208 .accept_registration(
2209 &guard(),
2210 &caller("unpl"),
2211 ®istration_full(&["unpl"], "", "", &["charge"]),
2212 tx_a,
2213 )
2214 .await?;
2215 let (tx_b, _rx_b) = mpsc::channel(1);
2216 let _reg_b = registry
2217 .accept_registration(
2218 &guard(),
2219 &caller("pref"),
2220 ®istration_full(&["pref"], "", "", &["charge"]),
2221 tx_b,
2222 )
2223 .await?;
2224
2225 assert_eq!(
2226 registry
2227 .workers_for("unpl", DEFAULT_TASK_QUEUE, "charge", None)?
2228 .len(),
2229 1,
2230 "an Unplaced namespace admits a node-less worker unchanged"
2231 );
2232 assert_eq!(
2233 registry
2234 .workers_for("pref", DEFAULT_TASK_QUEUE, "charge", None)?
2235 .len(),
2236 1,
2237 "a Prefer namespace admits a node-less worker unchanged (only Pinned gates)"
2238 );
2239 Ok(())
2240 }
2241
2242 #[tokio::test]
2246 async fn no_minter_registry_skips_the_placement_gate() -> Result<(), ServerError> {
2247 let registry = ConnectedWorkerRegistry::default();
2248 let (tx, _rx) = mpsc::channel(1);
2249 let _registration = registry
2250 .accept_registration(
2251 &guard(),
2252 &caller("plain"),
2253 ®istration_full(&["plain"], "", "", &["charge"]),
2254 tx,
2255 )
2256 .await?;
2257 assert_eq!(
2258 registry
2259 .workers_for("plain", DEFAULT_TASK_QUEUE, "charge", None)?
2260 .len(),
2261 1,
2262 "with no minter the placement gate is a no-op — registration is unchanged"
2263 );
2264 Ok(())
2265 }
2266
2267 #[tokio::test]
2268 async fn concurrent_registrations_for_a_new_namespace_create_exactly_one_record()
2269 -> Result<(), ServerError> {
2270 let store = namespace_store();
2271 let registry = minting_registry(&store, AutoCreate::Open);
2272
2273 let mut handles = Vec::new();
2274 for _ in 0..8 {
2275 let registry = registry.clone();
2276 handles.push(tokio::spawn(async move {
2277 let (tx, rx) = mpsc::channel(1);
2278 let outcome = registry
2279 .accept_registration(
2280 &guard(),
2281 &caller("rush"),
2282 ®istration("rush", &["charge"]),
2283 tx,
2284 )
2285 .await;
2286 drop(rx);
2288 outcome.map(|registration| registration.worker_id())
2289 }));
2290 }
2291 for handle in handles {
2292 handle
2293 .await
2294 .map_err(|_| ServerError::lock_poisoned("registration task panicked"))??;
2295 }
2296
2297 let all = store.list_namespaces().await?;
2298 assert_eq!(
2299 all.iter().filter(|r| r.name == "rush").count(),
2300 1,
2301 "racing registrations must converge on exactly one durable record"
2302 );
2303 Ok(())
2304 }
2305
2306 #[tokio::test]
2307 async fn closed_rejects_unknown_namespace_and_does_not_create_it() -> Result<(), ServerError> {
2308 let store = namespace_store();
2309 let registry = minting_registry(&store, AutoCreate::Closed);
2310 let (tx, _rx) = mpsc::channel(1);
2311
2312 let denied = registry
2313 .accept_registration(
2314 &guard(),
2315 &caller("ghost"),
2316 ®istration("ghost", &["charge"]),
2317 tx,
2318 )
2319 .await;
2320 assert!(
2321 matches!(denied, Err(ServerError::Namespace { .. })),
2322 "closed policy must reject an unknown namespace"
2323 );
2324 assert!(
2325 store.get_namespace("ghost").await?.is_none(),
2326 "closed policy must NOT create the namespace it rejected"
2327 );
2328 let tq = DEFAULT_TASK_QUEUE;
2329 assert!(
2330 registry
2331 .workers_for("ghost", tq, "charge", None)?
2332 .is_empty(),
2333 "a rejected registration must not insert a worker"
2334 );
2335 Ok(())
2336 }
2337
2338 #[tokio::test]
2339 async fn closed_admits_a_known_namespace() -> Result<(), ServerError> {
2340 let store = namespace_store();
2341 store
2343 .register_namespace("known", NamespaceOrigin::Explicit)
2344 .await?;
2345 let registry = minting_registry(&store, AutoCreate::Closed);
2346 let (tx, _rx) = mpsc::channel(1);
2347
2348 let _registration = registry
2350 .accept_registration(
2351 &guard(),
2352 &caller("known"),
2353 ®istration("known", &["charge"]),
2354 tx,
2355 )
2356 .await?;
2357 let tq = DEFAULT_TASK_QUEUE;
2358 assert_eq!(
2359 registry.workers_for("known", tq, "charge", None)?.len(),
2360 1,
2361 "a known namespace must register under closed policy"
2362 );
2363 Ok(())
2364 }
2365
2366 #[tokio::test]
2367 async fn no_minter_leaves_registration_untouched() -> Result<(), ServerError> {
2368 let registry = ConnectedWorkerRegistry::default();
2371 let (tx, _rx) = mpsc::channel(1);
2372 let _registration = registry
2373 .accept_registration(
2374 &guard(),
2375 &caller("orders"),
2376 ®istration("orders", &["charge"]),
2377 tx,
2378 )
2379 .await?;
2380 let tq = DEFAULT_TASK_QUEUE;
2381 assert_eq!(registry.workers_for("orders", tq, "charge", None)?.len(), 1);
2382 Ok(())
2383 }
2384
2385 #[test]
2388 fn a_never_served_address_censuses_empty_with_no_poller_age() -> Result<(), ServerError> {
2389 let registry = ConnectedWorkerRegistry::default();
2390 let census = registry.pool_census("default", "general", "greet", None)?;
2391 assert_eq!(census.workers_in_pool, 0);
2392 assert_eq!(census.workers_serving_activity, 0);
2393 assert_eq!(census.compatible_workers, 0);
2394 assert_eq!(census.last_compatible_poller_age, None);
2395 assert!(!census.is_served());
2396 Ok(())
2397 }
2398
2399 #[test]
2403 fn the_census_separates_pool_activity_and_node_coverage() -> Result<(), ServerError> {
2404 let registry = ConnectedWorkerRegistry::default();
2405 let (tx, _rx) = mpsc::channel(1);
2406 let _worker = registry.register_namespaces(
2407 [String::from("default")],
2408 "general",
2409 Some(String::from("n1")),
2410 [String::from("greet")].iter(),
2411 tx,
2412 )?;
2413
2414 let unpinned = registry.pool_census("default", "general", "greet", None)?;
2415 assert_eq!(unpinned.workers_in_pool, 1);
2416 assert_eq!(unpinned.workers_serving_activity, 1);
2417 assert_eq!(unpinned.compatible_workers, 1);
2418 assert_eq!(unpinned.last_compatible_poller_age, Some(Duration::ZERO));
2419
2420 let other_activity = registry.pool_census("default", "general", "settle", None)?;
2423 assert_eq!(other_activity.workers_in_pool, 1);
2424 assert_eq!(other_activity.workers_serving_activity, 0);
2425 assert_eq!(other_activity.compatible_workers, 0);
2426
2427 let wrong_node = registry.pool_census("default", "general", "greet", Some("n2"))?;
2430 assert_eq!(wrong_node.workers_serving_activity, 1);
2431 assert_eq!(wrong_node.compatible_workers, 0);
2432 assert_eq!(wrong_node.last_compatible_poller_age, None);
2433 Ok(())
2434 }
2435
2436 #[test]
2457 fn a_pump_alive_worker_the_server_cannot_reach_is_not_selected() -> Result<(), ServerError> {
2458 const WINDOW: Duration = Duration::from_secs(30);
2459
2460 let registry = ConnectedWorkerRegistry::default();
2461 let tracker = HeartbeatTracker::new(WINDOW);
2462 let (tx, _rx) = mpsc::channel(1);
2463 let worker = registry.register_namespaces(
2464 [String::from("default")],
2465 "general",
2466 None,
2467 [String::from("greet")].iter(),
2468 tx,
2469 )?;
2470 let Some(worker_id) = worker.worker_id() else {
2471 return Err(test_failure("registration carries an id"));
2472 };
2473 let start = Instant::now();
2474 tracker.register_connection(worker_id, start)?;
2475
2476 let publish = |now: Instant| -> Result<(), ServerError> {
2481 registry.set_dispatch_ineligible(
2482 tracker
2483 .unreachable_workers(now)?
2484 .into_iter()
2485 .map(|excluded| excluded.worker_id)
2486 .collect(),
2487 )
2488 };
2489
2490 for _ in 0..DISPATCH_PROBATION_PINGS {
2495 assert!(
2496 tracker.record_dispatch_reachability(worker_id, start)?,
2497 "the worker is tracked while it serves its probation"
2498 );
2499 }
2500 publish(start)?;
2501 assert!(
2502 registry
2503 .select_worker("default", "general", "greet", None)?
2504 .is_some(),
2505 "a worker that has answered a full run of pings must be selectable: this is the \
2506 control, and without it a registry selecting NOBODY would satisfy every assertion \
2507 below"
2508 );
2509
2510 let much_later = start + WINDOW * 4;
2516 assert!(
2517 tracker.record_connection_activity(worker_id, much_later)?,
2518 "the worker is still tracked; its process is plainly alive"
2519 );
2520 assert!(
2521 tracker.record_dispatch_unreachable(worker_id)?,
2522 "the probe fired and went unanswered"
2523 );
2524 publish(much_later)?;
2525
2526 assert!(
2527 registry.is_dispatch_ineligible(worker_id)?,
2528 "a worker the server cannot reach must be marked ineligible however alive its \
2529 process looks"
2530 );
2531 assert!(
2532 registry
2533 .select_worker("default", "general", "greet", None)?
2534 .is_none(),
2535 "an ineligible worker must not be SELECTED: selecting one produces a dispatch that \
2536 can only fail, and on the liminal transport it fails by consuming connection \
2537 capacity — making the unreachability worse"
2538 );
2539
2540 assert!(
2545 tracker.record_dispatch_reachability(worker_id, much_later)?,
2546 "the worker is still tracked"
2547 );
2548 publish(much_later)?;
2549 assert!(
2550 registry.is_dispatch_ineligible(worker_id)?,
2551 "one answer part-way through a fresh probation must NOT restore eligibility: a link \
2552 answering one probe in three would otherwise flap in and out of selection"
2553 );
2554
2555 for _ in 1..DISPATCH_PROBATION_PINGS {
2556 assert!(
2557 tracker.record_dispatch_reachability(worker_id, much_later)?,
2558 "the worker is still tracked"
2559 );
2560 }
2561 publish(much_later)?;
2562 assert!(
2563 !registry.is_dispatch_ineligible(worker_id)?,
2564 "an answered ping must clear the exclusion"
2565 );
2566 assert!(
2567 registry
2568 .select_worker("default", "general", "greet", None)?
2569 .is_some(),
2570 "and the worker must be selectable again"
2571 );
2572 Ok(())
2573 }
2574
2575 #[test]
2579 fn a_departed_worker_leaves_a_last_compatible_poller_age() -> Result<(), ServerError> {
2580 let registry = ConnectedWorkerRegistry::default();
2581 let (tx, _rx) = mpsc::channel(1);
2582 let worker = registry.register_namespaces(
2583 [String::from("default")],
2584 "general",
2585 None,
2586 [String::from("greet")].iter(),
2587 tx,
2588 )?;
2589 worker.deregister()?;
2590
2591 let census = registry.pool_census("default", "general", "greet", None)?;
2592 assert_eq!(census.workers_in_pool, 0);
2593 assert_eq!(census.compatible_workers, 0);
2594 let age = census
2595 .last_compatible_poller_age
2596 .ok_or_else(|| test_failure("a departed worker must leave an age behind"))?;
2597 assert!(
2598 age < Duration::from_secs(60),
2599 "the recorded departure is implausibly old: {age:?}"
2600 );
2601
2602 let (tx, _rx) = mpsc::channel(1);
2605 let _back = registry.register_namespaces(
2606 [String::from("default")],
2607 "general",
2608 None,
2609 [String::from("greet")].iter(),
2610 tx,
2611 )?;
2612 let served = registry.pool_census("default", "general", "greet", None)?;
2613 assert_eq!(served.last_compatible_poller_age, Some(Duration::ZERO));
2614 assert!(served.is_served());
2615 Ok(())
2616 }
2617}