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