Skip to main content

aion_server/worker/
registry.rs

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