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