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