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