Skip to main content

aion_server/
cluster.rs

1//! SS-5b: automatic multi-node failover detection.
2//!
3//! [`ClusterSupervisor`] is the production counterpart to the manual
4//! `Engine::adopt_shards` trigger proven in the SS-5 demo. It runs a background
5//! task that watches the liveness of every peer that owns shards and, when a
6//! peer's replication link drops and stays down past a debounce threshold,
7//! calls `adopt_shards` for that peer's shards ITSELF — no human in the loop.
8//!
9//! ## How peer-down is detected
10//!
11//! The liveness signal is the haematite distribution link state
12//! ([`HaematiteStore::peer_connected`]): beamr's OTP distribution tears the
13//! connection down (read-loop EOF → deregister) the instant the peer's process
14//! dies, so `peer_connected` flips to `false` on a real `kill -9` exactly as it
15//! does on a graceful drop. It is a true socket-liveness signal, not a heartbeat
16//! heuristic.
17//!
18//! ## Debounce
19//!
20//! A single missed poll is not a death: a transient blip must not trigger a
21//! disruptive shard adoption. The supervisor requires `confirmations`
22//! CONSECUTIVE polls observing the peer disconnected before it acts. Any single
23//! reconnect observation resets the counter. Once a peer's shards are adopted it
24//! is marked handled and not re-adopted while it stays down (adoption is itself
25//! idempotent, but re-running it every tick would be wasteful); a later reconnect
26//! clears the handled mark so a flapping peer that genuinely dies again is
27//! re-adopted.
28//!
29//! ## Scope
30//!
31//! Only ever constructed for a distributed (`[store.cluster]`) boot. A single-node / non-clustered server
32//! never spawns it, so default behaviour is unchanged.
33
34use std::collections::BTreeMap;
35use std::sync::Arc;
36use std::time::Duration;
37
38use aion::Engine;
39use aion_core::ClusterEvent;
40
41use crate::cluster_publisher::ClusterEventPublisher;
42
43/// The liveness signal the supervisor polls. Implemented by [`HaematiteStore`]
44/// in production and by a fake in tests, so the debounce/adopt logic is verified
45/// without standing up a real cluster every time.
46pub trait PeerLiveness: Send + Sync + 'static {
47    /// Whether the peer named `peer_name` currently holds a live replication link.
48    fn peer_connected(&self, peer_name: &str) -> bool;
49
50    /// The distribution name currently RECORDED as `shard`'s owner in the cluster
51    /// shard-owner directory (SS-3), or `None` when no record exists. Used by the
52    /// adopt pre-check to detect a shard already adopted-and-published by another
53    /// survivor, so this supervisor does not race a second adoption of it. Mirrors
54    /// `routing::directory::resolve_from_record`'s down-owner detection: a record
55    /// naming a LIVE peer means handled-elsewhere (skip); a record naming a peer
56    /// that is itself down is adoptable (the recorded owner has since died).
57    ///
58    /// A failed read returns `None` ("no directory opinion") so a transient read
59    /// failure never strands a dead peer's shards.
60    fn read_shard_owner(&self, shard: usize) -> Option<String>;
61}
62
63impl PeerLiveness for aion_store_haematite::HaematiteStore {
64    fn peer_connected(&self, peer_name: &str) -> bool {
65        Self::peer_connected(self, peer_name)
66    }
67
68    fn read_shard_owner(&self, shard: usize) -> Option<String> {
69        // A failed read is "no directory opinion": fall through to adoption.
70        Self::read_shard_owner(self, shard).ok().flatten()
71    }
72}
73
74/// The failover action the supervisor invokes when a peer is confirmed down.
75/// Implemented by [`Engine`] in production and by a fake in tests.
76#[async_trait::async_trait]
77pub trait ShardAdopter: Send + Sync + 'static {
78    /// Adopt `shards` from a dead peer: elect + union-merge + resume.
79    async fn adopt_shards(&self, shards: &[usize]) -> Result<(), String>;
80}
81
82#[async_trait::async_trait]
83impl ShardAdopter for Engine {
84    async fn adopt_shards(&self, shards: &[usize]) -> Result<(), String> {
85        Engine::adopt_shards(self, shards)
86            .await
87            .map_err(|error| error.to_string())
88    }
89}
90
91/// [`ShardAdopter`] over the live engine that, when the durable outbox is
92/// commissioned, re-runs the terminal-workflow outbox settlement sweep (#253)
93/// after each successful adoption.
94///
95/// The adoption fence has already widened this node's owned-shard scope by the
96/// time `Engine::adopt_shards` returns (mirroring the paused-runs `extend` the
97/// engine runs at the same point), so the sweep now enumerates the adopted
98/// shards' unsettled rows and settles those whose workflow is terminal —
99/// closing the failover half of the incident window: a dead node's stranded
100/// row for a terminal workflow must not be re-armed and redelivered by its
101/// adopter. A sweep failure is loud but never fails the adoption (the shards
102/// are durably adopted; the reconciler liveness gate remains the backstop).
103pub struct OutboxSettlingAdopter {
104    engine: Arc<Engine>,
105    outbox_store: Option<Arc<dyn aion_store::OutboxStore>>,
106}
107
108impl OutboxSettlingAdopter {
109    /// Build an adopter over the live engine; `outbox_store` is `Some` exactly
110    /// when the durable outbox is commissioned (there is nothing to settle
111    /// otherwise).
112    #[must_use]
113    pub fn new(
114        engine: Arc<Engine>,
115        outbox_store: Option<Arc<dyn aion_store::OutboxStore>>,
116    ) -> Self {
117        Self {
118            engine,
119            outbox_store,
120        }
121    }
122}
123
124#[async_trait::async_trait]
125impl ShardAdopter for OutboxSettlingAdopter {
126    async fn adopt_shards(&self, shards: &[usize]) -> Result<(), String> {
127        ShardAdopter::adopt_shards(self.engine.as_ref(), shards).await?;
128        let Some(outbox_store) = &self.outbox_store else {
129            return Ok(());
130        };
131        match crate::worker::settle_terminal_outbox_rows(
132            self.engine.store().as_ref(),
133            outbox_store.as_ref(),
134        )
135        .await
136        {
137            Ok(settled) if settled.is_empty() => {}
138            Ok(settled) => {
139                tracing::info!(
140                    ?shards,
141                    settled = settled.len(),
142                    "adoption sweep settled stranded outbox rows for terminal workflows"
143                );
144            }
145            Err(error) => {
146                tracing::error!(
147                    ?shards,
148                    %error,
149                    "adoption sweep failed to settle terminal workflows' outbox rows; \
150                     the reconciler liveness gate remains the backstop"
151                );
152            }
153        }
154        Ok(())
155    }
156}
157
158/// One peer the supervisor watches: its distribution name and the shards it owns
159/// (which this node will adopt if the peer dies).
160#[derive(Clone, Debug, PartialEq, Eq)]
161pub struct WatchedPeer {
162    /// The peer's globally-unique distribution name.
163    pub name: String,
164    /// The shards this peer owns; adopted on confirmed death.
165    pub owned_shards: Vec<usize>,
166}
167
168/// Tuning for the supervisor's poll loop.
169#[derive(Clone, Copy, Debug)]
170pub struct SupervisorConfig {
171    /// Interval between liveness polls.
172    pub poll_interval: Duration,
173    /// Consecutive disconnected observations required before adopting (debounce).
174    /// Must be at least one.
175    pub confirmations: u32,
176}
177
178/// Per-peer debounce state tracked across poll ticks.
179#[derive(Default)]
180struct PeerState {
181    /// Consecutive ticks this peer has been observed disconnected.
182    consecutive_down: u32,
183    /// Whether this peer's shards have already been adopted while down.
184    adopted: bool,
185}
186
187/// Watches peer liveness and auto-adopts a dead peer's shards (SS-5b).
188pub struct ClusterSupervisor<L: PeerLiveness, A: ShardAdopter> {
189    liveness: Arc<L>,
190    adopter: Arc<A>,
191    peers: Vec<WatchedPeer>,
192    config: SupervisorConfig,
193    state: BTreeMap<String, PeerState>,
194    /// WS3 cluster-event sink. `None` keeps every existing test compiling and
195    /// keeps a non-ops-console boot silent; when present, `tick()` emits a delta at
196    /// each of its existing branch points. The publisher fans out to live
197    /// ops console subscribers and is a no-op with none attached.
198    publisher: Option<Arc<ClusterEventPublisher>>,
199    /// This node's distribution name, stamped into `ShardAdopted.adopted_by` and
200    /// the supervisor lifecycle events. Empty when unknown (no emit honesty cost:
201    /// the field is still the real configured value or absent).
202    self_node: String,
203}
204
205impl<L: PeerLiveness, A: ShardAdopter> ClusterSupervisor<L, A> {
206    /// Build a supervisor over `peers`, polling `liveness` and calling
207    /// `adopter.adopt_shards` on confirmed peer death. Peers with no owned shards
208    /// are dropped from the watch set (nothing to adopt for them).
209    #[must_use]
210    pub fn new(
211        liveness: Arc<L>,
212        adopter: Arc<A>,
213        peers: Vec<WatchedPeer>,
214        config: SupervisorConfig,
215    ) -> Self {
216        let peers: Vec<WatchedPeer> = peers
217            .into_iter()
218            .filter(|peer| !peer.owned_shards.is_empty())
219            .collect();
220        let state = peers
221            .iter()
222            .map(|peer| (peer.name.clone(), PeerState::default()))
223            .collect();
224        Self {
225            liveness,
226            adopter,
227            peers,
228            config,
229            state,
230            publisher: None,
231            self_node: String::new(),
232        }
233    }
234
235    /// Attach the WS3 cluster-event publisher and this node's name so `tick()`
236    /// emits topology deltas. Pure builder addition — a supervisor without it
237    /// behaves exactly as before (every existing test passes `new` only).
238    #[must_use]
239    pub fn with_publisher(
240        mut self,
241        publisher: Arc<ClusterEventPublisher>,
242        self_node: impl Into<String>,
243    ) -> Self {
244        self.publisher = Some(publisher);
245        self.self_node = self_node.into();
246        self
247    }
248
249    /// Emit a cluster event through the attached publisher, if any. The `build`
250    /// closure receives the publisher-stamped meta; with no publisher attached
251    /// this is a no-op.
252    fn emit<F>(&self, build: F)
253    where
254        F: FnOnce(aion_core::ClusterEventMeta) -> ClusterEvent,
255    {
256        if let Some(publisher) = &self.publisher {
257            drop(publisher.emit(build));
258        }
259    }
260
261    /// Whether this supervisor watches any peer (false when no peer declared
262    /// owned shards — the loop would do nothing, so the caller can skip spawning).
263    #[must_use]
264    pub fn watches_any(&self) -> bool {
265        !self.peers.is_empty()
266    }
267
268    /// Borrow the adopter (the engine, in production) this supervisor drives.
269    /// Lets a test inspect the engine it auto-adopts onto after the loss.
270    #[must_use]
271    pub fn adopter(&self) -> &A {
272        &self.adopter
273    }
274
275    /// Run ONE poll tick: observe every watched peer's liveness, advance the
276    /// debounce counters, and adopt the shards of any peer that has now been down
277    /// for `confirmations` consecutive ticks and is not yet adopted.
278    ///
279    /// Returned is the list of peer names adopted on THIS tick (empty on a quiet
280    /// tick), so a test can assert exactly when adoption fires. Extracted from the
281    /// loop so the debounce decision is unit-testable without real time.
282    pub async fn tick(&mut self) -> Vec<String> {
283        let mut adopted_now = Vec::new();
284        // Collect emits to fire AFTER the borrow of `self.state` ends: the emit
285        // path borrows `&self` (for the publisher) while the loop holds `&mut
286        // self.state` via `entry`, so deltas are queued and flushed post-loop.
287        let mut pending: Vec<ClusterEvent> = Vec::new();
288        let confirmations = self.config.confirmations;
289        for peer in &self.peers {
290            let connected = self.liveness.peer_connected(&peer.name);
291            let entry = self.state.entry(peer.name.clone()).or_default();
292            if connected {
293                // RECOVERY EMIT: capture the prior-down signal BEFORE the reset,
294                // or every tick would look freshly connected and no recovery
295                // event would ever fire.
296                let was_down = entry.consecutive_down > 0 || entry.adopted;
297                entry.consecutive_down = 0;
298                entry.adopted = false;
299                if was_down {
300                    pending.push(ClusterEvent::PeerConnected {
301                        meta: placeholder_meta(),
302                        peer_name: peer.name.clone(),
303                        forward_addr: None,
304                    });
305                }
306                continue;
307            }
308            entry.consecutive_down = entry.consecutive_down.saturating_add(1);
309            let consecutive_down = entry.consecutive_down;
310            let confirmed = consecutive_down >= confirmations;
311            // Every tick a peer is observed down is a delta; `confirmed` flips
312            // once the debounce threshold authorizes adoption.
313            pending.push(ClusterEvent::PeerDisconnected {
314                meta: placeholder_meta(),
315                peer_name: peer.name.clone(),
316                consecutive_down,
317                confirmed,
318            });
319            if entry.adopted || consecutive_down < confirmations {
320                continue;
321            }
322            // Pre-check: skip any of this peer's shards already published to a
323            // DIFFERENT live owner — another survivor has adopted them, so racing
324            // a second adoption would be wasted work (the fence would drop us
325            // anyway). A record naming a peer that is itself down is adoptable (the
326            // recorded owner has since died); no record is adoptable too. Mirrors
327            // routing::directory::resolve_from_record's down-owner detection.
328            if Self::all_shards_handled_elsewhere(
329                self.liveness.as_ref(),
330                &peer.name,
331                &peer.owned_shards,
332            ) {
333                // Every shard is already served by a live owner: mark handled so
334                // the supervisor does NOT retry-loop on shards another node owns.
335                entry.adopted = true;
336                let held_by = Self::live_owner_of(self.liveness.as_ref(), &peer.owned_shards)
337                    .unwrap_or_default();
338                pending.push(ClusterEvent::ShardAdoptionSkipped {
339                    meta: placeholder_meta(),
340                    shards: peer.owned_shards.clone(),
341                    from_peer: peer.name.clone(),
342                    held_by,
343                });
344                tracing::info!(
345                    peer = %peer.name,
346                    shards = ?peer.owned_shards,
347                    "downed peer's shards already adopted by another live owner; skipping"
348                );
349                continue;
350            }
351            match self.adopter.adopt_shards(&peer.owned_shards).await {
352                Ok(()) => {
353                    entry.adopted = true;
354                    adopted_now.push(peer.name.clone());
355                    pending.push(ClusterEvent::ShardAdopted {
356                        meta: placeholder_meta(),
357                        shards: peer.owned_shards.clone(),
358                        from_peer: peer.name.clone(),
359                        adopted_by: self.self_node.clone(),
360                    });
361                    tracing::info!(
362                        peer = %peer.name,
363                        shards = ?peer.owned_shards,
364                        "cluster supervisor adopted a downed peer's shards (SS-5b auto-failover)"
365                    );
366                }
367                Err(error) => {
368                    pending.push(ClusterEvent::ShardAdoptionFailed {
369                        meta: placeholder_meta(),
370                        shards: peer.owned_shards.clone(),
371                        from_peer: peer.name.clone(),
372                        error: error.clone(),
373                    });
374                    // Leave `adopted` false so the next tick retries: a
375                    // quorum-unavailable / transport adopt error must not strand
376                    // the dead peer's shards forever (the retry contract). Note a
377                    // fenced (NotOwner) shard is NOT surfaced here — the engine's
378                    // clean-partial adopt drops a deposed shard internally and
379                    // returns Ok, and the pre-check above already short-circuits a
380                    // shard another LIVE owner holds, so this arm is reached only
381                    // for genuinely retryable faults.
382                    tracing::warn!(
383                        peer = %peer.name,
384                        shards = ?peer.owned_shards,
385                        %error,
386                        "cluster supervisor failed to adopt a downed peer's shards; will retry"
387                    );
388                }
389            }
390        }
391        // Flush queued deltas now that the `&mut self.state` borrow is released:
392        // each is re-stamped with a real publisher seq+instant (the placeholder
393        // meta is discarded). With no publisher attached this is a no-op.
394        for event in pending {
395            self.emit(|meta| with_meta(event, meta));
396        }
397        adopted_now
398    }
399
400    /// The live owner currently recorded for the first of `shards` that names a
401    /// connected third party, for the `ShardAdoptionSkipped.held_by` field. Reads
402    /// only real directory records; returns `None` if none is live-held.
403    fn live_owner_of(liveness: &L, shards: &[usize]) -> Option<String> {
404        shards.iter().find_map(|&shard| {
405            liveness
406                .read_shard_owner(shard)
407                .filter(|owner| liveness.peer_connected(owner))
408        })
409    }
410
411    /// Whether EVERY shard in `shards` is already published to a DIFFERENT live
412    /// owner — i.e. another survivor has adopted them, so this supervisor has
413    /// nothing left to do for the dead `peer_name`. A shard is "handled elsewhere"
414    /// only when its directory record names a peer that is BOTH not the dead peer
415    /// AND currently connected; a record naming the dead peer (or a peer now down)
416    /// or no record at all means the shard is still adoptable. Empty `shards`
417    /// is vacuously handled, but such peers are filtered out at construction.
418    fn all_shards_handled_elsewhere(liveness: &L, peer_name: &str, shards: &[usize]) -> bool {
419        !shards.is_empty()
420            && shards.iter().all(|&shard| {
421                liveness.read_shard_owner(shard).is_some_and(|owner| {
422                    // The recorded owner is a LIVE third party (not the dead peer):
423                    // that survivor serves it. A record naming the dead peer itself
424                    // is stale (it has since died) and remains adoptable.
425                    owner != peer_name && liveness.peer_connected(&owner)
426                })
427            })
428    }
429
430    /// Drive the poll loop until `shutdown` flips true, ticking every
431    /// `poll_interval`. Consumes `self`; spawn it as a background task.
432    pub async fn run(mut self, mut shutdown: tokio::sync::watch::Receiver<bool>) {
433        // Lifecycle EMIT: the supervisor is running on this node (ADR-019 calm
434        // state distinguishes "running, all healthy" from "not running").
435        let self_node = self.self_node.clone();
436        self.emit(|meta| ClusterEvent::SupervisorStarted {
437            meta,
438            node: self_node.clone(),
439        });
440        let mut interval = tokio::time::interval(self.config.poll_interval);
441        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
442        loop {
443            tokio::select! {
444                _ = interval.tick() => {
445                    drop(self.tick().await);
446                }
447                changed = shutdown.changed() => {
448                    if changed.is_err() || *shutdown.borrow() {
449                        break;
450                    }
451                }
452            }
453        }
454        // Lifecycle EMIT: clean drain/shutdown — the ops console can distinguish a
455        // stopped supervisor from "all peers healthy" (ADR-019).
456        let self_node = self.self_node.clone();
457        self.emit(|meta| ClusterEvent::SupervisorStopped {
458            meta,
459            node: self_node.clone(),
460        });
461    }
462}
463
464/// A placeholder meta used while a [`ClusterEvent`] is queued inside `tick()`'s
465/// `&mut self.state` borrow; it is ALWAYS replaced by the publisher-stamped meta
466/// in [`with_meta`] at flush time, so a placeholder seq never reaches the wire.
467fn placeholder_meta() -> aion_core::ClusterEventMeta {
468    aion_core::ClusterEventMeta {
469        cluster_seq: 0,
470        observed_at: chrono::Utc::now(),
471    }
472}
473
474/// Replace a queued event's placeholder meta with the publisher-stamped one.
475///
476/// The peer/shard topology arms (the only events `tick()` queues with a
477/// placeholder) are handled here; worker-lifecycle and the
478/// publisher-direct-emit variants delegate to [`with_meta_worker_lifecycle`] to
479/// keep each function under the house line limit.
480fn with_meta(event: ClusterEvent, meta: aion_core::ClusterEventMeta) -> ClusterEvent {
481    match event {
482        ClusterEvent::PeerAdded {
483            peer_name,
484            forward_addr,
485            ..
486        } => ClusterEvent::PeerAdded {
487            meta,
488            peer_name,
489            forward_addr,
490        },
491        ClusterEvent::PeerConnected {
492            peer_name,
493            forward_addr,
494            ..
495        } => ClusterEvent::PeerConnected {
496            meta,
497            peer_name,
498            forward_addr,
499        },
500        ClusterEvent::PeerDisconnected {
501            peer_name,
502            consecutive_down,
503            confirmed,
504            ..
505        } => ClusterEvent::PeerDisconnected {
506            meta,
507            peer_name,
508            consecutive_down,
509            confirmed,
510        },
511        ClusterEvent::ShardAdopted {
512            shards,
513            from_peer,
514            adopted_by,
515            ..
516        } => ClusterEvent::ShardAdopted {
517            meta,
518            shards,
519            from_peer,
520            adopted_by,
521        },
522        ClusterEvent::ShardAdoptionFailed {
523            shards,
524            from_peer,
525            error,
526            ..
527        } => ClusterEvent::ShardAdoptionFailed {
528            meta,
529            shards,
530            from_peer,
531            error,
532        },
533        ClusterEvent::ShardAdoptionSkipped {
534            shards,
535            from_peer,
536            held_by,
537            ..
538        } => ClusterEvent::ShardAdoptionSkipped {
539            meta,
540            shards,
541            from_peer,
542            held_by,
543        },
544        other => with_meta_worker_lifecycle(other, meta),
545    }
546}
547
548/// Meta re-stamp for the worker-lifecycle, supervisor, and `NamespaceCreated`
549/// variants (the tail of [`with_meta`]'s exhaustive match).
550///
551/// `NamespaceCreated` is emitted directly through the publisher (which stamps the
552/// real meta), never queued inside `tick()` with a placeholder, so its arm is
553/// unreachable in practice; it re-stamps faithfully to keep the match exhaustive
554/// without a wildcard that could silently swallow a future variant. The
555/// peer/shard variants never reach here ([`with_meta`] handles them), so they are
556/// `unreachable!` rather than silently mis-stamped.
557fn with_meta_worker_lifecycle(
558    event: ClusterEvent,
559    meta: aion_core::ClusterEventMeta,
560) -> ClusterEvent {
561    match event {
562        ClusterEvent::WorkerConnected {
563            worker_id,
564            namespaces,
565            task_queue,
566            transport,
567            node,
568            deployment,
569            deployment_association,
570            ..
571        } => ClusterEvent::WorkerConnected {
572            meta,
573            worker_id,
574            namespaces,
575            task_queue,
576            transport,
577            node,
578            deployment,
579            deployment_association,
580        },
581        ClusterEvent::WorkerDisconnected {
582            worker_id,
583            namespaces,
584            reason,
585            ..
586        } => ClusterEvent::WorkerDisconnected {
587            meta,
588            worker_id,
589            namespaces,
590            reason,
591        },
592        ClusterEvent::SupervisorStarted { node, .. } => {
593            ClusterEvent::SupervisorStarted { meta, node }
594        }
595        ClusterEvent::SupervisorStopped { node, .. } => {
596            ClusterEvent::SupervisorStopped { meta, node }
597        }
598        ClusterEvent::NamespaceCreated {
599            name,
600            created_at,
601            origin,
602            ..
603        } => ClusterEvent::NamespaceCreated {
604            meta,
605            name,
606            created_at,
607            origin,
608        },
609        // Like `NamespaceCreated`, emitted directly through the publisher (which
610        // stamps the real meta), never queued in `tick()`; re-stamped faithfully
611        // to keep the match exhaustive without a swallowing wildcard.
612        ClusterEvent::NamespacePlacementChanged {
613            name, placement, ..
614        } => ClusterEvent::NamespacePlacementChanged {
615            meta,
616            name,
617            placement,
618        },
619        // Like `NamespaceCreated`, emitted directly through the publisher (the
620        // throttled quota-snapshot task) which stamps the real meta, never queued
621        // in `tick()`; re-stamped faithfully to keep the match exhaustive.
622        ClusterEvent::NamespaceQuotaState {
623            namespace,
624            in_flight,
625            ceiling,
626            ..
627        } => ClusterEvent::NamespaceQuotaState {
628            meta,
629            namespace,
630            in_flight,
631            ceiling,
632        },
633        parked_event @ ClusterEvent::DispatchParked { .. } => {
634            with_meta_dispatch_parked(parked_event, meta)
635        }
636        deployment_event @ (ClusterEvent::WorkerDeploymentPut { .. }
637        | ClusterEvent::WorkerDeploymentDesiredStateChanged { .. }
638        | ClusterEvent::WorkerDeploymentDeleted { .. }) => {
639            with_meta_worker_deployment(deployment_event, meta)
640        }
641        ClusterEvent::PeerAdded { .. }
642        | ClusterEvent::PeerConnected { .. }
643        | ClusterEvent::PeerDisconnected { .. }
644        | ClusterEvent::ShardAdopted { .. }
645        | ClusterEvent::ShardAdoptionFailed { .. }
646        | ClusterEvent::ShardAdoptionSkipped { .. } => {
647            unreachable!("peer/shard variants are re-stamped by with_meta, never delegated here")
648        }
649    }
650}
651
652/// Meta re-stamp for the `DispatchParked` variant, the widest single re-stamp
653/// (split from [`with_meta_worker_lifecycle`] so neither match outgrows a
654/// readable dispatch). Like `NamespaceCreated`, it is emitted directly through
655/// the publisher (the queue-service wait path) which stamps the real meta,
656/// never queued in `tick()` with a placeholder; re-stamped faithfully to keep
657/// the exhaustive-match discipline without a swallowing wildcard.
658fn with_meta_dispatch_parked(
659    event: ClusterEvent,
660    meta: aion_core::ClusterEventMeta,
661) -> ClusterEvent {
662    match event {
663        ClusterEvent::DispatchParked {
664            namespace,
665            task_queue,
666            activity_type,
667            node,
668            reason,
669            policy,
670            workflow_id,
671            activity_id,
672            waited_ms,
673            workers_in_pool,
674            workers_serving_activity,
675            compatible_workers,
676            last_compatible_poller_age_ms,
677            ..
678        } => ClusterEvent::DispatchParked {
679            meta,
680            namespace,
681            task_queue,
682            activity_type,
683            node,
684            reason,
685            policy,
686            workflow_id,
687            activity_id,
688            waited_ms,
689            workers_in_pool,
690            workers_serving_activity,
691            compatible_workers,
692            last_compatible_poller_age_ms,
693        },
694        ClusterEvent::WorkerConnected { .. }
695        | ClusterEvent::WorkerDisconnected { .. }
696        | ClusterEvent::SupervisorStarted { .. }
697        | ClusterEvent::SupervisorStopped { .. }
698        | ClusterEvent::NamespaceCreated { .. }
699        | ClusterEvent::NamespacePlacementChanged { .. }
700        | ClusterEvent::NamespaceQuotaState { .. }
701        | ClusterEvent::WorkerDeploymentPut { .. }
702        | ClusterEvent::WorkerDeploymentDesiredStateChanged { .. }
703        | ClusterEvent::WorkerDeploymentDeleted { .. }
704        | ClusterEvent::PeerAdded { .. }
705        | ClusterEvent::PeerConnected { .. }
706        | ClusterEvent::PeerDisconnected { .. }
707        | ClusterEvent::ShardAdopted { .. }
708        | ClusterEvent::ShardAdoptionFailed { .. }
709        | ClusterEvent::ShardAdoptionSkipped { .. } => {
710            unreachable!("only DispatchParked is delegated here")
711        }
712    }
713}
714
715fn with_meta_worker_deployment(
716    event: ClusterEvent,
717    meta: aion_core::ClusterEventMeta,
718) -> ClusterEvent {
719    match event {
720        ClusterEvent::WorkerDeploymentPut {
721            name,
722            outcome,
723            desired_state,
724            binary_version,
725            binary_content_hash,
726            ..
727        } => ClusterEvent::WorkerDeploymentPut {
728            meta,
729            name,
730            outcome,
731            desired_state,
732            binary_version,
733            binary_content_hash,
734        },
735        ClusterEvent::WorkerDeploymentDesiredStateChanged {
736            name,
737            desired_state,
738            ..
739        } => ClusterEvent::WorkerDeploymentDesiredStateChanged {
740            meta,
741            name,
742            desired_state,
743        },
744        ClusterEvent::WorkerDeploymentDeleted { name, .. } => {
745            ClusterEvent::WorkerDeploymentDeleted { meta, name }
746        }
747        ClusterEvent::WorkerConnected { .. }
748        | ClusterEvent::WorkerDisconnected { .. }
749        | ClusterEvent::SupervisorStarted { .. }
750        | ClusterEvent::SupervisorStopped { .. }
751        | ClusterEvent::NamespaceCreated { .. }
752        | ClusterEvent::NamespacePlacementChanged { .. }
753        | ClusterEvent::NamespaceQuotaState { .. }
754        | ClusterEvent::DispatchParked { .. }
755        | ClusterEvent::PeerAdded { .. }
756        | ClusterEvent::PeerConnected { .. }
757        | ClusterEvent::PeerDisconnected { .. }
758        | ClusterEvent::ShardAdopted { .. }
759        | ClusterEvent::ShardAdoptionFailed { .. }
760        | ClusterEvent::ShardAdoptionSkipped { .. } => {
761            unreachable!("only worker-deployment variants are delegated here")
762        }
763    }
764}
765
766#[cfg(test)]
767mod tests {
768    use std::sync::Mutex;
769    use std::sync::atomic::{AtomicBool, Ordering};
770
771    use super::*;
772
773    /// A liveness fake whose verdict is flipped by the test. `connected` is the
774    /// verdict for ALL queried peers EXCEPT names explicitly registered as live
775    /// third-party owners via `set_live_owner`, which always report connected and
776    /// can be recorded as a shard's owner via `publish`.
777    struct FakeLiveness {
778        connected: AtomicBool,
779        /// shard -> recorded owner name (the SS-3 directory record).
780        owners: Mutex<std::collections::BTreeMap<usize, String>>,
781        /// peer names that always report connected (live third-party survivors).
782        live_owners: Mutex<std::collections::BTreeSet<String>>,
783    }
784
785    impl FakeLiveness {
786        fn new(connected: bool) -> Self {
787            Self {
788                connected: AtomicBool::new(connected),
789                owners: Mutex::new(std::collections::BTreeMap::new()),
790                live_owners: Mutex::new(std::collections::BTreeSet::new()),
791            }
792        }
793        fn set(&self, connected: bool) {
794            self.connected.store(connected, Ordering::SeqCst);
795        }
796        /// Record `owner` as `shard`'s directory owner and (if `live`) mark it as
797        /// a connected third-party survivor.
798        fn publish(&self, shard: usize, owner: &str, live: bool) {
799            self.owners
800                .lock()
801                .unwrap_or_else(std::sync::PoisonError::into_inner)
802                .insert(shard, owner.to_owned());
803            if live {
804                self.live_owners
805                    .lock()
806                    .unwrap_or_else(std::sync::PoisonError::into_inner)
807                    .insert(owner.to_owned());
808            }
809        }
810    }
811
812    impl PeerLiveness for FakeLiveness {
813        fn peer_connected(&self, peer_name: &str) -> bool {
814            if self
815                .live_owners
816                .lock()
817                .unwrap_or_else(std::sync::PoisonError::into_inner)
818                .contains(peer_name)
819            {
820                return true;
821            }
822            self.connected.load(Ordering::SeqCst)
823        }
824
825        fn read_shard_owner(&self, shard: usize) -> Option<String> {
826            self.owners
827                .lock()
828                .unwrap_or_else(std::sync::PoisonError::into_inner)
829                .get(&shard)
830                .cloned()
831        }
832    }
833
834    /// An adopter fake recording every adopt call, optionally failing the first.
835    struct FakeAdopter {
836        calls: Mutex<Vec<Vec<usize>>>,
837        fail_first: AtomicBool,
838    }
839
840    impl FakeAdopter {
841        fn new(fail_first: bool) -> Self {
842            Self {
843                calls: Mutex::new(Vec::new()),
844                fail_first: AtomicBool::new(fail_first),
845            }
846        }
847        fn calls(&self) -> Vec<Vec<usize>> {
848            self.calls
849                .lock()
850                .unwrap_or_else(std::sync::PoisonError::into_inner)
851                .clone()
852        }
853    }
854
855    #[async_trait::async_trait]
856    impl ShardAdopter for FakeAdopter {
857        async fn adopt_shards(&self, shards: &[usize]) -> Result<(), String> {
858            if self.fail_first.swap(false, Ordering::SeqCst) {
859                return Err("simulated election failure".to_owned());
860            }
861            self.calls
862                .lock()
863                .unwrap_or_else(std::sync::PoisonError::into_inner)
864                .push(shards.to_vec());
865            Ok(())
866        }
867    }
868
869    fn supervisor(
870        liveness: Arc<FakeLiveness>,
871        adopter: Arc<FakeAdopter>,
872        confirmations: u32,
873    ) -> ClusterSupervisor<FakeLiveness, FakeAdopter> {
874        ClusterSupervisor::new(
875            liveness,
876            adopter,
877            vec![WatchedPeer {
878                name: "node-1@127.0.0.1".to_owned(),
879                owned_shards: vec![1],
880            }],
881            SupervisorConfig {
882                poll_interval: Duration::from_millis(1),
883                confirmations,
884            },
885        )
886    }
887
888    /// #253 adoption sweep: after a (single-node no-op) adoption, the
889    /// outbox-settling adopter runs the terminal-workflow settlement over the
890    /// widened scope — a terminal workflow's stranded Claimed row is settled
891    /// to Cancelled by the adoption itself, before any dispatcher can re-arm
892    /// or redeliver it on the adopting node.
893    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
894    async fn outbox_settling_adopter_settles_terminal_rows_after_adoption()
895    -> Result<(), Box<dyn std::error::Error>> {
896        use aion::{EngineBuilder, RuntimeHandle, SignalRouter};
897        use aion_core::{Event, EventEnvelope};
898        use aion_store::{OutboxRow, OutboxStatus, OutboxStore, WritableEventStore, WriteToken};
899        use aion_store_haematite::HaematiteStore;
900
901        let db_path = std::env::temp_dir().join(format!(
902            "aion-adopter-settle-{}-{}.db",
903            std::process::id(),
904            uuid::Uuid::new_v4()
905        ));
906
907        // Seed the incident state: a terminal (Failed) workflow owning one
908        // stranded Claimed outbox row.
909        let seeder = Arc::new(
910            HaematiteStore::open_or_create(
911                db_path.clone(),
912                haematite::NodeCacheBudget::Unlimited,
913                // test-ruled patience: 250ms covers the measured 93-150ms fork window; not a default.
914                aion_store_haematite::LockAcquisitionRuling::from_millis(250, 5)?,
915            )
916            .await?,
917        );
918        let workflow_id = aion_core::WorkflowId::new_v4();
919        let envelope = |seq: u64| EventEnvelope {
920            seq,
921            recorded_at: chrono::Utc::now(),
922            workflow_id: workflow_id.clone(),
923        };
924        let events = vec![
925            Event::WorkflowStarted {
926                envelope: envelope(1),
927                workflow_type: String::from("dev_brief"),
928                input: aion_core::Payload::from_json(&serde_json::json!({}))?,
929                run_id: aion_core::RunId::new_v4(),
930                parent_run_id: None,
931                parent_workflow_id: None,
932                package_version: aion_core::PackageVersion::new("a".repeat(64)),
933            },
934            Event::WorkflowFailed {
935                envelope: envelope(2),
936                error: aion_core::WorkflowError {
937                    message: String::from("boom"),
938                    details: None,
939                },
940            },
941        ];
942        seeder
943            .append(WriteToken::recorder(), &workflow_id, &events, 0)
944            .await?;
945        let row = OutboxRow::pending(
946            workflow_id.clone(),
947            0,
948            String::from("norn_round"),
949            aion_core::Payload::from_json(&serde_json::json!({}))?,
950            chrono::Utc::now(),
951        );
952        let dispatch_key = row.dispatch_key.clone();
953        seeder
954            .append_outbox_batch(std::slice::from_ref(&row))
955            .await?;
956        assert_eq!(seeder.claim_outbox_rows(1).await?.len(), 1);
957
958        let event_store: Arc<dyn aion_store::EventStore> = Arc::clone(&seeder) as _;
959        let engine = Arc::new(
960            EngineBuilder::new()
961                .store_arc(event_store)
962                .in_memory_visibility()
963                .scheduler_threads(1)
964                .signal_router_factory(|runtime: Arc<RuntimeHandle>, handoff| {
965                    Arc::new(aion::signal::ConcreteSignalRouter::new(runtime, handoff))
966                        as Arc<dyn SignalRouter>
967                })
968                .build()
969                .await?,
970        );
971        let outbox_store: Arc<dyn OutboxStore> = Arc::clone(&seeder) as _;
972        let adopter =
973            OutboxSettlingAdopter::new(Arc::clone(&engine), Some(Arc::clone(&outbox_store)));
974
975        ShardAdopter::adopt_shards(&adopter, &[42])
976            .await
977            .map_err(|error| format!("adoption must succeed: {error}"))?;
978
979        let state = seeder
980            .outbox_row_state(&dispatch_key)
981            .await?
982            .ok_or("the stranded row must still exist")?;
983        assert_eq!(
984            state.status,
985            OutboxStatus::Cancelled,
986            "the adoption sweep must settle the terminal workflow's stranded row"
987        );
988        engine.shutdown()?;
989        Ok(())
990    }
991
992    #[tokio::test]
993    async fn does_not_adopt_while_peer_connected() {
994        let liveness = Arc::new(FakeLiveness::new(true));
995        let adopter = Arc::new(FakeAdopter::new(false));
996        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 2);
997        for _ in 0..5 {
998            assert!(sup.tick().await.is_empty());
999        }
1000        assert!(adopter.calls().is_empty(), "no adoption while peer is up");
1001    }
1002
1003    #[tokio::test]
1004    async fn debounce_requires_consecutive_down_before_adopting() {
1005        let liveness = Arc::new(FakeLiveness::new(true));
1006        let adopter = Arc::new(FakeAdopter::new(false));
1007        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 3);
1008
1009        liveness.set(false);
1010        assert!(sup.tick().await.is_empty(), "tick 1 down: below threshold");
1011        // A blip back up resets the counter.
1012        liveness.set(true);
1013        assert!(sup.tick().await.is_empty());
1014        liveness.set(false);
1015        assert!(
1016            sup.tick().await.is_empty(),
1017            "down again, counter reset to 1"
1018        );
1019        assert!(sup.tick().await.is_empty(), "2 consecutive: still below 3");
1020        let fired = sup.tick().await;
1021        assert_eq!(
1022            fired,
1023            vec!["node-1@127.0.0.1".to_owned()],
1024            "3rd consecutive triggers"
1025        );
1026        assert_eq!(adopter.calls(), vec![vec![1]]);
1027    }
1028
1029    #[tokio::test]
1030    async fn adopts_once_then_stays_quiet_while_down() {
1031        let liveness = Arc::new(FakeLiveness::new(false));
1032        let adopter = Arc::new(FakeAdopter::new(false));
1033        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1034        assert_eq!(sup.tick().await.len(), 1, "first down tick adopts");
1035        for _ in 0..5 {
1036            assert!(sup.tick().await.is_empty(), "no re-adopt while still down");
1037        }
1038        assert_eq!(adopter.calls(), vec![vec![1]], "adopted exactly once");
1039    }
1040
1041    #[tokio::test]
1042    async fn failed_adoption_is_retried_next_tick() {
1043        let liveness = Arc::new(FakeLiveness::new(false));
1044        let adopter = Arc::new(FakeAdopter::new(true)); // first adopt fails
1045        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1046        assert!(
1047            sup.tick().await.is_empty(),
1048            "first adopt fails, not recorded"
1049        );
1050        assert!(adopter.calls().is_empty());
1051        assert_eq!(sup.tick().await.len(), 1, "retry succeeds next tick");
1052        assert_eq!(adopter.calls(), vec![vec![1]]);
1053    }
1054
1055    #[tokio::test]
1056    async fn peer_with_no_shards_is_not_watched() {
1057        let liveness = Arc::new(FakeLiveness::new(false));
1058        let adopter = Arc::new(FakeAdopter::new(false));
1059        let mut sup = ClusterSupervisor::new(
1060            Arc::clone(&liveness),
1061            Arc::clone(&adopter),
1062            vec![WatchedPeer {
1063                name: "node-2@127.0.0.1".to_owned(),
1064                owned_shards: vec![],
1065            }],
1066            SupervisorConfig {
1067                poll_interval: Duration::from_millis(1),
1068                confirmations: 1,
1069            },
1070        );
1071        assert!(!sup.watches_any());
1072        assert!(sup.tick().await.is_empty());
1073        assert!(adopter.calls().is_empty());
1074    }
1075
1076    /// PRE-CHECK: a downed peer whose shard is ALREADY published to a DIFFERENT
1077    /// LIVE owner is NOT adopted — another survivor holds it. The supervisor marks
1078    /// the peer handled (no retry-loop) and never calls the adopter.
1079    #[tokio::test]
1080    async fn shard_already_published_to_live_owner_is_not_adopted() {
1081        let liveness = Arc::new(FakeLiveness::new(false));
1082        // Shard 1 (the watched peer's shard) is recorded as owned by a live third
1083        // party, node-9 — it adopted the shard already.
1084        liveness.publish(1, "node-9@127.0.0.1", true);
1085        let adopter = Arc::new(FakeAdopter::new(false));
1086        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1087
1088        // The peer is down past the threshold, but its shard is handled elsewhere.
1089        assert!(
1090            sup.tick().await.is_empty(),
1091            "no adoption fires for a shard a live owner already holds"
1092        );
1093        assert!(
1094            adopter.calls().is_empty(),
1095            "the adopter is never invoked for an already-handled shard"
1096        );
1097        // Subsequent ticks stay quiet: marked handled, no retry-loop.
1098        for _ in 0..3 {
1099            assert!(sup.tick().await.is_empty());
1100        }
1101        assert!(adopter.calls().is_empty());
1102    }
1103
1104    /// A directory record naming a peer that is itself DOWN is NOT "handled
1105    /// elsewhere": the recorded owner has since died, so the shard remains
1106    /// adoptable and the supervisor adopts it.
1107    #[tokio::test]
1108    async fn shard_published_to_a_down_owner_is_still_adopted() {
1109        let liveness = Arc::new(FakeLiveness::new(false));
1110        // Shard 1 recorded as owned by node-9, but node-9 is NOT live (not
1111        // registered as a live owner) — `connected=false` applies to it.
1112        liveness.publish(1, "node-9@127.0.0.1", false);
1113        let adopter = Arc::new(FakeAdopter::new(false));
1114        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1115
1116        assert_eq!(
1117            sup.tick().await.len(),
1118            1,
1119            "a shard whose recorded owner is itself down is adoptable"
1120        );
1121        assert_eq!(adopter.calls(), vec![vec![1]]);
1122    }
1123
1124    /// WS3 EMIT: with a publisher attached, a down tick emits `PeerDisconnected`
1125    /// (confirmed flipping at the threshold) and the adoption tick emits
1126    /// `ShardAdopted` carrying this node's name. The recovery EMIT then fires
1127    /// `PeerConnected` — proving the capture-before-reset: a freshly-reconnected
1128    /// peer that was previously down/adopted yields exactly one recovery event.
1129    #[tokio::test]
1130    async fn tick_emits_topology_deltas_through_the_publisher()
1131    -> Result<(), Box<dyn std::error::Error>> {
1132        use std::num::NonZeroUsize;
1133
1134        use aion_core::ClusterEvent;
1135        use futures::StreamExt;
1136
1137        use crate::cluster_publisher::ClusterEventPublisher;
1138
1139        let capacity = NonZeroUsize::new(64).ok_or("non-zero")?;
1140        let publisher = Arc::new(ClusterEventPublisher::new(capacity));
1141        let mut subscription = publisher.subscribe(0);
1142
1143        let liveness = Arc::new(FakeLiveness::new(true));
1144        let adopter = Arc::new(FakeAdopter::new(false));
1145        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 2)
1146            .with_publisher(Arc::clone(&publisher), "node-self@127.0.0.1");
1147
1148        // Tick 1 down: PeerDisconnected{confirmed=false} (below threshold 2).
1149        liveness.set(false);
1150        drop(sup.tick().await);
1151        // Tick 2 down: PeerDisconnected{confirmed=true} then ShardAdopted.
1152        let fired = sup.tick().await;
1153        assert_eq!(fired, vec!["node-1@127.0.0.1".to_owned()]);
1154
1155        // Drain the three emitted deltas in order.
1156        let first = next_event(&mut subscription).await?;
1157        assert!(
1158            matches!(
1159                &first,
1160                ClusterEvent::PeerDisconnected {
1161                    confirmed: false,
1162                    consecutive_down: 1,
1163                    ..
1164                }
1165            ),
1166            "first delta must be an unconfirmed down: {first:?}"
1167        );
1168        let second = next_event(&mut subscription).await?;
1169        assert!(
1170            matches!(
1171                &second,
1172                ClusterEvent::PeerDisconnected {
1173                    confirmed: true,
1174                    consecutive_down: 2,
1175                    ..
1176                }
1177            ),
1178            "second delta must be the confirmed down: {second:?}"
1179        );
1180        let third = next_event(&mut subscription).await?;
1181        let ClusterEvent::ShardAdopted {
1182            shards,
1183            adopted_by,
1184            from_peer,
1185            ..
1186        } = &third
1187        else {
1188            return Err(format!("third delta must be ShardAdopted: {third:?}").into());
1189        };
1190        assert_eq!(shards, &vec![1]);
1191        assert_eq!(adopted_by, "node-self@127.0.0.1");
1192        assert_eq!(from_peer, "node-1@127.0.0.1");
1193
1194        // RECOVERY: peer comes back up. The capture-before-reset must fire exactly
1195        // one PeerConnected for the now-recovered (previously adopted) peer.
1196        liveness.set(true);
1197        drop(sup.tick().await);
1198        let recovery = next_event(&mut subscription).await?;
1199        assert!(
1200            matches!(&recovery, ClusterEvent::PeerConnected { .. }),
1201            "recovery delta must be PeerConnected: {recovery:?}"
1202        );
1203
1204        // A second connected tick (already reset) must NOT re-emit recovery: the
1205        // next delta is whatever a subsequent down produces, never a duplicate
1206        // PeerConnected. Quiet tick yields nothing.
1207        let quiet = sup.tick().await;
1208        assert!(quiet.is_empty());
1209        // No further event is buffered (no spurious recovery re-emit).
1210        assert!(
1211            tokio::time::timeout(std::time::Duration::from_millis(50), subscription.next())
1212                .await
1213                .is_err(),
1214            "a steady connected peer must not re-emit PeerConnected every tick"
1215        );
1216        Ok(())
1217    }
1218
1219    async fn next_event(
1220        subscription: &mut futures::stream::BoxStream<
1221            'static,
1222            Result<aion_core::ClusterEvent, crate::cluster_publisher::ClusterStreamLagged>,
1223        >,
1224    ) -> Result<aion_core::ClusterEvent, Box<dyn std::error::Error>> {
1225        use futures::StreamExt;
1226        tokio::time::timeout(std::time::Duration::from_secs(1), subscription.next())
1227            .await?
1228            .ok_or("cluster subscription ended")?
1229            .map_err(|lag| format!("unexpected lag: {lag:?}").into())
1230    }
1231
1232    /// A record naming the DEAD peer itself (the steady-state declared owner) is
1233    /// stale and does NOT block adoption.
1234    #[tokio::test]
1235    async fn shard_published_to_the_dead_peer_itself_is_adopted() {
1236        let liveness = Arc::new(FakeLiveness::new(false));
1237        // The directory still names the (now dead) declared owner of shard 1.
1238        liveness.publish(1, "node-1@127.0.0.1", false);
1239        let adopter = Arc::new(FakeAdopter::new(false));
1240        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1241
1242        assert_eq!(
1243            sup.tick().await.len(),
1244            1,
1245            "a record naming the dead peer itself is stale and still adoptable"
1246        );
1247        assert_eq!(adopter.calls(), vec![vec![1]]);
1248    }
1249}