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