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