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