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