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