aion_server/worker/registry.rs
1//! Connected-worker registry keyed by worker-pool address and activity type.
2
3use 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, ProtoLivenessPing, 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
22/// The literal task queue an empty/absent selector normalizes to.
23///
24/// A worker-pool address has two disjoint dimensions; the second one
25/// (`task_queue`) is a liveness selector, not a correctness boundary. An empty
26/// `task_queue` is normalized to this one named default pool so a producer that
27/// names no queue and a worker that advertises none both land on the same pool.
28///
29/// Re-exported from [`aion_core::DEFAULT_TASK_QUEUE`] so the server cannot drift
30/// from the canonical domain default; the name is kept stable here for existing
31/// call sites.
32pub use aion_core::DEFAULT_TASK_QUEUE;
33
34/// Server-side handle used to push activity tasks to a connected worker stream.
35pub type WorkerTaskSender = mpsc::Sender<WorkerMessage>;
36
37/// Transport through which the server delivers a dispatch to a registered worker.
38///
39/// A worker is selected the SAME way regardless of transport (`select_worker`
40/// over the `(namespace, task_queue, node)` pool key); only the delivery leg
41/// differs. The default gRPC path pushes a [`WorkerMessage`] onto the worker's
42/// stream `mpsc` ([`WorkerDelivery::Grpc`]); a liminal-connected worker is
43/// delivered to by pushing the dispatch out on its existing liminal connection
44/// ([`WorkerDelivery::Liminal`], feature-gated). This enum is the minimal
45/// transport-agnostic seam: the registry holds it on each [`WorkerHandle`], and
46/// the dispatch path reads the variant it needs. The gRPC variant carries exactly
47/// the `mpsc::Sender` it always did, so the gRPC dispatch path is unchanged.
48#[derive(Clone, Debug)]
49pub enum WorkerDelivery {
50 /// gRPC stream delivery: the dispatch path pushes a [`WorkerMessage`] onto
51 /// this `mpsc` sender, exactly as before this enum existed.
52 Grpc(WorkerTaskSender),
53 /// Liminal server-push delivery: the dispatch path pushes the serialized
54 /// dispatch out on the worker's existing liminal connection and awaits the
55 /// correlated reply. Carries the connection identity needed to address that
56 /// push.
57 #[cfg(feature = "liminal-transport")]
58 Liminal(crate::worker::liminal_transport::LiminalWorkerDelivery),
59}
60
61/// Message queued from server-side dispatch/shutdown into a worker stream writer.
62#[derive(Clone, Debug, Eq, PartialEq)]
63pub enum WorkerMessage {
64 /// Activity invocation pushed to a worker.
65 ActivityTask(Box<ProtoActivityTask>),
66 /// Graceful-shutdown notification; no new work will be dispatched.
67 DrainRequest,
68 /// Transport liveness ping pushed by the liveness probe (#197).
69 ///
70 /// It travels this channel — the SAME one dispatches travel — deliberately.
71 /// The fact the probe needs is whether the server can reach this worker's
72 /// DISPATCH path, and a ping on a parallel channel could be answered by a
73 /// process whose dispatch path is dead. Answered by the worker SDK RUNTIME,
74 /// never by action code.
75 LivenessPing(ProtoLivenessPing),
76}
77
78/// Address of a worker pool: the two disjoint routing dimensions that select a
79/// pool, before an `activity_type` is matched within it.
80///
81/// `namespace` is the correctness/isolation boundary — a workflow's activities
82/// only ever reach workers in the workflow's namespace, so crossing it is a bug.
83/// `task_queue` is the pool/flavour selector within that namespace (norn /
84/// claude / cpu / gpu) — a miss is a liveness issue, never a correctness one.
85///
86/// This is a named type rather than a `(String, String)` tuple so a `node`
87/// dimension (Tier 3 affinity) can be added later without re-threading every
88/// call site.
89#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
90pub struct PoolAddress {
91 namespace: String,
92 task_queue: String,
93}
94
95impl PoolAddress {
96 /// Build a pool address, normalizing an empty `task_queue` to the named
97 /// [`DEFAULT_TASK_QUEUE`] pool. The `namespace` is the authorization
98 /// boundary and is never normalized.
99 #[must_use]
100 pub fn new(namespace: impl Into<String>, task_queue: impl Into<String>) -> Self {
101 let task_queue = task_queue.into();
102 let task_queue = if task_queue.is_empty() {
103 String::from(DEFAULT_TASK_QUEUE)
104 } else {
105 task_queue
106 };
107 Self {
108 namespace: namespace.into(),
109 task_queue,
110 }
111 }
112
113 /// The correctness/isolation boundary of this pool.
114 #[must_use]
115 pub fn namespace(&self) -> &str {
116 &self.namespace
117 }
118
119 /// The pool/flavour selector within the namespace.
120 #[must_use]
121 pub fn task_queue(&self) -> &str {
122 &self.task_queue
123 }
124}
125
126/// Registry match key: a worker-pool address plus the activity type matched
127/// within that pool. A named type (not an anonymous tuple) so the routing
128/// identity stays self-describing and extensible.
129#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
130struct ActivityKey {
131 pool: PoolAddress,
132 activity_type: String,
133}
134
135impl ActivityKey {
136 fn new(pool: PoolAddress, activity_type: impl Into<String>) -> Self {
137 Self {
138 pool,
139 activity_type: activity_type.into(),
140 }
141 }
142}
143
144type WorkerMap = HashMap<WorkerId, WorkerHandle>;
145type RegistryMap = HashMap<ActivityKey, WorkerMap>;
146
147/// Stable identifier assigned to a connected worker stream.
148#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
149pub struct WorkerId(u64);
150
151impl WorkerId {
152 /// Build a worker id from the numeric value exposed by administrative and
153 /// wire surfaces.
154 #[must_use]
155 pub const fn from_value(value: u64) -> Self {
156 Self(value)
157 }
158
159 /// Raw numeric value, as carried by the wire `RegisterAck.worker_id` so
160 /// workers can correlate their logs with the server's.
161 #[must_use]
162 pub const fn value(self) -> u64 {
163 self.0
164 }
165}
166
167/// Worker-supplied identity for one instance of a durable deployment.
168#[derive(Clone, Debug, PartialEq, Eq)]
169pub struct WorkerInstanceIdentity {
170 /// Durable deployment name supplied on the wire.
171 pub deployment: String,
172 /// Operator/launcher-assigned instance identifier.
173 pub instance_id: String,
174 /// Deployment-store lookup result captured when registration was accepted.
175 pub association: DeploymentAssociation,
176}
177
178struct RegistrationOptions {
179 instance: Option<WorkerInstanceIdentity>,
180 intervention_capabilities: InterventionCapabilities,
181}
182
183/// Cloneable handle for a registered worker stream.
184///
185/// A worker serves a SET of namespaces under a single `task_queue`, so it is
186/// indexed under one `(namespace, task_queue, activity_type)` key per namespace
187/// in its set. `node` is an OPTIONAL locality affinity (a locality, not a
188/// process — many handles may share a node id) used as a within-pool filter at
189/// selection time; `None` means the worker advertised no locality.
190#[derive(Clone, Debug)]
191pub struct WorkerHandle {
192 id: WorkerId,
193 namespaces: BTreeSet<String>,
194 task_queue: String,
195 node: Option<String>,
196 activity_types: BTreeSet<String>,
197 instance: Option<WorkerInstanceIdentity>,
198 delivery: WorkerDelivery,
199 /// The neutral mid-run intervention primitives this worker's harness advertises
200 /// support for (NOI-6). The server gates every intervention command on THIS set
201 /// and NEVER routes an unadvertised primitive. Empty = observability-only (the
202 /// default for every non-agent worker), so a normal activity worker advertises
203 /// no controls and the intervention router refuses every command for it.
204 intervention_capabilities: InterventionCapabilities,
205}
206
207impl WorkerHandle {
208 /// Worker identifier assigned by this server process.
209 #[must_use]
210 pub const fn id(&self) -> WorkerId {
211 self.id
212 }
213
214 /// Namespaces authorized for this worker stream. The worker is reachable for
215 /// a dispatch only when its set includes the workflow's namespace.
216 #[must_use]
217 pub const fn namespaces(&self) -> &BTreeSet<String> {
218 &self.namespaces
219 }
220
221 /// Task queue (pool/flavour) this worker serves within each namespace.
222 #[must_use]
223 pub fn task_queue(&self) -> &str {
224 &self.task_queue
225 }
226
227 /// Optional locality affinity this worker advertised. `None` means the
228 /// worker carries no node and is reachable only for unpinned dispatches.
229 #[must_use]
230 pub fn node(&self) -> Option<&str> {
231 self.node.as_deref()
232 }
233
234 /// Activity types advertised by this worker.
235 #[must_use]
236 pub fn activity_types(&self) -> &BTreeSet<String> {
237 &self.activity_types
238 }
239
240 /// Optional durable deployment/instance association supplied at registration.
241 #[must_use]
242 pub const fn instance(&self) -> Option<&WorkerInstanceIdentity> {
243 self.instance.as_ref()
244 }
245
246 /// The transport this worker is delivered to through.
247 #[must_use]
248 pub const fn delivery(&self) -> &WorkerDelivery {
249 &self.delivery
250 }
251
252 /// The neutral intervention primitives this worker's harness advertises (NOI-6).
253 ///
254 /// The intervention router gates on this set and never routes an unadvertised
255 /// primitive. Empty (the default for a plain activity worker) means the worker
256 /// is observability-only: the router refuses every intervention command for it.
257 #[must_use]
258 pub const fn intervention_capabilities(&self) -> &InterventionCapabilities {
259 &self.intervention_capabilities
260 }
261
262 /// gRPC stream sender used by the gRPC dispatch path to push work, or `None`
263 /// when this worker is delivered to over a non-gRPC transport (liminal).
264 ///
265 /// The gRPC dispatch path registers every worker with a [`WorkerDelivery::Grpc`]
266 /// delivery, so this is always `Some` for a gRPC-registered worker — the
267 /// behaviour the path relied on before delivery became transport-agnostic.
268 #[must_use]
269 pub fn sender(&self) -> Option<&WorkerTaskSender> {
270 match &self.delivery {
271 WorkerDelivery::Grpc(sender) => Some(sender),
272 #[cfg(feature = "liminal-transport")]
273 WorkerDelivery::Liminal(_) => None,
274 }
275 }
276}
277
278#[derive(Debug)]
279struct RegistryState {
280 next_worker_id: u64,
281 workers: BTreeMap<WorkerId, WorkerHandle>,
282 by_activity: RegistryMap,
283 /// Round-robin cursor per `(namespace, task_queue, activity_type)` triple, so
284 /// each pool rotates independently of every other pool.
285 rotation: HashMap<ActivityKey, usize>,
286 /// When a worker advertising a given `(pool, activity_type)` and node last
287 /// LEFT service, per advertised node. Read by [`ConnectedWorkerRegistry::pool_census`]
288 /// to answer "how old is the last compatible poller" for an address that has
289 /// none right now (R1). Node-keyed because a node-pinned dispatch's notion
290 /// of "compatible" is node-specific: a worker still serving the activity on
291 /// another node must never make a pinned address look freshly served.
292 last_departure: HashMap<ActivityKey, BTreeMap<Option<String>, Instant>>,
293 /// Workers the server currently cannot reach on the server-to-worker push
294 /// leg, and which are therefore excluded from selection while they remain
295 /// so.
296 ///
297 /// Held here rather than derived at selection time because selection must
298 /// not take the heartbeat tracker's lock: the liveness probe owns the
299 /// evidence and publishes the verdict, and selection only reads it.
300 ///
301 /// This is exclusion from DISPATCH, not deregistration. An unreachable
302 /// worker stays registered, keeps its in-flight work, and becomes eligible
303 /// again the moment a ping is answered — nothing about it is torn down on
304 /// the strength of a push failure.
305 dispatch_ineligible: BTreeSet<WorkerId>,
306}
307
308impl Default for RegistryState {
309 fn default() -> Self {
310 Self {
311 next_worker_id: 1,
312 workers: BTreeMap::new(),
313 by_activity: HashMap::new(),
314 rotation: HashMap::new(),
315 last_departure: HashMap::new(),
316 dispatch_ineligible: BTreeSet::new(),
317 }
318 }
319}
320
321/// Cloneable registry of currently connected worker streams.
322#[derive(Clone)]
323pub struct ConnectedWorkerRegistry {
324 inner: Arc<Mutex<RegistryState>>,
325 metrics: Option<Metrics>,
326 /// WS3 cluster-event publisher: emits `WorkerConnected`/`WorkerDisconnected`
327 /// topology deltas on register/deregister. `None` keeps existing
328 /// constructions (and every test) silent, exactly like `metrics`.
329 cluster_publisher: Option<ClusterEventPublisher>,
330 /// Minted-on-use hook (Control-Plane Phase 1). `None` disables minting, so
331 /// registration is byte-identical to before the registry existed; `Some`
332 /// durably records (open) or gates (closed) each authorized namespace.
333 minter: Option<NamespaceMinter>,
334 /// Durable deployment store used only to classify instance associations.
335 deployment_store: Option<Arc<dyn WorkerDeploymentStore>>,
336 worker_arrived: Arc<Notify>,
337 /// The refusal side of this registry's ledger, shared by every transport.
338 ///
339 /// The registry records who was ADMITTED; this records who was turned away,
340 /// so a repeated identical refusal can be met with silence. It lives here
341 /// because both callers of the admission gate already hold the registry, so
342 /// one home serves both — and one shared record is what stops the two
343 /// transports drifting apart the way #147 found them.
344 audit: Arc<AdmissionAudit>,
345}
346
347impl std::fmt::Debug for ConnectedWorkerRegistry {
348 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
349 formatter
350 .debug_struct("ConnectedWorkerRegistry")
351 .field(
352 "deployment_store_attached",
353 &self.deployment_store.is_some(),
354 )
355 .finish_non_exhaustive()
356 }
357}
358
359impl Default for ConnectedWorkerRegistry {
360 fn default() -> Self {
361 Self {
362 inner: Arc::new(Mutex::new(RegistryState::default())),
363 metrics: None,
364 cluster_publisher: None,
365 minter: None,
366 deployment_store: None,
367 worker_arrived: Arc::new(Notify::new()),
368 audit: Arc::new(AdmissionAudit::new()),
369 }
370 }
371}
372
373impl ConnectedWorkerRegistry {
374 /// Build a registry that records connected-worker gauge updates.
375 #[must_use]
376 pub fn with_metrics(metrics: Metrics) -> Self {
377 Self {
378 inner: Arc::new(Mutex::new(RegistryState::default())),
379 metrics: Some(metrics),
380 cluster_publisher: None,
381 minter: None,
382 deployment_store: None,
383 worker_arrived: Arc::new(Notify::new()),
384 audit: Arc::new(AdmissionAudit::new()),
385 }
386 }
387
388 /// The admission audit every registration transport names its refusals
389 /// through. Clones of a registry share one, so a worker refused over gRPC
390 /// and then over liminal is one site, not two.
391 #[must_use]
392 pub fn admission_audit(&self) -> &AdmissionAudit {
393 &self.audit
394 }
395
396 /// Attach the WS3 cluster-event publisher so worker topology changes are
397 /// pushed to the dashboard. Pure builder addition.
398 #[must_use]
399 pub fn with_cluster_publisher(mut self, publisher: ClusterEventPublisher) -> Self {
400 self.cluster_publisher = Some(publisher);
401 self
402 }
403
404 /// Attach the durable worker-deployment store used for association lookup.
405 #[must_use]
406 pub fn with_worker_deployment_store(mut self, store: Arc<dyn WorkerDeploymentStore>) -> Self {
407 self.deployment_store = Some(store);
408 self
409 }
410
411 /// Install the minted-on-use namespace hook (Control-Plane Phase 1).
412 ///
413 /// After a registration is authorized and its namespace set scoped, each
414 /// authorized namespace is durably recorded ([`AutoCreate::Open`]) or gated
415 /// ([`AutoCreate::Closed`]) through `store`. Without this builder the
416 /// registry never touches the namespace registry, so registration stays
417 /// byte-identical to before the registry existed. Pure builder addition,
418 /// mirroring [`Self::with_cluster_publisher`].
419 ///
420 /// When a cluster publisher has already been attached
421 /// ([`Self::with_cluster_publisher`], called first on the boot path), it is
422 /// threaded into the minter so a first worker-mint emits the live
423 /// `namespace created` delta to the ops console (S8). Order-independence is
424 /// not assumed: callers wire the publisher before minting on the boot path.
425 #[must_use]
426 pub fn with_namespace_minting(
427 mut self,
428 store: Arc<dyn NamespaceStore>,
429 policy: AutoCreate,
430 ) -> Self {
431 let minter = NamespaceMinter::new(store, policy);
432 let minter = match &self.cluster_publisher {
433 Some(publisher) => minter.with_cluster_publisher(publisher.clone()),
434 None => minter,
435 };
436 self.minter = Some(minter);
437 self
438 }
439
440 /// Thread the boot's namespace-mint routing context into the registry's
441 /// minter, so a worker registering for a namespace whose registry shard this
442 /// node does not own mints through the shard's owner instead of being
443 /// refused `NotOwner` forever.
444 ///
445 /// The SECOND of the two minter construction sites (the first is
446 /// [`ServerState::namespace_minter`](crate::ServerState::namespace_minter),
447 /// which serves the gRPC and HTTP start seams). Called after
448 /// [`Self::with_namespace_minting`] on the boot path; a no-op when no minter
449 /// is installed, and never called at all off-cluster, so default/test
450 /// registries stay byte-identical.
451 #[must_use]
452 pub fn with_namespace_routing(mut self, routing: crate::namespace::NamespaceRouting) -> Self {
453 self.minter = self.minter.map(|minter| minter.with_routing(routing));
454 self
455 }
456
457 /// Authorize a worker registration and insert it into the connected-worker registry.
458 ///
459 /// # Errors
460 ///
461 /// Returns [`ServerError`] if namespace authorization fails or the registry lock is poisoned.
462 pub async fn accept_registration(
463 &self,
464 guard: &NamespaceGuard,
465 caller: &CallerIdentity,
466 registration: &ProtoRegisterWorker,
467 sender: WorkerTaskSender,
468 ) -> Result<WorkerRegistration, ServerError> {
469 // Verify the operation against the guard's worker-registration policy,
470 // then authorize EACH namespace in the worker's set: a worker serves a
471 // SET of correctness boundaries, so the registration is denied unless
472 // the caller is granted every one. The wire's empty `node` carries no
473 // locality affinity; a non-empty value is the worker's advertised node.
474 guard
475 .scope(caller, &NamespaceOperation::register_worker(registration))
476 .await?;
477 let namespaces = guard.scope_worker_namespaces(caller, ®istration.namespaces)?;
478 // MINT HOOK (Control-Plane Phase 1). This runs strictly AFTER the
479 // per-namespace authorization above (`scope` + `scope_worker_namespaces`),
480 // so it can only ever mint a namespace the caller is already authorized
481 // for — the mint is auth-scoped by construction (CVE-2025-14986: open
482 // minting and namespace isolation only coexist when minting is
483 // auth-gated). It runs BEFORE the worker is inserted, so a `closed`
484 // rejection never leaves a half-registered worker behind.
485 self.mint_or_gate_namespaces(&namespaces).await?;
486 let node = optional_node(®istration.node);
487 // PLACEMENT-ADMISSION GATE (Control-Plane Phase 2, P2-I1). Runs strictly
488 // AFTER the mint hook (so every authorized namespace has a durable record to
489 // read a placement from) and with BOTH the worker's advertised `node` and
490 // the full authorized namespace set in scope. It rejects the WHOLE
491 // registration (Open Decision 6) when the worker's node violates any
492 // `Pinned{L}` namespace it would serve, so only L-node workers ever enter a
493 // hard-pinned namespace's pool. Auth-scoped by construction (it only ever
494 // gates a namespace already authorized above); a no-op with no minter
495 // installed, so default/test registries stay byte-identical.
496 self.enforce_pinned_placement(&namespaces, node.as_deref())
497 .await?;
498 let instance = self
499 .resolve_instance_identity(registration.instance.as_ref())
500 .await?;
501 self.register_delivery(
502 namespaces,
503 registration.task_queue.clone(),
504 node,
505 instance,
506 registration.activity_types.iter(),
507 WorkerDelivery::Grpc(sender),
508 )
509 }
510
511 /// Resolve a wire instance identity without refusing unknown deployments.
512 ///
513 /// Workers connect inward, forever: the server never holds an exhaustive
514 /// roster of launchers, and a worker may outlive the deployment record
515 /// that spawned it (record deleted, store swapped, or the worker predates
516 /// the record entirely). Refusing an unknown deployment name here would
517 /// orphan real, servable workers on exactly the recovery paths where they
518 /// matter most. The association is therefore observational — the record's
519 /// existence is captured as a Known/Absent/Unchecked state for supervision
520 /// to read (Unchecked means no store was attached) —
521 /// and never an admission gate. Admission is decided solely by the
522 /// contract checks that precede this call.
523 ///
524 /// # Errors
525 ///
526 /// Returns [`ServerError`] when durable deployment lookup fails.
527 pub async fn resolve_instance_identity(
528 &self,
529 instance: Option<&aion_proto::ProtoWorkerInstanceIdentity>,
530 ) -> Result<Option<WorkerInstanceIdentity>, ServerError> {
531 let Some(instance) = instance else {
532 return Ok(None);
533 };
534 let association = match &self.deployment_store {
535 Some(store) => {
536 if store
537 .get_worker_deployment(&instance.deployment)
538 .await
539 .map_err(ServerError::from)?
540 .is_some()
541 {
542 DeploymentAssociation::Known
543 } else {
544 DeploymentAssociation::Absent
545 }
546 }
547 None => DeploymentAssociation::Unchecked,
548 };
549 Ok(Some(WorkerInstanceIdentity {
550 deployment: instance.deployment.clone(),
551 instance_id: instance.instance_id.clone(),
552 association,
553 }))
554 }
555
556 /// Apply the minted-on-use policy to an already-authorized namespace set.
557 ///
558 /// A no-op when no minter is installed (every default/test registry), so
559 /// registration stays byte-identical. With a minter, the work is delegated
560 /// to the shared [`NamespaceMinter::mint_or_gate`] — the single
561 /// transport-agnostic implementation reused by the workflow-start safety net
562 /// — with [`NamespaceOrigin::WorkerMint`] so a first mint is attributed to
563 /// worker registration. See that method for the open/closed policy, the
564 /// idempotent "namespace created" event, and the retryable `NotOwner`
565 /// surface.
566 ///
567 /// # Errors
568 ///
569 /// Returns [`ServerError::StoreBackend`] if a durable upsert/lookup fails
570 /// (including a retryable `NotOwner` fence), or [`ServerError::Namespace`]
571 /// when `closed` rejects an unknown namespace.
572 async fn mint_or_gate_namespaces(&self, namespaces: &[String]) -> Result<(), ServerError> {
573 let Some(minter) = &self.minter else {
574 return Ok(());
575 };
576 minter
577 .mint_or_gate(namespaces, NamespaceOrigin::WorkerMint)
578 .await
579 }
580
581 /// Reject the whole registration when the worker's advertised `node` violates
582 /// any `Pinned{L}` namespace it would serve (Control-Plane Phase 2, P2-I1).
583 ///
584 /// For each authorized namespace whose placement is [`NamespacePlacement::Pinned`],
585 /// the worker's advertised `node` must be `Some(n)` with `n ∈ L`; a `None` node
586 /// or an `n ∉ L` is a loud, whole-registration rejection naming the namespace,
587 /// the node, and the required set. This guarantees only L-node workers ever
588 /// serve a hard-pinned namespace's pool, which is exactly what lets the
589 /// `Some(N ∉ L)` composition case (§2.2) resolve to the correct isolation stall
590 /// at dispatch rather than needing a start-time enumeration of future nodes.
591 ///
592 /// Non-`Pinned` placements ([`NamespacePlacement::Unplaced`]/[`NamespacePlacement::Prefer`])
593 /// are UNAFFECTED — byte-identical registration. A no-op when no minter is
594 /// installed (every default/test registry), so those stay behaviour-identical:
595 /// the gate reads placement from the SAME registry record the minter/placement
596 /// endpoint writes, never a second source of truth.
597 ///
598 /// # Errors
599 ///
600 /// Returns [`ServerError::Namespace`] (placement-admission denial) when the
601 /// worker's node violates a `Pinned` namespace, or [`ServerError::StoreBackend`]
602 /// if a placement read fails at the backend.
603 async fn enforce_pinned_placement(
604 &self,
605 namespaces: &[String],
606 node: Option<&str>,
607 ) -> Result<(), ServerError> {
608 let Some(minter) = &self.minter else {
609 return Ok(());
610 };
611 for namespace in namespaces {
612 let NamespacePlacement::Pinned { nodes } = minter.placement_of(namespace).await? else {
613 continue;
614 };
615 let admitted = node.is_some_and(|n| nodes.contains(n));
616 if !admitted {
617 return Err(ServerError::placement_admission_denied(
618 namespace, node, &nodes,
619 ));
620 }
621 }
622 Ok(())
623 }
624
625 /// Insert an already-authorized worker stream into the default task queue of
626 /// a single `namespace`, with no node affinity.
627 ///
628 /// Convenience over [`Self::register_namespaces`] for callers that serve one
629 /// namespace and do not select a task queue (notably tests of the default
630 /// pool).
631 ///
632 /// # Errors
633 ///
634 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
635 pub fn register<'a>(
636 &self,
637 namespace: impl Into<String>,
638 activity_types: impl IntoIterator<Item = &'a String>,
639 sender: WorkerTaskSender,
640 ) -> Result<WorkerRegistration, ServerError> {
641 self.register_namespaces(
642 [namespace.into()],
643 String::from(DEFAULT_TASK_QUEUE),
644 None,
645 activity_types,
646 sender,
647 )
648 }
649
650 /// Insert an already-authorized worker stream into one explicit worker pool
651 /// (single namespace + task queue), with no node affinity.
652 ///
653 /// # Errors
654 ///
655 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
656 pub fn register_pool<'a>(
657 &self,
658 pool: PoolAddress,
659 activity_types: impl IntoIterator<Item = &'a String>,
660 sender: WorkerTaskSender,
661 ) -> Result<WorkerRegistration, ServerError> {
662 let PoolAddress {
663 namespace,
664 task_queue,
665 } = pool;
666 self.register_namespaces([namespace], task_queue, None, activity_types, sender)
667 }
668
669 /// Insert an already-authorized worker stream serving a SET of namespaces
670 /// under one `task_queue`, with an optional `node` locality affinity.
671 ///
672 /// The worker is indexed under one `(namespace, task_queue, activity_type)`
673 /// key per namespace in its set, so a dispatch in any of those namespaces
674 /// can reach it. `node` is recorded on the handle and used only as a
675 /// within-pool filter at selection time — it is NOT part of [`PoolAddress`].
676 ///
677 /// # Errors
678 ///
679 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
680 pub fn register_namespaces<'a>(
681 &self,
682 namespaces: impl IntoIterator<Item = String>,
683 task_queue: impl Into<String>,
684 node: Option<String>,
685 activity_types: impl IntoIterator<Item = &'a String>,
686 sender: WorkerTaskSender,
687 ) -> Result<WorkerRegistration, ServerError> {
688 self.register_delivery(
689 namespaces,
690 task_queue,
691 node,
692 None,
693 activity_types,
694 WorkerDelivery::Grpc(sender),
695 )
696 }
697
698 /// Insert an already-authorized worker serving a SET of namespaces under one
699 /// `task_queue` and optional `node`, delivered to through an explicit
700 /// [`WorkerDelivery`] transport.
701 ///
702 /// This is the transport-agnostic registration core: [`Self::register_namespaces`]
703 /// is the gRPC façade over it (it wraps the stream sender in
704 /// [`WorkerDelivery::Grpc`]). Selection (`select_worker`/`workers_for`) is
705 /// identical across transports; only the held delivery differs.
706 ///
707 /// # Errors
708 ///
709 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
710 pub fn register_delivery<'a>(
711 &self,
712 namespaces: impl IntoIterator<Item = String>,
713 task_queue: impl Into<String>,
714 node: Option<String>,
715 instance: Option<WorkerInstanceIdentity>,
716 activity_types: impl IntoIterator<Item = &'a String>,
717 delivery: WorkerDelivery,
718 ) -> Result<WorkerRegistration, ServerError> {
719 self.register_delivery_inner(
720 namespaces,
721 task_queue,
722 node,
723 activity_types,
724 delivery,
725 RegistrationOptions {
726 instance,
727 intervention_capabilities: InterventionCapabilities::none(),
728 },
729 )
730 }
731
732 /// Insert an already-authorized worker exactly like [`Self::register_delivery`],
733 /// additionally recording the neutral [`InterventionCapabilities`] its harness
734 /// advertises (NOI-6).
735 ///
736 /// This is the capability-carrying registration core: [`Self::register_delivery`]
737 /// is the façade over it that advertises the empty set (observability-only), so
738 /// every existing caller stays byte-identical. Selection is unchanged — the
739 /// capability set is metadata the intervention router gates on, never a routing
740 /// dimension.
741 ///
742 /// # Errors
743 ///
744 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
745 pub fn register_delivery_with_capabilities<'a>(
746 &self,
747 namespaces: impl IntoIterator<Item = String>,
748 task_queue: impl Into<String>,
749 node: Option<String>,
750 activity_types: impl IntoIterator<Item = &'a String>,
751 delivery: WorkerDelivery,
752 intervention_capabilities: InterventionCapabilities,
753 ) -> Result<WorkerRegistration, ServerError> {
754 self.register_delivery_inner(
755 namespaces,
756 task_queue,
757 node,
758 activity_types,
759 delivery,
760 RegistrationOptions {
761 instance: None,
762 intervention_capabilities,
763 },
764 )
765 }
766
767 fn register_delivery_inner<'a>(
768 &self,
769 namespaces: impl IntoIterator<Item = String>,
770 task_queue: impl Into<String>,
771 node: Option<String>,
772 activity_types: impl IntoIterator<Item = &'a String>,
773 delivery: WorkerDelivery,
774 options: RegistrationOptions,
775 ) -> Result<WorkerRegistration, ServerError> {
776 let namespaces = namespaces.into_iter().collect::<BTreeSet<_>>();
777 let task_queue = task_queue.into();
778 let activity_types = activity_types.into_iter().cloned().collect::<BTreeSet<_>>();
779 let mut state = self.state()?;
780 let worker_id = WorkerId(state.next_worker_id);
781 state.next_worker_id = state.next_worker_id.saturating_add(1);
782
783 // Capture the node affinity for the WS3 WorkerConnected delta before the
784 // handle moves it.
785 let node_for_event = node.clone();
786 let instance_for_event = options.instance.clone();
787 let handle = WorkerHandle {
788 id: worker_id,
789 namespaces: namespaces.clone(),
790 task_queue: task_queue.clone(),
791 node,
792 activity_types: activity_types.clone(),
793 instance: options.instance,
794 delivery,
795 intervention_capabilities: options.intervention_capabilities,
796 };
797
798 for namespace in &namespaces {
799 let pool = PoolAddress::new(namespace.clone(), task_queue.clone());
800 for activity_type in &activity_types {
801 let key = ActivityKey::new(pool.clone(), activity_type.clone());
802 // This address is served again, so its departure record has no
803 // remaining meaning — drop it in lockstep with the insert. The
804 // departure map is therefore bounded by the addresses real
805 // workers have served and left, never by caller-supplied
806 // dispatch strings (which never create an entry at all).
807 if let Some(by_node) = state.last_departure.get_mut(&key) {
808 by_node.remove(&handle.node);
809 if by_node.is_empty() {
810 state.last_departure.remove(&key);
811 }
812 }
813 state
814 .by_activity
815 .entry(key)
816 .or_default()
817 .insert(worker_id, handle.clone());
818 }
819 }
820 let transport = transport_of(&handle.delivery);
821 state.workers.insert(worker_id, handle);
822 drop(state);
823
824 if let Some(metrics) = &self.metrics {
825 for namespace in &namespaces {
826 metrics.worker_connected(namespace);
827 }
828 }
829
830 // WS3: one WorkerConnected delta carrying the full namespace set (the
831 // event is namespace-list-valued; the deploy-scoped cluster channel sees
832 // it whole). Edge-triggered by the real insert, never a poll.
833 if let Some(publisher) = &self.cluster_publisher {
834 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
835 let task_queue_owned = task_queue.clone();
836 drop(publisher.emit(|meta| {
837 ClusterEvent::WorkerConnected {
838 meta,
839 worker_id: worker_id.value().to_string(),
840 namespaces: namespaces_vec,
841 task_queue: task_queue_owned,
842 transport,
843 node: node_for_event,
844 deployment: instance_for_event
845 .as_ref()
846 .map(|identity| identity.deployment.clone()),
847 deployment_association: instance_for_event.map(|identity| identity.association),
848 }
849 }));
850 }
851
852 self.worker_arrived.notify_waiters();
853
854 Ok(WorkerRegistration {
855 registry: self.clone(),
856 parts: Some(WorkerRegistrationParts {
857 worker_id,
858 namespaces,
859 task_queue,
860 activity_types,
861 }),
862 })
863 }
864
865 /// Wait until at least one new worker registers.
866 ///
867 /// Returns immediately if a registration occurred since the last call.
868 /// Callers should re-check the registry after waking — the newly arrived
869 /// worker may not serve the namespace or activity type the caller needs.
870 pub async fn wait_for_worker(&self) {
871 self.worker_arrived.notified().await;
872 }
873
874 /// Return a snapshot of workers registered for the
875 /// `(namespace, task_queue, activity_type)` pool, ordered by worker id and
876 /// then rotated so each call starts from the next worker in the pool. The
877 /// rotation cursor is per triple, so each pool round-robins independently.
878 ///
879 /// When `node` is `Some`, the result is filtered to workers whose advertised
880 /// node equals it — a dispatch pinned to a node reaches only workers on that
881 /// node (NODE affinity = require). When `node` is `None`, the behaviour is
882 /// exactly the unpinned pool: every worker in the `(namespace, task_queue)`
883 /// pool is a candidate regardless of locality. node is a within-pool filter,
884 /// NOT part of the pool key, so the per-triple rotation cursor is shared
885 /// across pinned and unpinned lookups of the same pool.
886 ///
887 /// The id sort matters: `by_activity` holds workers in a `HashMap`, whose
888 /// iteration order is unspecified. Sorting first makes the rotation below
889 /// the sole, deterministic source of ordering — true round-robin across
890 /// calls with the same membership, not a wobble layered on hash order.
891 ///
892 /// # Errors
893 ///
894 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
895 pub fn workers_for(
896 &self,
897 namespace: &str,
898 task_queue: &str,
899 activity_type: &str,
900 node: Option<&str>,
901 ) -> Result<Vec<WorkerHandle>, ServerError> {
902 let mut state = self.state()?;
903 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
904 let mut workers: Vec<WorkerHandle> = state
905 .by_activity
906 .get(&key)
907 .map(|workers| {
908 workers
909 .values()
910 .filter(|worker| worker_matches_node(worker, node))
911 .cloned()
912 .collect()
913 })
914 .unwrap_or_default();
915 if workers.is_empty() {
916 return Ok(workers);
917 }
918 workers.sort_by_key(WorkerHandle::id);
919 let idx = state.rotation.entry(key).or_insert(0);
920 let start = *idx % workers.len();
921 *idx = idx.wrapping_add(1);
922 let mut rotated = Vec::with_capacity(workers.len());
923 rotated.extend_from_slice(&workers[start..]);
924 rotated.extend_from_slice(&workers[..start]);
925 Ok(rotated)
926 }
927
928 /// Return a snapshot of every connected worker stream.
929 ///
930 /// # Errors
931 ///
932 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
933 pub fn all_workers(&self) -> Result<Vec<WorkerHandle>, ServerError> {
934 let state = self.state()?;
935 Ok(state.workers.values().cloned().collect())
936 }
937
938 /// Return the handle for a worker by id, or `None` when it is not registered.
939 ///
940 /// The intervention router resolves the owning worker of a target attempt by id
941 /// (NOI-6): the attempt-owner back-index stores a [`WorkerId`], and the router
942 /// reads back the live handle to gate on its advertised capabilities and select
943 /// its delivery. A `None` result means the owner disconnected — the router
944 /// treats that as the attempt-scoped no-op.
945 ///
946 /// # Errors
947 ///
948 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
949 pub fn worker_by_id(&self, worker_id: WorkerId) -> Result<Option<WorkerHandle>, ServerError> {
950 Ok(self.state()?.workers.get(&worker_id).cloned())
951 }
952
953 /// Replace the advertised intervention capabilities of a registered worker
954 /// (NOI-6). The liminal registration frame cannot carry capabilities, so a
955 /// liminal agent worker announces them on the reserved capabilities channel
956 /// right after registering, and this applies the announcement to the live
957 /// handle the intervention router gates on. Returns `false` when the worker
958 /// is no longer registered (a disconnect racing the announcement — benign).
959 ///
960 /// # Errors
961 ///
962 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
963 pub fn set_intervention_capabilities(
964 &self,
965 worker_id: WorkerId,
966 capabilities: &InterventionCapabilities,
967 ) -> Result<bool, ServerError> {
968 let mut state = self.state()?;
969 if !state.workers.contains_key(&worker_id) {
970 return Ok(false);
971 }
972 if let Some(handle) = state.workers.get_mut(&worker_id) {
973 handle.intervention_capabilities = capabilities.clone();
974 }
975 // The selection index holds handle clones; keep them capability-consistent
976 // even though capabilities are never a routing dimension.
977 for workers in state.by_activity.values_mut() {
978 if let Some(handle) = workers.get_mut(&worker_id) {
979 handle.intervention_capabilities = capabilities.clone();
980 }
981 }
982 Ok(true)
983 }
984
985 /// Broadcast a graceful drain request to every connected worker stream.
986 /// Workers are removed from routing before any transport signal is attempted.
987 /// A liminal worker has no drain control frame, so it is force-fenced and
988 /// deregistered with an error-level identity-bearing log instead of being
989 /// silently skipped.
990 ///
991 /// # Errors
992 ///
993 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
994 pub fn broadcast_drain(&self) -> Result<usize, ServerError> {
995 let workers = self.all_workers()?;
996 let mut delivered = 0usize;
997 for worker in workers {
998 if self.drain_worker(worker.id())? {
999 delivered = delivered.saturating_add(1);
1000 }
1001 }
1002 Ok(delivered)
1003 }
1004
1005 /// Stop assigning work to one worker and request a graceful transport drain.
1006 ///
1007 /// Returns `false` if the worker is not registered or if its transport has no
1008 /// drain channel. In the latter case the worker is deregistered immediately,
1009 /// after an error-level log names the worker and its transport limitation.
1010 ///
1011 /// # Errors
1012 ///
1013 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1014 pub fn drain_worker(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1015 let worker = {
1016 let mut state = self.state()?;
1017 let Some(worker) = state.workers.get(&worker_id).cloned() else {
1018 return Ok(false);
1019 };
1020 Self::remove_worker_from_service(&mut state, &worker);
1021 worker
1022 };
1023 match worker.delivery() {
1024 WorkerDelivery::Grpc(sender) => {
1025 if sender.try_send(WorkerMessage::DrainRequest).is_ok() {
1026 tracing::info!(worker_id = worker_id.value(), "worker drain requested");
1027 Ok(true)
1028 } else {
1029 tracing::error!(
1030 worker_id = worker_id.value(),
1031 "worker drain signal failed; force-deregistering closed transport"
1032 );
1033 self.deregister(worker_id)?;
1034 Ok(false)
1035 }
1036 }
1037 #[cfg(feature = "liminal-transport")]
1038 WorkerDelivery::Liminal(delivery) => {
1039 tracing::error!(
1040 worker_id = worker_id.value(),
1041 connection_pid = delivery.pid(),
1042 "liminal transport has no drain control channel; worker fenced and \
1043 deregistered at drain start"
1044 );
1045 self.deregister(worker_id)?;
1046 Ok(false)
1047 }
1048 }
1049 }
1050
1051 /// Select one worker for the `(namespace, task_queue, activity_type)` pool.
1052 ///
1053 /// When `node` is `Some`, only workers whose advertised node equals it are
1054 /// considered (NODE affinity = require); `None` considers every worker in
1055 /// the pool. node is a within-pool filter, NOT part of the pool key.
1056 ///
1057 /// # Errors
1058 ///
1059 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1060 pub fn select_worker(
1061 &self,
1062 namespace: &str,
1063 task_queue: &str,
1064 activity_type: &str,
1065 node: Option<&str>,
1066 ) -> Result<Option<WorkerHandle>, ServerError> {
1067 let state = self.state()?;
1068 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1069 Ok(state.by_activity.get(&key).and_then(|workers| {
1070 workers
1071 .values()
1072 .filter(|worker| worker_matches_node(worker, node))
1073 // Reachability is a dispatch PRECONDITION, so it is enforced
1074 // here rather than discovered at push time. Selecting a worker
1075 // the server cannot push to produces a dispatch that can only
1076 // fail, and on the liminal transport it fails by consuming
1077 // connection capacity — so an unreachable worker chosen anyway
1078 // makes its own unreachability worse.
1079 .filter(|worker| !state.dispatch_ineligible.contains(&worker.id))
1080 .min_by_key(|worker| worker.id)
1081 .cloned()
1082 }))
1083 }
1084
1085 /// Publish the liveness probe's reachability verdict: the set of workers the
1086 /// server cannot currently reach on the push leg, which
1087 /// [`Self::select_worker`] then skips.
1088 ///
1089 /// Replaces the whole set rather than toggling one worker, so the published
1090 /// verdict is always exactly one round's evidence and a worker can never be
1091 /// left excluded by a stale entry nobody cleared.
1092 ///
1093 /// # Errors
1094 ///
1095 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1096 pub fn set_dispatch_ineligible(
1097 &self,
1098 unreachable: BTreeSet<WorkerId>,
1099 ) -> Result<(), ServerError> {
1100 self.state()?.dispatch_ineligible = unreachable;
1101 Ok(())
1102 }
1103
1104 /// Whether a worker is currently excluded from dispatch selection for
1105 /// unreachability.
1106 ///
1107 /// # Errors
1108 ///
1109 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1110 pub fn is_dispatch_ineligible(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1111 Ok(self.state()?.dispatch_ineligible.contains(&worker_id))
1112 }
1113
1114 /// Every gRPC-delivered worker the liveness probe must ping this round,
1115 /// paired with the stream sender the ping rides (#197).
1116 ///
1117 /// The liminal half of the same census is
1118 /// [`LiminalConnectionNotifier::liveness_targets`](crate::worker::LiminalConnectionNotifier::liveness_targets),
1119 /// which enumerates CONNECTIONS. This one enumerates REGISTRATIONS, because
1120 /// a gRPC worker's only server-side identity is its registry handle: the
1121 /// stream is owned by a tonic task and reachable solely through the sender
1122 /// the registration carries.
1123 ///
1124 /// Deliberately not filtered by current eligibility. A worker excluded from
1125 /// dispatch is precisely the worker whose next answered ping restores it,
1126 /// so skipping the excluded set would make exclusion permanent — the exact
1127 /// defect this lane exists to remove.
1128 ///
1129 /// # Errors
1130 ///
1131 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1132 pub fn grpc_liveness_targets(
1133 &self,
1134 ) -> Result<Vec<super::grpc_liveness::GrpcLivenessTarget>, ServerError> {
1135 Ok(self
1136 .state()?
1137 .workers
1138 .values()
1139 .filter_map(|worker| match &worker.delivery {
1140 WorkerDelivery::Grpc(sender) => Some(super::grpc_liveness::GrpcLivenessTarget {
1141 worker_id: worker.id,
1142 sender: sender.clone(),
1143 }),
1144 #[cfg(feature = "liminal-transport")]
1145 WorkerDelivery::Liminal(_) => None,
1146 })
1147 .collect())
1148 }
1149
1150 /// How many workers that selection COULD have chosen are currently excluded
1151 /// from it by the liveness verdict (#197 R3).
1152 ///
1153 /// `tiers` is the ordered sequence of node filters the selection actually
1154 /// walked, and taking it as an argument — rather than re-deriving one here
1155 /// — is the whole point of the method. Selection over a `Pinned{L}`
1156 /// namespace walks the required labels and NEVER spills to a `None`
1157 /// any-node tier; a census taken over the row's own node instead would pass
1158 /// `None`, match every worker in the pool, and report an unlabelled worker
1159 /// that was never a candidate as the reason the row found nobody. The
1160 /// refusal built on that count then tells an operator not to start the
1161 /// labelled worker that is the only remedy.
1162 ///
1163 /// Counted as a UNION over the tiers and over DISTINCT workers: the pool is
1164 /// walked once and a worker admissible to more than one tier is counted
1165 /// once, so the number is a headcount rather than a sum of overlapping
1166 /// matches.
1167 ///
1168 /// Read under the SAME lock and with the SAME `worker_matches_node` filter
1169 /// [`Self::select_worker`] applies, so a refusal quoting this count
1170 /// describes the fleet selection actually saw.
1171 ///
1172 /// # Errors
1173 ///
1174 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1175 pub fn ineligible_workers_over_tiers(
1176 &self,
1177 namespace: &str,
1178 task_queue: &str,
1179 activity_type: &str,
1180 tiers: &[Option<String>],
1181 ) -> Result<usize, ServerError> {
1182 let state = self.state()?;
1183 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1184 Ok(state.by_activity.get(&key).map_or(0, |workers| {
1185 workers
1186 .values()
1187 .filter(|worker| state.dispatch_ineligible.contains(&worker.id))
1188 .filter(|worker| {
1189 tiers
1190 .iter()
1191 .any(|tier| worker_matches_node(worker, tier.as_deref()))
1192 })
1193 .count()
1194 }))
1195 }
1196
1197 /// The currently published exclusion set, read as a whole.
1198 ///
1199 /// The liveness probe reads this BEFORE publishing a round's verdict, so it
1200 /// can announce the workers that just LEFT the set. Without the whole
1201 /// previous set there is no way to name a recovery: per-worker queries can
1202 /// only be asked about workers the new verdict already mentions, and a
1203 /// recovered worker is precisely the one it does not.
1204 ///
1205 /// # Errors
1206 ///
1207 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1208 pub fn dispatch_ineligible(&self) -> Result<BTreeSet<WorkerId>, ServerError> {
1209 Ok(self.state()?.dispatch_ineligible.clone())
1210 }
1211
1212 /// Census the live fleet for one dispatch address (R1).
1213 ///
1214 /// Taken under the SAME lock discipline as [`Self::select_worker`] and read
1215 /// immediately after a selection miss, so the taxonomy verdict describes
1216 /// the fleet selection actually saw. Three nested counts — the pool, the
1217 /// activity coverage within it, the node coverage within that — are what
1218 /// separate `NO_LIVE_POLLERS` from `POLLERS_INCOMPATIBLE`.
1219 ///
1220 /// `last_compatible_poller_age` is zero while a compatible worker is
1221 /// connected, the elapsed time since the most recent compatible departure
1222 /// when one has left, and `None` when this server has never had one.
1223 ///
1224 /// # Errors
1225 ///
1226 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1227 pub fn pool_census(
1228 &self,
1229 namespace: &str,
1230 task_queue: &str,
1231 activity_type: &str,
1232 node: Option<&str>,
1233 ) -> Result<super::queue_service::PoolCensus, ServerError> {
1234 let state = self.state()?;
1235 let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
1236 let workers_in_pool = state
1237 .workers
1238 .values()
1239 .filter(|worker| {
1240 worker.task_queue == task_queue && worker.namespaces.contains(namespace)
1241 })
1242 .count();
1243 let serving = state.by_activity.get(&key);
1244 let workers_serving_activity = serving.map_or(0, HashMap::len);
1245 let compatible_workers = serving.map_or(0, |workers| {
1246 workers
1247 .values()
1248 .filter(|worker| worker_matches_node(worker, node))
1249 .count()
1250 });
1251 let last_compatible_poller_age = if compatible_workers > 0 {
1252 Some(Duration::ZERO)
1253 } else {
1254 state
1255 .last_departure
1256 .get(&key)
1257 .and_then(|by_node| {
1258 by_node
1259 .iter()
1260 .filter(|(departed_node, _)| match node {
1261 None => true,
1262 Some(node) => departed_node.as_deref() == Some(node),
1263 })
1264 .map(|(_, departed_at)| *departed_at)
1265 .max()
1266 })
1267 .map(|departed_at| departed_at.elapsed())
1268 };
1269 Ok(super::queue_service::PoolCensus {
1270 workers_in_pool,
1271 workers_serving_activity,
1272 compatible_workers,
1273 last_compatible_poller_age,
1274 })
1275 }
1276
1277 /// Return whether a worker stream is currently registered.
1278 ///
1279 /// The activity dispatch path uses this after queuing a task to detect a
1280 /// worker whose stream tore down concurrently: a sweep that ran before
1281 /// the dispatch tracked its task can never complete it, so the dispatch
1282 /// must fail the activity itself instead of waiting forever.
1283 ///
1284 /// # Errors
1285 ///
1286 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1287 pub fn is_registered(&self, worker_id: WorkerId) -> Result<bool, ServerError> {
1288 Ok(self.state()?.workers.contains_key(&worker_id))
1289 }
1290
1291 /// Remove a worker by id from every namespace/activity index it advertised.
1292 ///
1293 /// Emits a WS3 [`WorkerDeathReason::Disconnect`] delta — the truthful default
1294 /// for a removed worker whose stream/registration went away. Callers that can
1295 /// PROVE a finer reason (a liveness-timeout sweep) call
1296 /// [`Self::deregister_with_reason`] instead, so the dashboard never sees a
1297 /// fabricated distinction.
1298 ///
1299 /// # Errors
1300 ///
1301 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1302 pub fn deregister(&self, worker_id: WorkerId) -> Result<(), ServerError> {
1303 self.deregister_with_reason(worker_id, WorkerDeathReason::Disconnect)
1304 }
1305
1306 /// Remove a worker by id, attributing the departure to an explicit
1307 /// [`WorkerDeathReason`] the caller can prove at its call site (for example a
1308 /// heartbeat sweep passes [`WorkerDeathReason::Timeout`]).
1309 ///
1310 /// # Errors
1311 ///
1312 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1313 pub fn deregister_with_reason(
1314 &self,
1315 worker_id: WorkerId,
1316 reason: WorkerDeathReason,
1317 ) -> Result<(), ServerError> {
1318 let mut state = self.state()?;
1319 let removed_namespaces = Self::remove_worker(&mut state, worker_id);
1320 drop(state);
1321
1322 let Some(namespaces) = removed_namespaces else {
1323 // Already gone: no metrics double-count, no duplicate delta.
1324 return Ok(());
1325 };
1326
1327 if let Some(metrics) = &self.metrics {
1328 for namespace in &namespaces {
1329 metrics.worker_disconnected(namespace);
1330 }
1331 }
1332 self.emit_worker_disconnected(worker_id, &namespaces, reason);
1333
1334 Ok(())
1335 }
1336
1337 /// Emit a WS3 `WorkerDisconnected` delta if a publisher is attached.
1338 fn emit_worker_disconnected(
1339 &self,
1340 worker_id: WorkerId,
1341 namespaces: &BTreeSet<String>,
1342 reason: WorkerDeathReason,
1343 ) {
1344 if let Some(publisher) = &self.cluster_publisher {
1345 let namespaces_vec: Vec<String> = namespaces.iter().cloned().collect();
1346 drop(publisher.emit(|meta| ClusterEvent::WorkerDisconnected {
1347 meta,
1348 worker_id: worker_id.value().to_string(),
1349 namespaces: namespaces_vec,
1350 reason,
1351 }));
1352 }
1353 }
1354
1355 /// Remove a worker from every `(namespace, task_queue, activity_type)` index
1356 /// it advertised. Returns the namespace set it served (for metrics), or
1357 /// `None` if the worker was already gone.
1358 fn remove_worker(state: &mut RegistryState, worker_id: WorkerId) -> Option<BTreeSet<String>> {
1359 let handle = state.workers.remove(&worker_id)?;
1360 Self::remove_worker_from_service(state, &handle);
1361 Some(handle.namespaces)
1362 }
1363
1364 /// Fence a worker from every routing index while retaining its connection
1365 /// handle so an in-progress graceful drain can finish and tear down normally.
1366 fn remove_worker_from_service(state: &mut RegistryState, handle: &WorkerHandle) {
1367 let departed_at = Instant::now();
1368 for namespace in &handle.namespaces {
1369 let pool = PoolAddress::new(namespace.clone(), handle.task_queue.clone());
1370 for activity_type in &handle.activity_types {
1371 let key = ActivityKey::new(pool.clone(), activity_type.clone());
1372 // Departure time is recorded for EVERY address this worker
1373 // served, so a later census can age the last compatible poller
1374 // instead of reporting a bare "nobody" (R1).
1375 state
1376 .last_departure
1377 .entry(key.clone())
1378 .or_default()
1379 .insert(handle.node.clone(), departed_at);
1380 if let Some(workers) = state.by_activity.get_mut(&key) {
1381 workers.remove(&handle.id);
1382 if workers.is_empty() {
1383 state.by_activity.remove(&key);
1384 // Prune the round-robin cursor in lockstep: the cursor
1385 // map is keyed on arbitrary caller-supplied strings and
1386 // is lazily created by `workers_for`, so leaving stale
1387 // entries behind leaks memory unboundedly on a
1388 // never-dying server. When the last worker for a triple
1389 // leaves, its cursor has no remaining meaning.
1390 state.rotation.remove(&key);
1391 }
1392 }
1393 }
1394 }
1395 }
1396
1397 fn state(&self) -> Result<MutexGuard<'_, RegistryState>, ServerError> {
1398 self.inner
1399 .lock()
1400 .map_err(|_| ServerError::lock_poisoned("connected worker registry"))
1401 }
1402}
1403
1404/// Map a held [`WorkerDelivery`] to the wire [`WorkerTransport`] discriminant for
1405/// the WS3 `WorkerConnected` delta.
1406const fn transport_of(delivery: &WorkerDelivery) -> WorkerTransport {
1407 match delivery {
1408 WorkerDelivery::Grpc(_) => WorkerTransport::Grpc,
1409 #[cfg(feature = "liminal-transport")]
1410 WorkerDelivery::Liminal(_) => WorkerTransport::Liminal,
1411 }
1412}
1413
1414/// Normalize a wire `node` string into an optional locality affinity: an empty
1415/// value (the proto3 default) carries no node, anything else is the worker's
1416/// advertised node id.
1417///
1418/// Shared with contract admission ([`super::contracts::validate_worker_contracts`])
1419/// rather than restated there: admission decides which of a package's actions a
1420/// connection owes from the same locality this registry then routes by, and two
1421/// independent normalizations of the wire default would be a place for that
1422/// agreement to drift silently.
1423pub(crate) fn optional_node(node: &str) -> Option<String> {
1424 if node.is_empty() {
1425 None
1426 } else {
1427 Some(node.to_owned())
1428 }
1429}
1430
1431/// Whether a worker satisfies an optional node filter. `None` (unpinned) matches
1432/// every worker; `Some(node)` matches only a worker advertising that exact node
1433/// (NODE affinity = require). A worker with no advertised node never matches a
1434/// pinned dispatch.
1435fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
1436 match node {
1437 None => true,
1438 Some(node) => worker.node() == Some(node),
1439 }
1440}
1441
1442#[derive(Clone, Debug)]
1443struct WorkerRegistrationParts {
1444 worker_id: WorkerId,
1445 namespaces: BTreeSet<String>,
1446 task_queue: String,
1447 activity_types: BTreeSet<String>,
1448}
1449
1450/// Registration token owned by the worker stream task.
1451///
1452/// Dropping the token performs best-effort cleanup for disconnect paths. Call
1453/// [`WorkerRegistration::deregister`] when the caller needs a typed poison error.
1454#[derive(Debug)]
1455pub struct WorkerRegistration {
1456 registry: ConnectedWorkerRegistry,
1457 parts: Option<WorkerRegistrationParts>,
1458}
1459
1460impl WorkerRegistration {
1461 /// Worker id assigned to this registration.
1462 #[must_use]
1463 pub fn worker_id(&self) -> Option<WorkerId> {
1464 self.parts.as_ref().map(|parts| parts.worker_id)
1465 }
1466
1467 /// Authorized namespace set for this registration.
1468 #[must_use]
1469 pub fn namespaces(&self) -> Option<&BTreeSet<String>> {
1470 self.parts.as_ref().map(|parts| &parts.namespaces)
1471 }
1472
1473 /// Task queue (pool/flavour) this registration serves within each namespace.
1474 #[must_use]
1475 pub fn task_queue(&self) -> Option<&str> {
1476 self.parts.as_ref().map(|parts| parts.task_queue.as_str())
1477 }
1478
1479 /// Activity types advertised by this registration.
1480 #[must_use]
1481 pub fn activity_types(&self) -> Option<&BTreeSet<String>> {
1482 self.parts.as_ref().map(|parts| &parts.activity_types)
1483 }
1484
1485 /// Explicitly remove this worker from the registry.
1486 ///
1487 /// # Errors
1488 ///
1489 /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
1490 pub fn deregister(mut self) -> Result<(), ServerError> {
1491 let Some(parts) = self.parts.take() else {
1492 return Ok(());
1493 };
1494 self.registry.deregister(parts.worker_id)
1495 }
1496}
1497
1498impl Drop for WorkerRegistration {
1499 fn drop(&mut self) {
1500 let Some(parts) = self.parts.take() else {
1501 return;
1502 };
1503 let removed_namespaces = self.registry.inner.lock().ok().and_then(|mut state| {
1504 ConnectedWorkerRegistry::remove_worker(&mut state, parts.worker_id)
1505 });
1506 if let Some(namespaces) = removed_namespaces {
1507 if let Some(metrics) = &self.registry.metrics {
1508 for namespace in &namespaces {
1509 metrics.worker_disconnected(namespace);
1510 }
1511 }
1512 // A dropped registration token means the worker's stream/connection
1513 // went away — the truthful reason is Disconnect, not a fabricated
1514 // timeout/deregister distinction this path cannot prove.
1515 self.registry.emit_worker_disconnected(
1516 parts.worker_id,
1517 &namespaces,
1518 WorkerDeathReason::Disconnect,
1519 );
1520 }
1521 }
1522}
1523
1524#[cfg(test)]
1525mod tests {
1526 use crate::config::NamespaceMode;
1527 use crate::namespace::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
1528 use crate::worker::heartbeat::{DISPATCH_PROBATION_PINGS, HeartbeatTracker};
1529
1530 use super::*;
1531
1532 fn guard() -> NamespaceGuard {
1533 NamespaceGuard::new(NamespaceResolver::authorization_only(
1534 NamespaceMode::SharedEngine,
1535 StaticWorkflowNamespaces::default(),
1536 StaticScheduleNamespaces::default(),
1537 ))
1538 }
1539
1540 fn caller(namespace: &str) -> CallerIdentity {
1541 CallerIdentity::new("worker", [namespace.to_owned()])
1542 }
1543
1544 /// A test-expectation failure as a `ServerError`, so `Result`-returning
1545 /// tests can fail on an unexpected `None` without `panic!`/`expect`.
1546 fn test_failure(message: &str) -> ServerError {
1547 ServerError::worker_dispatch("default".to_owned(), "test".to_owned(), message.to_owned())
1548 }
1549
1550 fn registration(namespace: &str, activity_types: &[&str]) -> ProtoRegisterWorker {
1551 registration_with_queue(namespace, "", activity_types)
1552 }
1553
1554 fn registration_with_queue(
1555 namespace: &str,
1556 task_queue: &str,
1557 activity_types: &[&str],
1558 ) -> ProtoRegisterWorker {
1559 registration_full(&[namespace], task_queue, "", activity_types)
1560 }
1561
1562 fn registration_full(
1563 namespaces: &[&str],
1564 task_queue: &str,
1565 node: &str,
1566 activity_types: &[&str],
1567 ) -> ProtoRegisterWorker {
1568 ProtoRegisterWorker {
1569 namespaces: namespaces.iter().map(|value| (*value).to_owned()).collect(),
1570 activity_types: activity_types
1571 .iter()
1572 .map(|value| (*value).to_owned())
1573 .collect(),
1574 task_queue: task_queue.to_owned(),
1575 node: node.to_owned(),
1576 activities: Vec::new(),
1577 identity: String::new(),
1578 instance: None,
1579 }
1580 }
1581
1582 fn multi_caller(namespaces: &[&str]) -> CallerIdentity {
1583 CallerIdentity::new("worker", namespaces.iter().map(|value| (*value).to_owned()))
1584 }
1585
1586 /// `set_intervention_capabilities` replaces a live worker's advertised set —
1587 /// the announcement path a liminal agent worker takes after its in-band
1588 /// registration (the registration frame cannot carry capabilities) — and
1589 /// reports an unknown worker as `false` (an announcement racing a
1590 /// disconnect is benign, never an error).
1591 #[tokio::test]
1592 async fn set_intervention_capabilities_updates_live_worker() -> Result<(), ServerError> {
1593 let registry = ConnectedWorkerRegistry::default();
1594 let (sender, _receiver) = mpsc::channel(1);
1595 let types = ["scout".to_owned()];
1596 let guard = registry.register_delivery_with_capabilities(
1597 ["default".to_owned()],
1598 "default",
1599 None,
1600 types.iter(),
1601 WorkerDelivery::Grpc(sender),
1602 InterventionCapabilities::none(),
1603 )?;
1604 let Some(worker_id) = guard.worker_id() else {
1605 return Err(test_failure("registration carries an id"));
1606 };
1607
1608 let announced = InterventionCapabilities {
1609 supported: vec![aion_core::InterventionPrimitive::InjectMessage],
1610 };
1611 assert!(
1612 registry.set_intervention_capabilities(worker_id, &announced)?,
1613 "a live worker's capabilities must be updatable"
1614 );
1615 let Some(handle) = registry.worker_by_id(worker_id)? else {
1616 return Err(test_failure("worker stays registered"));
1617 };
1618 assert_eq!(handle.intervention_capabilities(), &announced);
1619
1620 assert!(
1621 !registry.set_intervention_capabilities(WorkerId(u64::MAX), &announced)?,
1622 "an unknown worker reports false, never an error"
1623 );
1624 Ok(())
1625 }
1626
1627 #[tokio::test]
1628 async fn register_and_deregister_are_namespace_isolated() -> Result<(), ServerError> {
1629 let registry = ConnectedWorkerRegistry::default();
1630 let (tenant_a_tx, _tenant_a_rx) = mpsc::channel(1);
1631 let (tenant_b_tx, _tenant_b_rx) = mpsc::channel(1);
1632
1633 let tenant_a = registry
1634 .accept_registration(
1635 &guard(),
1636 &caller("tenant-a"),
1637 ®istration("tenant-a", &["charge", "charge"]),
1638 tenant_a_tx,
1639 )
1640 .await?;
1641 let tenant_b = registry
1642 .accept_registration(
1643 &guard(),
1644 &caller("tenant-b"),
1645 ®istration("tenant-b", &["charge"]),
1646 tenant_b_tx,
1647 )
1648 .await?;
1649
1650 let tq = DEFAULT_TASK_QUEUE;
1651 assert_eq!(
1652 registry.workers_for("tenant-a", tq, "charge", None)?.len(),
1653 1
1654 );
1655 assert_eq!(
1656 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1657 1
1658 );
1659 assert!(
1660 registry
1661 .workers_for("tenant-a", tq, "missing", None)?
1662 .is_empty()
1663 );
1664
1665 let tenant_a_id = tenant_a.worker_id();
1666 tenant_a.deregister()?;
1667
1668 assert!(
1669 registry
1670 .workers_for("tenant-a", tq, "charge", None)?
1671 .is_empty()
1672 );
1673 assert_eq!(
1674 registry.workers_for("tenant-b", tq, "charge", None)?.len(),
1675 1
1676 );
1677 assert_ne!(tenant_a_id, tenant_b.worker_id());
1678
1679 tenant_b.deregister()?;
1680 assert!(
1681 registry
1682 .workers_for("tenant-b", tq, "charge", None)?
1683 .is_empty()
1684 );
1685 Ok(())
1686 }
1687
1688 #[tokio::test]
1689 async fn denied_namespace_is_not_registered() -> Result<(), ServerError> {
1690 let registry = ConnectedWorkerRegistry::default();
1691 let (tx, _rx) = mpsc::channel(1);
1692 let denied = registry
1693 .accept_registration(
1694 &guard(),
1695 &caller("tenant-a"),
1696 ®istration("tenant-b", &["charge"]),
1697 tx,
1698 )
1699 .await;
1700
1701 assert!(denied.is_err());
1702 assert!(
1703 registry
1704 .workers_for("tenant-b", DEFAULT_TASK_QUEUE, "charge", None)?
1705 .is_empty()
1706 );
1707 Ok(())
1708 }
1709
1710 #[tokio::test]
1711 async fn task_queues_partition_disjoint_pools_within_one_namespace() -> Result<(), ServerError>
1712 {
1713 // Same namespace + same activity_type, two DIFFERENT task queues: the
1714 // pools are disjoint, a lookup for one queue never returns the other's
1715 // worker, and round-robin holds independently per (ns, tq, type) triple.
1716 let registry = ConnectedWorkerRegistry::default();
1717 let (norn_tx, _norn_rx) = mpsc::channel(1);
1718 let (claude_a_tx, _claude_a_rx) = mpsc::channel(1);
1719 let (claude_b_tx, _claude_b_rx) = mpsc::channel(1);
1720
1721 let norn = registry
1722 .accept_registration(
1723 &guard(),
1724 &caller("local"),
1725 ®istration_with_queue("local", "norn", &["dev"]),
1726 norn_tx,
1727 )
1728 .await?;
1729 // Two workers on the SAME (local, claude) pool to exercise round-robin.
1730 let claude_a = registry
1731 .accept_registration(
1732 &guard(),
1733 &caller("local"),
1734 ®istration_with_queue("local", "claude", &["dev"]),
1735 claude_a_tx,
1736 )
1737 .await?;
1738 let claude_b = registry
1739 .accept_registration(
1740 &guard(),
1741 &caller("local"),
1742 ®istration_with_queue("local", "claude", &["dev"]),
1743 claude_b_tx,
1744 )
1745 .await?;
1746
1747 let norn_pool = registry.workers_for("local", "norn", "dev", None)?;
1748 assert_eq!(norn_pool.len(), 1, "norn pool has exactly its one worker");
1749 let norn_id = norn.worker_id().ok_or_else(missing_id)?;
1750 assert_eq!(norn_pool[0].id(), norn_id);
1751
1752 let claude_pool = registry.workers_for("local", "claude", "dev", None)?;
1753 assert_eq!(
1754 claude_pool.len(),
1755 2,
1756 "claude pool sees only its two workers"
1757 );
1758 let claude_ids: BTreeSet<WorkerId> = claude_pool.iter().map(WorkerHandle::id).collect();
1759 assert!(
1760 !claude_ids.contains(&norn_id),
1761 "the norn worker must never appear in the claude pool"
1762 );
1763
1764 // A dispatch targeting `norn` never reaches a `claude` worker, and vice
1765 // versa: the disjoint key is the boundary.
1766 assert!(
1767 !registry
1768 .workers_for("local", "norn", "dev", None)?
1769 .iter()
1770 .any(|worker| claude_ids.contains(&worker.id()))
1771 );
1772
1773 // Round-robin per triple: the (local, claude, dev) cursor advances
1774 // independently and cycles through both claude workers, while the
1775 // (local, norn, dev) cursor keeps returning its single worker.
1776 let first = registry.workers_for("local", "claude", "dev", None)?[0].id();
1777 let second = registry.workers_for("local", "claude", "dev", None)?[0].id();
1778 assert_ne!(
1779 first, second,
1780 "claude pool round-robins across both workers"
1781 );
1782 assert_eq!(
1783 registry.workers_for("local", "norn", "dev", None)?[0].id(),
1784 norn_id,
1785 "the norn pool rotation is unaffected by claude traffic"
1786 );
1787
1788 norn.deregister()?;
1789 claude_a.deregister()?;
1790 claude_b.deregister()?;
1791 Ok(())
1792 }
1793
1794 #[tokio::test]
1795 async fn same_task_queue_in_different_namespaces_is_isolated() -> Result<(), ServerError> {
1796 // Same task_queue string, two DIFFERENT namespaces: namespace is the
1797 // correctness boundary, so the pools are isolated.
1798 let registry = ConnectedWorkerRegistry::default();
1799 let (local_tx, _local_rx) = mpsc::channel(1);
1800 let (remote_tx, _remote_rx) = mpsc::channel(1);
1801
1802 let local = registry
1803 .accept_registration(
1804 &guard(),
1805 &caller("local"),
1806 ®istration_with_queue("local", "gpu", &["render"]),
1807 local_tx,
1808 )
1809 .await?;
1810 let remote = registry
1811 .accept_registration(
1812 &guard(),
1813 &caller("remote"),
1814 ®istration_with_queue("remote", "gpu", &["render"]),
1815 remote_tx,
1816 )
1817 .await?;
1818
1819 let local_pool = registry.workers_for("local", "gpu", "render", None)?;
1820 let remote_pool = registry.workers_for("remote", "gpu", "render", None)?;
1821 assert_eq!(local_pool.len(), 1);
1822 assert_eq!(remote_pool.len(), 1);
1823 assert_ne!(
1824 local_pool[0].id(),
1825 remote_pool[0].id(),
1826 "a shared task_queue string does not merge two namespaces"
1827 );
1828
1829 local.deregister()?;
1830 assert!(
1831 registry
1832 .workers_for("local", "gpu", "render", None)?
1833 .is_empty(),
1834 "deregistering the local worker leaves the remote namespace untouched"
1835 );
1836 assert_eq!(
1837 registry.workers_for("remote", "gpu", "render", None)?.len(),
1838 1
1839 );
1840
1841 remote.deregister()?;
1842 Ok(())
1843 }
1844
1845 #[tokio::test]
1846 async fn worker_serving_a_namespace_set_is_reachable_in_each() -> Result<(), ServerError> {
1847 // A worker advertising {a, b} is reachable for dispatch in BOTH a and b;
1848 // a worker in {a} is NOT reachable in b.
1849 let registry = ConnectedWorkerRegistry::default();
1850 let (ab_tx, _ab_rx) = mpsc::channel(1);
1851 let (a_tx, _a_rx) = mpsc::channel(1);
1852
1853 let worker_ab = registry
1854 .accept_registration(
1855 &guard(),
1856 &multi_caller(&["a", "b"]),
1857 ®istration_full(&["a", "b"], "default", "", &["dev"]),
1858 ab_tx,
1859 )
1860 .await?;
1861 let worker_a = registry
1862 .accept_registration(
1863 &guard(),
1864 &caller("a"),
1865 ®istration_full(&["a"], "default", "", &["dev"]),
1866 a_tx,
1867 )
1868 .await?;
1869
1870 let in_a = registry.workers_for("a", "default", "dev", None)?;
1871 let in_b = registry.workers_for("b", "default", "dev", None)?;
1872 let both_id = worker_ab.worker_id().ok_or_else(missing_id)?;
1873 let only_a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1874
1875 // Namespace a sees BOTH workers; namespace b sees ONLY the {a, b} worker.
1876 let a_ids: BTreeSet<WorkerId> = in_a.iter().map(WorkerHandle::id).collect();
1877 assert_eq!(a_ids, BTreeSet::from([both_id, only_a_id]));
1878 assert_eq!(in_b.len(), 1, "only the {{a, b}} worker is reachable in b");
1879 assert_eq!(in_b[0].id(), both_id);
1880 assert!(
1881 !in_b.iter().any(|worker| worker.id() == only_a_id),
1882 "the {{a}}-only worker must not be reachable in b"
1883 );
1884
1885 // Deregistering the {a, b} worker removes it from BOTH buckets.
1886 worker_ab.deregister()?;
1887 assert!(
1888 registry
1889 .workers_for("b", "default", "dev", None)?
1890 .is_empty()
1891 );
1892 assert_eq!(registry.workers_for("a", "default", "dev", None)?.len(), 1);
1893
1894 worker_a.deregister()?;
1895 Ok(())
1896 }
1897
1898 #[tokio::test]
1899 async fn node_pin_filters_within_pool() -> Result<(), ServerError> {
1900 // Two workers in the same (namespace, task_queue) pool on different
1901 // nodes: unpinned round-robins across both; pinned to node N reaches
1902 // ONLY the worker(s) on N; pinned to a node with no worker finds none.
1903 let registry = ConnectedWorkerRegistry::default();
1904 let (n1_tx, _n1_rx) = mpsc::channel(1);
1905 let (n2_tx, _n2_rx) = mpsc::channel(1);
1906
1907 let on_n1 = registry
1908 .accept_registration(
1909 &guard(),
1910 &caller("ns"),
1911 ®istration_full(&["ns"], "tq", "n1", &["dev"]),
1912 n1_tx,
1913 )
1914 .await?;
1915 let on_n2 = registry
1916 .accept_registration(
1917 &guard(),
1918 &caller("ns"),
1919 ®istration_full(&["ns"], "tq", "n2", &["dev"]),
1920 n2_tx,
1921 )
1922 .await?;
1923 let n1_id = on_n1.worker_id().ok_or_else(missing_id)?;
1924 let n2_id = on_n2.worker_id().ok_or_else(missing_id)?;
1925
1926 // Unpinned: both workers are candidates and round-robin advances.
1927 let unpinned = registry.workers_for("ns", "tq", "dev", None)?;
1928 assert_eq!(unpinned.len(), 2, "unpinned reaches the whole pool");
1929 let first = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1930 let second = registry.workers_for("ns", "tq", "dev", None)?[0].id();
1931 assert_ne!(first, second, "unpinned round-robins across both nodes");
1932
1933 // Pinned to n1: only the n1 worker; pinned to n2: only the n2 worker.
1934 let pinned_n1 = registry.workers_for("ns", "tq", "dev", Some("n1"))?;
1935 assert_eq!(pinned_n1.len(), 1);
1936 assert_eq!(pinned_n1[0].id(), n1_id);
1937 let pinned_n2 = registry.workers_for("ns", "tq", "dev", Some("n2"))?;
1938 assert_eq!(pinned_n2.len(), 1);
1939 assert_eq!(pinned_n2[0].id(), n2_id);
1940
1941 // select_worker honours the same filter.
1942 assert_eq!(
1943 registry
1944 .select_worker("ns", "tq", "dev", Some("n1"))?
1945 .map(|worker| worker.id()),
1946 Some(n1_id)
1947 );
1948
1949 // Pinned to a node with no worker finds no candidate (the dispatcher
1950 // then waits via the same no-worker path the existing test exercises).
1951 assert!(
1952 registry
1953 .workers_for("ns", "tq", "dev", Some("absent"))?
1954 .is_empty(),
1955 "a pin to a node with no worker yields no candidate"
1956 );
1957 assert!(
1958 registry
1959 .select_worker("ns", "tq", "dev", Some("absent"))?
1960 .is_none()
1961 );
1962
1963 on_n1.deregister()?;
1964 on_n2.deregister()?;
1965 Ok(())
1966 }
1967
1968 #[tokio::test]
1969 async fn shared_node_id_round_robins_across_workers() -> Result<(), ServerError> {
1970 // Two workers SHARING a node id in the same pool: a dispatch pinned to
1971 // that node round-robins across BOTH (node is locality, not process).
1972 let registry = ConnectedWorkerRegistry::default();
1973 let (a_tx, _a_rx) = mpsc::channel(1);
1974 let (b_tx, _b_rx) = mpsc::channel(1);
1975
1976 let worker_a = registry
1977 .accept_registration(
1978 &guard(),
1979 &caller("ns"),
1980 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1981 a_tx,
1982 )
1983 .await?;
1984 let worker_b = registry
1985 .accept_registration(
1986 &guard(),
1987 &caller("ns"),
1988 ®istration_full(&["ns"], "tq", "shared", &["dev"]),
1989 b_tx,
1990 )
1991 .await?;
1992 let a_id = worker_a.worker_id().ok_or_else(missing_id)?;
1993 let b_id = worker_b.worker_id().ok_or_else(missing_id)?;
1994
1995 let pinned = registry.workers_for("ns", "tq", "dev", Some("shared"))?;
1996 assert_eq!(
1997 pinned.len(),
1998 2,
1999 "both workers on the shared node are candidates"
2000 );
2001 let pinned_ids: BTreeSet<WorkerId> = pinned.iter().map(WorkerHandle::id).collect();
2002 assert_eq!(pinned_ids, BTreeSet::from([a_id, b_id]));
2003
2004 let first = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
2005 let second = registry.workers_for("ns", "tq", "dev", Some("shared"))?[0].id();
2006 assert_ne!(
2007 first, second,
2008 "a pin to a shared node round-robins across both workers on it"
2009 );
2010
2011 worker_a.deregister()?;
2012 worker_b.deregister()?;
2013 Ok(())
2014 }
2015
2016 #[tokio::test]
2017 async fn rotation_cursor_is_pruned_when_last_worker_leaves() -> Result<(), ServerError> {
2018 // The round-robin cursor is keyed on arbitrary caller-supplied strings;
2019 // it must not outlive the pool it rotates, or a never-dying server leaks
2020 // memory. After the last worker for a triple deregisters, no cursor for
2021 // that triple may remain.
2022 let registry = ConnectedWorkerRegistry::default();
2023 let (tx, _rx) = mpsc::channel(1);
2024 let worker = registry
2025 .accept_registration(
2026 &guard(),
2027 &caller("ns"),
2028 ®istration_full(&["ns"], "tq", "", &["dev"]),
2029 tx,
2030 )
2031 .await?;
2032
2033 // Drive the lazy cursor insert.
2034 let _ = registry.workers_for("ns", "tq", "dev", None)?;
2035 let key = ActivityKey::new(PoolAddress::new("ns", "tq"), "dev");
2036 assert!(
2037 registry.state()?.rotation.contains_key(&key),
2038 "a lookup must have created the rotation cursor"
2039 );
2040
2041 worker.deregister()?;
2042 let state = registry.state()?;
2043 assert!(
2044 !state.rotation.contains_key(&key),
2045 "the rotation cursor must be pruned once the last worker leaves"
2046 );
2047 assert!(
2048 !state.by_activity.contains_key(&key),
2049 "the activity bucket must also be gone"
2050 );
2051 Ok(())
2052 }
2053
2054 fn missing_id() -> ServerError {
2055 ServerError::lock_poisoned("registration unexpectedly missing a worker id")
2056 }
2057
2058 // ---- Minted-on-use (Control-Plane Phase 1) -----------------------------
2059
2060 fn namespace_store() -> Arc<dyn NamespaceStore> {
2061 Arc::new(aion_store::InMemoryStore::default())
2062 }
2063
2064 fn minting_registry(
2065 store: &Arc<dyn NamespaceStore>,
2066 policy: AutoCreate,
2067 ) -> ConnectedWorkerRegistry {
2068 ConnectedWorkerRegistry::default().with_namespace_minting(Arc::clone(store), policy)
2069 }
2070
2071 #[tokio::test]
2072 async fn open_register_mints_durable_record_and_is_idempotent() -> Result<(), ServerError> {
2073 let store = namespace_store();
2074 let registry = minting_registry(&store, AutoCreate::Open);
2075
2076 // First registration mints the namespace.
2077 let (tx_one, _rx_one) = mpsc::channel(1);
2078 let first = registry
2079 .accept_registration(
2080 &guard(),
2081 &caller("orders"),
2082 ®istration("orders", &["charge"]),
2083 tx_one,
2084 )
2085 .await?;
2086 let record = store
2087 .get_namespace("orders")
2088 .await?
2089 .ok_or_else(|| ServerError::namespace_denied("expected a minted record"))?;
2090 assert_eq!(record.name, "orders");
2091 assert_eq!(record.origin, NamespaceOrigin::WorkerMint);
2092
2093 // Re-registering the same namespace is idempotent: no duplicate row,
2094 // and the prior worker is unaffected.
2095 let (tx_two, _rx_two) = mpsc::channel(1);
2096 registry
2097 .accept_registration(
2098 &guard(),
2099 &caller("orders"),
2100 ®istration("orders", &["refund"]),
2101 tx_two,
2102 )
2103 .await?;
2104 let all = store.list_namespaces().await?;
2105 assert_eq!(
2106 all.iter().filter(|r| r.name == "orders").count(),
2107 1,
2108 "re-register must not create a duplicate namespace row"
2109 );
2110 drop(first);
2111 Ok(())
2112 }
2113
2114 #[tokio::test]
2115 async fn open_register_mints_each_namespace_in_a_multi_namespace_worker()
2116 -> Result<(), ServerError> {
2117 let store = namespace_store();
2118 let registry = minting_registry(&store, AutoCreate::Open);
2119 let (tx, _rx) = mpsc::channel(1);
2120
2121 registry
2122 .accept_registration(
2123 &guard(),
2124 &multi_caller(&["alpha", "beta"]),
2125 ®istration_full(&["alpha", "beta"], "", "", &["charge"]),
2126 tx,
2127 )
2128 .await?;
2129
2130 assert!(store.get_namespace("alpha").await?.is_some());
2131 assert!(store.get_namespace("beta").await?.is_some());
2132 Ok(())
2133 }
2134
2135 // ---- Placement admission (Control-Plane Phase 2, P2-I1) -----------------
2136
2137 /// Pre-mint `namespace` and set its placement to `Pinned{nodes}`, returning a
2138 /// minting registry over the same store so `accept_registration` reads the
2139 /// placement from the SAME durable record.
2140 async fn pinned_registry(
2141 store: &Arc<dyn NamespaceStore>,
2142 namespace: &str,
2143 nodes: &[&str],
2144 ) -> Result<ConnectedWorkerRegistry, ServerError> {
2145 store
2146 .register_namespace(namespace, NamespaceOrigin::Explicit)
2147 .await?;
2148 store
2149 .set_namespace_placement(
2150 namespace,
2151 NamespacePlacement::Pinned {
2152 nodes: nodes.iter().map(|n| (*n).to_owned()).collect(),
2153 },
2154 )
2155 .await?;
2156 Ok(minting_registry(store, AutoCreate::Open))
2157 }
2158
2159 /// A worker on a node IN the required set registers successfully into a
2160 /// `Pinned{n1}` namespace, and is reachable in the pool.
2161 #[tokio::test]
2162 async fn pinned_admits_a_worker_on_a_required_node() -> Result<(), ServerError> {
2163 let store = namespace_store();
2164 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2165 let (tx, _rx) = mpsc::channel(1);
2166
2167 let _registration = registry
2168 .accept_registration(
2169 &guard(),
2170 &caller("iso"),
2171 ®istration_full(&["iso"], "", "n1", &["charge"]),
2172 tx,
2173 )
2174 .await?;
2175
2176 assert_eq!(
2177 registry
2178 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", Some("n1"))?
2179 .len(),
2180 1,
2181 "an n1 worker must be admitted into the Pinned{{n1}} namespace's pool"
2182 );
2183 Ok(())
2184 }
2185
2186 /// A worker on a node NOT in the required set is rejected — the WHOLE
2187 /// registration fails (loud) and no worker is inserted. This would FAIL under
2188 /// no admission gate (the worker would join and steal Pinned dispatches).
2189 #[tokio::test]
2190 async fn pinned_rejects_a_wrong_node_worker() -> Result<(), ServerError> {
2191 let store = namespace_store();
2192 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2193 let (tx, _rx) = mpsc::channel(1);
2194
2195 let denied = registry
2196 .accept_registration(
2197 &guard(),
2198 &caller("iso"),
2199 ®istration_full(&["iso"], "", "n2", &["charge"]),
2200 tx,
2201 )
2202 .await;
2203 assert!(
2204 matches!(denied, Err(ServerError::Namespace { .. })),
2205 "a wrong-node (n2) worker must be rejected from a Pinned{{n1}} namespace"
2206 );
2207 assert!(
2208 registry
2209 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
2210 .is_empty(),
2211 "a rejected registration must not insert a worker on any node"
2212 );
2213 Ok(())
2214 }
2215
2216 /// A worker advertising NO node (`node == ""` → `None`) is rejected from a
2217 /// `Pinned{n1}` namespace: an unlabelled worker can never satisfy a hard pin.
2218 #[tokio::test]
2219 async fn pinned_rejects_a_node_less_worker() -> Result<(), ServerError> {
2220 let store = namespace_store();
2221 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2222 let (tx, _rx) = mpsc::channel(1);
2223
2224 let denied = registry
2225 .accept_registration(
2226 &guard(),
2227 &caller("iso"),
2228 ®istration_full(&["iso"], "", "", &["charge"]),
2229 tx,
2230 )
2231 .await;
2232 assert!(
2233 matches!(denied, Err(ServerError::Namespace { .. })),
2234 "a node-less worker must be rejected from a Pinned{{n1}} namespace"
2235 );
2236 assert!(
2237 registry
2238 .workers_for("iso", DEFAULT_TASK_QUEUE, "charge", None)?
2239 .is_empty(),
2240 "a rejected node-less registration must not insert a worker"
2241 );
2242 Ok(())
2243 }
2244
2245 /// Reject-WHOLE-registration (Open Decision 6): a worker serving BOTH a
2246 /// non-isolated namespace and a `Pinned{n1}` namespace on a wrong node is
2247 /// rejected entirely — the compliant namespace does NOT get a partial admit.
2248 #[tokio::test]
2249 async fn pinned_violation_rejects_the_whole_multi_namespace_registration()
2250 -> Result<(), ServerError> {
2251 let store = namespace_store();
2252 let registry = pinned_registry(&store, "iso", &["n1"]).await?;
2253 let (tx, _rx) = mpsc::channel(1);
2254
2255 let denied = registry
2256 .accept_registration(
2257 &guard(),
2258 &multi_caller(&["free", "iso"]),
2259 ®istration_full(&["free", "iso"], "", "n2", &["charge"]),
2260 tx,
2261 )
2262 .await;
2263 assert!(
2264 matches!(denied, Err(ServerError::Namespace { .. })),
2265 "a wrong-node worker serving a Pinned namespace fails the WHOLE registration"
2266 );
2267 assert!(
2268 registry
2269 .workers_for("free", DEFAULT_TASK_QUEUE, "charge", None)?
2270 .is_empty(),
2271 "the compliant namespace must NOT be partially admitted"
2272 );
2273 Ok(())
2274 }
2275
2276 /// Unplaced and Prefer namespaces are UNAFFECTED: a node-less worker registers
2277 /// normally (byte-identical to the pre-P2-I1 behaviour). Only Pinned gates.
2278 #[tokio::test]
2279 async fn unplaced_and_prefer_admission_is_unaffected_by_the_pinned_gate()
2280 -> Result<(), ServerError> {
2281 let store = namespace_store();
2282 // `unpl` is left Unplaced (default); `pref` is Prefer{n1}. A node-less
2283 // worker must be admitted into BOTH.
2284 store
2285 .register_namespace("pref", NamespaceOrigin::Explicit)
2286 .await?;
2287 store
2288 .set_namespace_placement(
2289 "pref",
2290 NamespacePlacement::Prefer {
2291 nodes: ["n1".to_owned()].into_iter().collect(),
2292 },
2293 )
2294 .await?;
2295 let registry = minting_registry(&store, AutoCreate::Open);
2296
2297 let (tx_a, _rx_a) = mpsc::channel(1);
2298 let _reg_a = registry
2299 .accept_registration(
2300 &guard(),
2301 &caller("unpl"),
2302 ®istration_full(&["unpl"], "", "", &["charge"]),
2303 tx_a,
2304 )
2305 .await?;
2306 let (tx_b, _rx_b) = mpsc::channel(1);
2307 let _reg_b = registry
2308 .accept_registration(
2309 &guard(),
2310 &caller("pref"),
2311 ®istration_full(&["pref"], "", "", &["charge"]),
2312 tx_b,
2313 )
2314 .await?;
2315
2316 assert_eq!(
2317 registry
2318 .workers_for("unpl", DEFAULT_TASK_QUEUE, "charge", None)?
2319 .len(),
2320 1,
2321 "an Unplaced namespace admits a node-less worker unchanged"
2322 );
2323 assert_eq!(
2324 registry
2325 .workers_for("pref", DEFAULT_TASK_QUEUE, "charge", None)?
2326 .len(),
2327 1,
2328 "a Prefer namespace admits a node-less worker unchanged (only Pinned gates)"
2329 );
2330 Ok(())
2331 }
2332
2333 /// A default (no-minter) registry is byte-identical: the placement gate is a
2334 /// no-op with no minter installed, so a node-less worker registers freely even
2335 /// though there is no way to have set a placement in the first place.
2336 #[tokio::test]
2337 async fn no_minter_registry_skips_the_placement_gate() -> Result<(), ServerError> {
2338 let registry = ConnectedWorkerRegistry::default();
2339 let (tx, _rx) = mpsc::channel(1);
2340 let _registration = registry
2341 .accept_registration(
2342 &guard(),
2343 &caller("plain"),
2344 ®istration_full(&["plain"], "", "", &["charge"]),
2345 tx,
2346 )
2347 .await?;
2348 assert_eq!(
2349 registry
2350 .workers_for("plain", DEFAULT_TASK_QUEUE, "charge", None)?
2351 .len(),
2352 1,
2353 "with no minter the placement gate is a no-op — registration is unchanged"
2354 );
2355 Ok(())
2356 }
2357
2358 #[tokio::test]
2359 async fn concurrent_registrations_for_a_new_namespace_create_exactly_one_record()
2360 -> Result<(), ServerError> {
2361 let store = namespace_store();
2362 let registry = minting_registry(&store, AutoCreate::Open);
2363
2364 let mut handles = Vec::new();
2365 for _ in 0..8 {
2366 let registry = registry.clone();
2367 handles.push(tokio::spawn(async move {
2368 let (tx, rx) = mpsc::channel(1);
2369 let outcome = registry
2370 .accept_registration(
2371 &guard(),
2372 &caller("rush"),
2373 ®istration("rush", &["charge"]),
2374 tx,
2375 )
2376 .await;
2377 // Keep the receiver alive for the duration of the registration.
2378 drop(rx);
2379 outcome.map(|registration| registration.worker_id())
2380 }));
2381 }
2382 for handle in handles {
2383 handle
2384 .await
2385 .map_err(|_| ServerError::lock_poisoned("registration task panicked"))??;
2386 }
2387
2388 let all = store.list_namespaces().await?;
2389 assert_eq!(
2390 all.iter().filter(|r| r.name == "rush").count(),
2391 1,
2392 "racing registrations must converge on exactly one durable record"
2393 );
2394 Ok(())
2395 }
2396
2397 #[tokio::test]
2398 async fn closed_rejects_unknown_namespace_and_does_not_create_it() -> Result<(), ServerError> {
2399 let store = namespace_store();
2400 let registry = minting_registry(&store, AutoCreate::Closed);
2401 let (tx, _rx) = mpsc::channel(1);
2402
2403 let denied = registry
2404 .accept_registration(
2405 &guard(),
2406 &caller("ghost"),
2407 ®istration("ghost", &["charge"]),
2408 tx,
2409 )
2410 .await;
2411 assert!(
2412 matches!(denied, Err(ServerError::Namespace { .. })),
2413 "closed policy must reject an unknown namespace"
2414 );
2415 assert!(
2416 store.get_namespace("ghost").await?.is_none(),
2417 "closed policy must NOT create the namespace it rejected"
2418 );
2419 let tq = DEFAULT_TASK_QUEUE;
2420 assert!(
2421 registry
2422 .workers_for("ghost", tq, "charge", None)?
2423 .is_empty(),
2424 "a rejected registration must not insert a worker"
2425 );
2426 Ok(())
2427 }
2428
2429 #[tokio::test]
2430 async fn closed_admits_a_known_namespace() -> Result<(), ServerError> {
2431 let store = namespace_store();
2432 // Pre-mint the namespace (the POST /namespaces escape hatch's effect).
2433 store
2434 .register_namespace("known", NamespaceOrigin::Explicit)
2435 .await?;
2436 let registry = minting_registry(&store, AutoCreate::Closed);
2437 let (tx, _rx) = mpsc::channel(1);
2438
2439 // Bind the registration token: dropping it deregisters the worker.
2440 let _registration = registry
2441 .accept_registration(
2442 &guard(),
2443 &caller("known"),
2444 ®istration("known", &["charge"]),
2445 tx,
2446 )
2447 .await?;
2448 let tq = DEFAULT_TASK_QUEUE;
2449 assert_eq!(
2450 registry.workers_for("known", tq, "charge", None)?.len(),
2451 1,
2452 "a known namespace must register under closed policy"
2453 );
2454 Ok(())
2455 }
2456
2457 #[tokio::test]
2458 async fn no_minter_leaves_registration_untouched() -> Result<(), ServerError> {
2459 // The default registry installs no minter: registration succeeds and
2460 // never touches any namespace registry (byte-identical legacy path).
2461 let registry = ConnectedWorkerRegistry::default();
2462 let (tx, _rx) = mpsc::channel(1);
2463 let _registration = registry
2464 .accept_registration(
2465 &guard(),
2466 &caller("orders"),
2467 ®istration("orders", &["charge"]),
2468 tx,
2469 )
2470 .await?;
2471 let tq = DEFAULT_TASK_QUEUE;
2472 assert_eq!(registry.workers_for("orders", tq, "charge", None)?.len(), 1);
2473 Ok(())
2474 }
2475
2476 /// R1 census: an address no worker has ever served reports an empty fleet
2477 /// and no poller age at all — "never seen" is not "seen long ago".
2478 #[test]
2479 fn a_never_served_address_censuses_empty_with_no_poller_age() -> Result<(), ServerError> {
2480 let registry = ConnectedWorkerRegistry::default();
2481 let census = registry.pool_census("default", "general", "greet", None)?;
2482 assert_eq!(census.workers_in_pool, 0);
2483 assert_eq!(census.workers_serving_activity, 0);
2484 assert_eq!(census.compatible_workers, 0);
2485 assert_eq!(census.last_compatible_poller_age, None);
2486 assert!(!census.is_served());
2487 Ok(())
2488 }
2489
2490 /// R1 census: pool membership, activity coverage, and node coverage are
2491 /// three separate counts — that separation is what tells `NO_LIVE_POLLERS`
2492 /// apart from `POLLERS_INCOMPATIBLE`.
2493 #[test]
2494 fn the_census_separates_pool_activity_and_node_coverage() -> Result<(), ServerError> {
2495 let registry = ConnectedWorkerRegistry::default();
2496 let (tx, _rx) = mpsc::channel(1);
2497 let _worker = registry.register_namespaces(
2498 [String::from("default")],
2499 "general",
2500 Some(String::from("n1")),
2501 [String::from("greet")].iter(),
2502 tx,
2503 )?;
2504
2505 let unpinned = registry.pool_census("default", "general", "greet", None)?;
2506 assert_eq!(unpinned.workers_in_pool, 1);
2507 assert_eq!(unpinned.workers_serving_activity, 1);
2508 assert_eq!(unpinned.compatible_workers, 1);
2509 assert_eq!(unpinned.last_compatible_poller_age, Some(Duration::ZERO));
2510
2511 // Same pool, activity nobody advertises: pollers are live but
2512 // incompatible.
2513 let other_activity = registry.pool_census("default", "general", "settle", None)?;
2514 assert_eq!(other_activity.workers_in_pool, 1);
2515 assert_eq!(other_activity.workers_serving_activity, 0);
2516 assert_eq!(other_activity.compatible_workers, 0);
2517
2518 // Same activity, wrong node: coverage exists in the pool but not for
2519 // this dispatch.
2520 let wrong_node = registry.pool_census("default", "general", "greet", Some("n2"))?;
2521 assert_eq!(wrong_node.workers_serving_activity, 1);
2522 assert_eq!(wrong_node.compatible_workers, 0);
2523 assert_eq!(wrong_node.last_compatible_poller_age, None);
2524 Ok(())
2525 }
2526
2527 /// ACCEPTANCE (b) FOR #58, wiring half: a worker whose PROCESS looks alive
2528 /// but which the server cannot REACH on the push leg must lose dispatch
2529 /// eligibility — and losing it must actually remove it from selection.
2530 ///
2531 /// This is the composition the lease split exists for. The tracker half is
2532 /// pinned in `heartbeat::reachability_tests`; this pins the half that turns
2533 /// that verdict into a dispatch decision, using the SAME two types the
2534 /// liveness probe wires together — a real [`HeartbeatTracker`] census
2535 /// feeding a real [`ConnectedWorkerRegistry`].
2536 ///
2537 /// The staging is Finding B exactly: the worker's own liveness pump keeps
2538 /// beating (`record_connection_activity`) well past the window, so every
2539 /// signal that says "this process is alive" says so — and NO ping is ever
2540 /// answered, so the one signal that says "the server can reach it" is
2541 /// absent. Before the split those were one lease and the pump's weaker
2542 /// positive evidence buried the ping's negative evidence.
2543 ///
2544 /// The control is the second half of the test: the same worker, reachable,
2545 /// must still be selected. Without it a registry that selected NOBODY would
2546 /// satisfy the first assertion and prove nothing.
2547 #[test]
2548 fn a_pump_alive_worker_the_server_cannot_reach_is_not_selected() -> Result<(), ServerError> {
2549 const WINDOW: Duration = Duration::from_secs(30);
2550
2551 let registry = ConnectedWorkerRegistry::default();
2552 let tracker = HeartbeatTracker::new(WINDOW);
2553 let (tx, _rx) = mpsc::channel(1);
2554 let worker = registry.register_namespaces(
2555 [String::from("default")],
2556 "general",
2557 None,
2558 [String::from("greet")].iter(),
2559 tx,
2560 )?;
2561 let Some(worker_id) = worker.worker_id() else {
2562 return Err(test_failure("registration carries an id"));
2563 };
2564 let start = Instant::now();
2565 tracker.register_connection(worker_id, start)?;
2566
2567 // Publishing the tracker's verdict is one motion the liveness probe
2568 // performs at the end of every round, and this test performs it four
2569 // times. Written once so the four publications cannot drift apart and
2570 // quietly stop testing the same thing.
2571 let publish = |now: Instant| -> Result<(), ServerError> {
2572 registry.set_dispatch_ineligible(
2573 tracker
2574 .unreachable_workers(now)?
2575 .into_iter()
2576 .map(|excluded| excluded.worker_id)
2577 .collect(),
2578 )
2579 };
2580
2581 // Baseline: the worker SERVES ITS PROBATION and is therefore selectable.
2582 // Registration alone grants nothing — the handshake ack is sent, never
2583 // confirmed received — so eligibility here is earned by answered pings.
2584 // If this did not hold, no worker could ever be dispatched to.
2585 for _ in 0..DISPATCH_PROBATION_PINGS {
2586 assert!(
2587 tracker.record_dispatch_reachability(worker_id, start)?,
2588 "the worker is tracked while it serves its probation"
2589 );
2590 }
2591 publish(start)?;
2592 assert!(
2593 registry
2594 .select_worker("default", "general", "greet", None)?
2595 .is_some(),
2596 "a worker that has answered a full run of pings must be selectable: this is the \
2597 control, and without it a registry selecting NOBODY would satisfy every assertion \
2598 below"
2599 );
2600
2601 // Now the poisoned-connection shape: the pump beats past the window and
2602 // not one ping is answered. Each unanswered probe is recorded exactly as
2603 // the liveness probe records it, because that is the trajectory a real
2604 // poisoned connection takes — staleness alone would be a shape only a
2605 // stopped probe could produce.
2606 let much_later = start + WINDOW * 4;
2607 assert!(
2608 tracker.record_connection_activity(worker_id, much_later)?,
2609 "the worker is still tracked; its process is plainly alive"
2610 );
2611 assert!(
2612 tracker.record_dispatch_unreachable(worker_id)?,
2613 "the probe fired and went unanswered"
2614 );
2615 publish(much_later)?;
2616
2617 assert!(
2618 registry.is_dispatch_ineligible(worker_id)?,
2619 "a worker the server cannot reach must be marked ineligible however alive its \
2620 process looks"
2621 );
2622 assert!(
2623 registry
2624 .select_worker("default", "general", "greet", None)?
2625 .is_none(),
2626 "an ineligible worker must not be SELECTED: selecting one produces a dispatch that \
2627 can only fail, and on the liminal transport it fails by consuming connection \
2628 capacity — making the unreachability worse"
2629 );
2630
2631 // The recovery path: exclusion must be WITHDRAWABLE, or a single bad
2632 // round would strand a healthy worker forever. It is withdrawn by a
2633 // served probation, not by one answer — and the intermediate assertion
2634 // below pins that distinction rather than assuming it.
2635 assert!(
2636 tracker.record_dispatch_reachability(worker_id, much_later)?,
2637 "the worker is still tracked"
2638 );
2639 publish(much_later)?;
2640 assert!(
2641 registry.is_dispatch_ineligible(worker_id)?,
2642 "one answer part-way through a fresh probation must NOT restore eligibility: a link \
2643 answering one probe in three would otherwise flap in and out of selection"
2644 );
2645
2646 for _ in 1..DISPATCH_PROBATION_PINGS {
2647 assert!(
2648 tracker.record_dispatch_reachability(worker_id, much_later)?,
2649 "the worker is still tracked"
2650 );
2651 }
2652 publish(much_later)?;
2653 assert!(
2654 !registry.is_dispatch_ineligible(worker_id)?,
2655 "an answered ping must clear the exclusion"
2656 );
2657 assert!(
2658 registry
2659 .select_worker("default", "general", "greet", None)?
2660 .is_some(),
2661 "and the worker must be selectable again"
2662 );
2663 Ok(())
2664 }
2665
2666 /// R1 census: after the last compatible worker leaves, the address reports
2667 /// how long ago it was served — the age an operator sees on the parked
2668 /// dispatch's WARN.
2669 #[test]
2670 fn a_departed_worker_leaves_a_last_compatible_poller_age() -> Result<(), ServerError> {
2671 let registry = ConnectedWorkerRegistry::default();
2672 let (tx, _rx) = mpsc::channel(1);
2673 let worker = registry.register_namespaces(
2674 [String::from("default")],
2675 "general",
2676 None,
2677 [String::from("greet")].iter(),
2678 tx,
2679 )?;
2680 worker.deregister()?;
2681
2682 let census = registry.pool_census("default", "general", "greet", None)?;
2683 assert_eq!(census.workers_in_pool, 0);
2684 assert_eq!(census.compatible_workers, 0);
2685 let age = census
2686 .last_compatible_poller_age
2687 .ok_or_else(|| test_failure("a departed worker must leave an age behind"))?;
2688 assert!(
2689 age < Duration::from_secs(60),
2690 "the recorded departure is implausibly old: {age:?}"
2691 );
2692
2693 // A worker returning to the address clears the departure record: the
2694 // map that answers "how long ago" never accumulates live addresses.
2695 let (tx, _rx) = mpsc::channel(1);
2696 let _back = registry.register_namespaces(
2697 [String::from("default")],
2698 "general",
2699 None,
2700 [String::from("greet")].iter(),
2701 tx,
2702 )?;
2703 let served = registry.pool_census("default", "general", "greet", None)?;
2704 assert_eq!(served.last_compatible_poller_age, Some(Duration::ZERO));
2705 assert!(served.is_served());
2706 Ok(())
2707 }
2708}