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            deployment,
571            deployment_association,
572            ..
573        } => ClusterEvent::WorkerConnected {
574            meta,
575            worker_id,
576            namespaces,
577            task_queue,
578            transport,
579            node,
580            deployment,
581            deployment_association,
582        },
583        ClusterEvent::WorkerDisconnected {
584            worker_id,
585            namespaces,
586            reason,
587            ..
588        } => ClusterEvent::WorkerDisconnected {
589            meta,
590            worker_id,
591            namespaces,
592            reason,
593        },
594        ClusterEvent::SupervisorStarted { node, .. } => {
595            ClusterEvent::SupervisorStarted { meta, node }
596        }
597        ClusterEvent::SupervisorStopped { node, .. } => {
598            ClusterEvent::SupervisorStopped { meta, node }
599        }
600        ClusterEvent::NamespaceCreated {
601            name,
602            created_at,
603            origin,
604            ..
605        } => ClusterEvent::NamespaceCreated {
606            meta,
607            name,
608            created_at,
609            origin,
610        },
611        // Like `NamespaceCreated`, emitted directly through the publisher (which
612        // stamps the real meta), never queued in `tick()`; re-stamped faithfully
613        // to keep the match exhaustive without a swallowing wildcard.
614        ClusterEvent::NamespacePlacementChanged {
615            name, placement, ..
616        } => ClusterEvent::NamespacePlacementChanged {
617            meta,
618            name,
619            placement,
620        },
621        // Like `NamespaceCreated`, emitted directly through the publisher (the
622        // throttled quota-snapshot task) which stamps the real meta, never queued
623        // in `tick()`; re-stamped faithfully to keep the match exhaustive.
624        ClusterEvent::NamespaceQuotaState {
625            namespace,
626            in_flight,
627            ceiling,
628            ..
629        } => ClusterEvent::NamespaceQuotaState {
630            meta,
631            namespace,
632            in_flight,
633            ceiling,
634        },
635        deployment_event @ (ClusterEvent::WorkerDeploymentPut { .. }
636        | ClusterEvent::WorkerDeploymentDesiredStateChanged { .. }
637        | ClusterEvent::WorkerDeploymentDeleted { .. }) => {
638            with_meta_worker_deployment(deployment_event, meta)
639        }
640        ClusterEvent::PeerAdded { .. }
641        | ClusterEvent::PeerConnected { .. }
642        | ClusterEvent::PeerDisconnected { .. }
643        | ClusterEvent::ShardAdopted { .. }
644        | ClusterEvent::ShardAdoptionFailed { .. }
645        | ClusterEvent::ShardAdoptionSkipped { .. } => {
646            unreachable!("peer/shard variants are re-stamped by with_meta, never delegated here")
647        }
648    }
649}
650
651fn with_meta_worker_deployment(
652    event: ClusterEvent,
653    meta: aion_core::ClusterEventMeta,
654) -> ClusterEvent {
655    match event {
656        ClusterEvent::WorkerDeploymentPut {
657            name,
658            outcome,
659            desired_state,
660            binary_version,
661            binary_content_hash,
662            ..
663        } => ClusterEvent::WorkerDeploymentPut {
664            meta,
665            name,
666            outcome,
667            desired_state,
668            binary_version,
669            binary_content_hash,
670        },
671        ClusterEvent::WorkerDeploymentDesiredStateChanged {
672            name,
673            desired_state,
674            ..
675        } => ClusterEvent::WorkerDeploymentDesiredStateChanged {
676            meta,
677            name,
678            desired_state,
679        },
680        ClusterEvent::WorkerDeploymentDeleted { name, .. } => {
681            ClusterEvent::WorkerDeploymentDeleted { meta, name }
682        }
683        ClusterEvent::WorkerConnected { .. }
684        | ClusterEvent::WorkerDisconnected { .. }
685        | ClusterEvent::SupervisorStarted { .. }
686        | ClusterEvent::SupervisorStopped { .. }
687        | ClusterEvent::NamespaceCreated { .. }
688        | ClusterEvent::NamespacePlacementChanged { .. }
689        | ClusterEvent::NamespaceQuotaState { .. }
690        | ClusterEvent::PeerAdded { .. }
691        | ClusterEvent::PeerConnected { .. }
692        | ClusterEvent::PeerDisconnected { .. }
693        | ClusterEvent::ShardAdopted { .. }
694        | ClusterEvent::ShardAdoptionFailed { .. }
695        | ClusterEvent::ShardAdoptionSkipped { .. } => {
696            unreachable!("only worker-deployment variants are delegated here")
697        }
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use std::sync::Mutex;
704    use std::sync::atomic::{AtomicBool, Ordering};
705
706    use super::*;
707
708    /// A liveness fake whose verdict is flipped by the test. `connected` is the
709    /// verdict for ALL queried peers EXCEPT names explicitly registered as live
710    /// third-party owners via `set_live_owner`, which always report connected and
711    /// can be recorded as a shard's owner via `publish`.
712    struct FakeLiveness {
713        connected: AtomicBool,
714        /// shard -> recorded owner name (the SS-3 directory record).
715        owners: Mutex<std::collections::BTreeMap<usize, String>>,
716        /// peer names that always report connected (live third-party survivors).
717        live_owners: Mutex<std::collections::BTreeSet<String>>,
718    }
719
720    impl FakeLiveness {
721        fn new(connected: bool) -> Self {
722            Self {
723                connected: AtomicBool::new(connected),
724                owners: Mutex::new(std::collections::BTreeMap::new()),
725                live_owners: Mutex::new(std::collections::BTreeSet::new()),
726            }
727        }
728        fn set(&self, connected: bool) {
729            self.connected.store(connected, Ordering::SeqCst);
730        }
731        /// Record `owner` as `shard`'s directory owner and (if `live`) mark it as
732        /// a connected third-party survivor.
733        fn publish(&self, shard: usize, owner: &str, live: bool) {
734            self.owners
735                .lock()
736                .unwrap_or_else(std::sync::PoisonError::into_inner)
737                .insert(shard, owner.to_owned());
738            if live {
739                self.live_owners
740                    .lock()
741                    .unwrap_or_else(std::sync::PoisonError::into_inner)
742                    .insert(owner.to_owned());
743            }
744        }
745    }
746
747    impl PeerLiveness for FakeLiveness {
748        fn peer_connected(&self, peer_name: &str) -> bool {
749            if self
750                .live_owners
751                .lock()
752                .unwrap_or_else(std::sync::PoisonError::into_inner)
753                .contains(peer_name)
754            {
755                return true;
756            }
757            self.connected.load(Ordering::SeqCst)
758        }
759
760        fn read_shard_owner(&self, shard: usize) -> Option<String> {
761            self.owners
762                .lock()
763                .unwrap_or_else(std::sync::PoisonError::into_inner)
764                .get(&shard)
765                .cloned()
766        }
767    }
768
769    /// An adopter fake recording every adopt call, optionally failing the first.
770    struct FakeAdopter {
771        calls: Mutex<Vec<Vec<usize>>>,
772        fail_first: AtomicBool,
773    }
774
775    impl FakeAdopter {
776        fn new(fail_first: bool) -> Self {
777            Self {
778                calls: Mutex::new(Vec::new()),
779                fail_first: AtomicBool::new(fail_first),
780            }
781        }
782        fn calls(&self) -> Vec<Vec<usize>> {
783            self.calls
784                .lock()
785                .unwrap_or_else(std::sync::PoisonError::into_inner)
786                .clone()
787        }
788    }
789
790    #[async_trait::async_trait]
791    impl ShardAdopter for FakeAdopter {
792        async fn adopt_shards(&self, shards: &[usize]) -> Result<(), String> {
793            if self.fail_first.swap(false, Ordering::SeqCst) {
794                return Err("simulated election failure".to_owned());
795            }
796            self.calls
797                .lock()
798                .unwrap_or_else(std::sync::PoisonError::into_inner)
799                .push(shards.to_vec());
800            Ok(())
801        }
802    }
803
804    fn supervisor(
805        liveness: Arc<FakeLiveness>,
806        adopter: Arc<FakeAdopter>,
807        confirmations: u32,
808    ) -> ClusterSupervisor<FakeLiveness, FakeAdopter> {
809        ClusterSupervisor::new(
810            liveness,
811            adopter,
812            vec![WatchedPeer {
813                name: "node-1@127.0.0.1".to_owned(),
814                owned_shards: vec![1],
815            }],
816            SupervisorConfig {
817                poll_interval: Duration::from_millis(1),
818                confirmations,
819            },
820        )
821    }
822
823    /// #253 adoption sweep: after a (single-node no-op) adoption, the
824    /// outbox-settling adopter runs the terminal-workflow settlement over the
825    /// widened scope — a terminal workflow's stranded Claimed row is settled
826    /// to Cancelled by the adoption itself, before any dispatcher can re-arm
827    /// or redeliver it on the adopting node.
828    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
829    async fn outbox_settling_adopter_settles_terminal_rows_after_adoption()
830    -> Result<(), Box<dyn std::error::Error>> {
831        use aion::{EngineBuilder, RuntimeHandle, SignalRouter};
832        use aion_core::{Event, EventEnvelope};
833        use aion_store::{OutboxRow, OutboxStatus, OutboxStore, WritableEventStore, WriteToken};
834        use aion_store_libsql::LibSqlStore;
835
836        let db_path = std::env::temp_dir().join(format!(
837            "aion-adopter-settle-{}-{}.db",
838            std::process::id(),
839            uuid::Uuid::new_v4()
840        ));
841
842        // Seed the incident state: a terminal (Failed) workflow owning one
843        // stranded Claimed outbox row.
844        let seeder = LibSqlStore::open(db_path.clone()).await?;
845        let workflow_id = aion_core::WorkflowId::new_v4();
846        let envelope = |seq: u64| EventEnvelope {
847            seq,
848            recorded_at: chrono::Utc::now(),
849            workflow_id: workflow_id.clone(),
850        };
851        let events = vec![
852            Event::WorkflowStarted {
853                envelope: envelope(1),
854                workflow_type: String::from("dev_brief"),
855                input: aion_core::Payload::from_json(&serde_json::json!({}))?,
856                run_id: aion_core::RunId::new_v4(),
857                parent_run_id: None,
858                package_version: aion_core::PackageVersion::new("a".repeat(64)),
859            },
860            Event::WorkflowFailed {
861                envelope: envelope(2),
862                error: aion_core::WorkflowError {
863                    message: String::from("boom"),
864                    details: None,
865                },
866            },
867        ];
868        seeder
869            .append(WriteToken::recorder(), &workflow_id, &events, 0)
870            .await?;
871        let row = OutboxRow::pending(
872            workflow_id.clone(),
873            0,
874            String::from("norn_round"),
875            aion_core::Payload::from_json(&serde_json::json!({}))?,
876            chrono::Utc::now(),
877        );
878        let dispatch_key = row.dispatch_key.clone();
879        seeder
880            .append_outbox_batch(std::slice::from_ref(&row))
881            .await?;
882        assert_eq!(seeder.claim_outbox_rows(1).await?.len(), 1);
883
884        let engine = Arc::new(
885            EngineBuilder::new()
886                .store_arc(Arc::new(LibSqlStore::open(db_path.clone()).await?))
887                .in_memory_visibility()
888                .scheduler_threads(1)
889                .signal_router_factory(|runtime: Arc<RuntimeHandle>, handoff| {
890                    Arc::new(aion::signal::ConcreteSignalRouter::new(runtime, handoff))
891                        as Arc<dyn SignalRouter>
892                })
893                .build()
894                .await?,
895        );
896        let outbox_store: Arc<dyn OutboxStore> =
897            Arc::new(LibSqlStore::open(db_path.clone()).await?);
898        let adopter =
899            OutboxSettlingAdopter::new(Arc::clone(&engine), Some(Arc::clone(&outbox_store)));
900
901        ShardAdopter::adopt_shards(&adopter, &[42])
902            .await
903            .map_err(|error| format!("adoption must succeed: {error}"))?;
904
905        let state = seeder
906            .outbox_row_state(&dispatch_key)
907            .await?
908            .ok_or("the stranded row must still exist")?;
909        assert_eq!(
910            state.status,
911            OutboxStatus::Cancelled,
912            "the adoption sweep must settle the terminal workflow's stranded row"
913        );
914        engine.shutdown()?;
915        Ok(())
916    }
917
918    #[tokio::test]
919    async fn does_not_adopt_while_peer_connected() {
920        let liveness = Arc::new(FakeLiveness::new(true));
921        let adopter = Arc::new(FakeAdopter::new(false));
922        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 2);
923        for _ in 0..5 {
924            assert!(sup.tick().await.is_empty());
925        }
926        assert!(adopter.calls().is_empty(), "no adoption while peer is up");
927    }
928
929    #[tokio::test]
930    async fn debounce_requires_consecutive_down_before_adopting() {
931        let liveness = Arc::new(FakeLiveness::new(true));
932        let adopter = Arc::new(FakeAdopter::new(false));
933        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 3);
934
935        liveness.set(false);
936        assert!(sup.tick().await.is_empty(), "tick 1 down: below threshold");
937        // A blip back up resets the counter.
938        liveness.set(true);
939        assert!(sup.tick().await.is_empty());
940        liveness.set(false);
941        assert!(
942            sup.tick().await.is_empty(),
943            "down again, counter reset to 1"
944        );
945        assert!(sup.tick().await.is_empty(), "2 consecutive: still below 3");
946        let fired = sup.tick().await;
947        assert_eq!(
948            fired,
949            vec!["node-1@127.0.0.1".to_owned()],
950            "3rd consecutive triggers"
951        );
952        assert_eq!(adopter.calls(), vec![vec![1]]);
953    }
954
955    #[tokio::test]
956    async fn adopts_once_then_stays_quiet_while_down() {
957        let liveness = Arc::new(FakeLiveness::new(false));
958        let adopter = Arc::new(FakeAdopter::new(false));
959        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
960        assert_eq!(sup.tick().await.len(), 1, "first down tick adopts");
961        for _ in 0..5 {
962            assert!(sup.tick().await.is_empty(), "no re-adopt while still down");
963        }
964        assert_eq!(adopter.calls(), vec![vec![1]], "adopted exactly once");
965    }
966
967    #[tokio::test]
968    async fn failed_adoption_is_retried_next_tick() {
969        let liveness = Arc::new(FakeLiveness::new(false));
970        let adopter = Arc::new(FakeAdopter::new(true)); // first adopt fails
971        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
972        assert!(
973            sup.tick().await.is_empty(),
974            "first adopt fails, not recorded"
975        );
976        assert!(adopter.calls().is_empty());
977        assert_eq!(sup.tick().await.len(), 1, "retry succeeds next tick");
978        assert_eq!(adopter.calls(), vec![vec![1]]);
979    }
980
981    #[tokio::test]
982    async fn peer_with_no_shards_is_not_watched() {
983        let liveness = Arc::new(FakeLiveness::new(false));
984        let adopter = Arc::new(FakeAdopter::new(false));
985        let mut sup = ClusterSupervisor::new(
986            Arc::clone(&liveness),
987            Arc::clone(&adopter),
988            vec![WatchedPeer {
989                name: "node-2@127.0.0.1".to_owned(),
990                owned_shards: vec![],
991            }],
992            SupervisorConfig {
993                poll_interval: Duration::from_millis(1),
994                confirmations: 1,
995            },
996        );
997        assert!(!sup.watches_any());
998        assert!(sup.tick().await.is_empty());
999        assert!(adopter.calls().is_empty());
1000    }
1001
1002    /// PRE-CHECK: a downed peer whose shard is ALREADY published to a DIFFERENT
1003    /// LIVE owner is NOT adopted — another survivor holds it. The supervisor marks
1004    /// the peer handled (no retry-loop) and never calls the adopter.
1005    #[tokio::test]
1006    async fn shard_already_published_to_live_owner_is_not_adopted() {
1007        let liveness = Arc::new(FakeLiveness::new(false));
1008        // Shard 1 (the watched peer's shard) is recorded as owned by a live third
1009        // party, node-9 — it adopted the shard already.
1010        liveness.publish(1, "node-9@127.0.0.1", true);
1011        let adopter = Arc::new(FakeAdopter::new(false));
1012        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1013
1014        // The peer is down past the threshold, but its shard is handled elsewhere.
1015        assert!(
1016            sup.tick().await.is_empty(),
1017            "no adoption fires for a shard a live owner already holds"
1018        );
1019        assert!(
1020            adopter.calls().is_empty(),
1021            "the adopter is never invoked for an already-handled shard"
1022        );
1023        // Subsequent ticks stay quiet: marked handled, no retry-loop.
1024        for _ in 0..3 {
1025            assert!(sup.tick().await.is_empty());
1026        }
1027        assert!(adopter.calls().is_empty());
1028    }
1029
1030    /// A directory record naming a peer that is itself DOWN is NOT "handled
1031    /// elsewhere": the recorded owner has since died, so the shard remains
1032    /// adoptable and the supervisor adopts it.
1033    #[tokio::test]
1034    async fn shard_published_to_a_down_owner_is_still_adopted() {
1035        let liveness = Arc::new(FakeLiveness::new(false));
1036        // Shard 1 recorded as owned by node-9, but node-9 is NOT live (not
1037        // registered as a live owner) — `connected=false` applies to it.
1038        liveness.publish(1, "node-9@127.0.0.1", false);
1039        let adopter = Arc::new(FakeAdopter::new(false));
1040        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1041
1042        assert_eq!(
1043            sup.tick().await.len(),
1044            1,
1045            "a shard whose recorded owner is itself down is adoptable"
1046        );
1047        assert_eq!(adopter.calls(), vec![vec![1]]);
1048    }
1049
1050    /// WS3 EMIT: with a publisher attached, a down tick emits `PeerDisconnected`
1051    /// (confirmed flipping at the threshold) and the adoption tick emits
1052    /// `ShardAdopted` carrying this node's name. The recovery EMIT then fires
1053    /// `PeerConnected` — proving the capture-before-reset: a freshly-reconnected
1054    /// peer that was previously down/adopted yields exactly one recovery event.
1055    #[tokio::test]
1056    async fn tick_emits_topology_deltas_through_the_publisher()
1057    -> Result<(), Box<dyn std::error::Error>> {
1058        use std::num::NonZeroUsize;
1059
1060        use aion_core::ClusterEvent;
1061        use futures::StreamExt;
1062
1063        use crate::cluster_publisher::ClusterEventPublisher;
1064
1065        let capacity = NonZeroUsize::new(64).ok_or("non-zero")?;
1066        let publisher = Arc::new(ClusterEventPublisher::new(capacity));
1067        let mut subscription = publisher.subscribe(0);
1068
1069        let liveness = Arc::new(FakeLiveness::new(true));
1070        let adopter = Arc::new(FakeAdopter::new(false));
1071        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 2)
1072            .with_publisher(Arc::clone(&publisher), "node-self@127.0.0.1");
1073
1074        // Tick 1 down: PeerDisconnected{confirmed=false} (below threshold 2).
1075        liveness.set(false);
1076        drop(sup.tick().await);
1077        // Tick 2 down: PeerDisconnected{confirmed=true} then ShardAdopted.
1078        let fired = sup.tick().await;
1079        assert_eq!(fired, vec!["node-1@127.0.0.1".to_owned()]);
1080
1081        // Drain the three emitted deltas in order.
1082        let first = next_event(&mut subscription).await?;
1083        assert!(
1084            matches!(
1085                &first,
1086                ClusterEvent::PeerDisconnected {
1087                    confirmed: false,
1088                    consecutive_down: 1,
1089                    ..
1090                }
1091            ),
1092            "first delta must be an unconfirmed down: {first:?}"
1093        );
1094        let second = next_event(&mut subscription).await?;
1095        assert!(
1096            matches!(
1097                &second,
1098                ClusterEvent::PeerDisconnected {
1099                    confirmed: true,
1100                    consecutive_down: 2,
1101                    ..
1102                }
1103            ),
1104            "second delta must be the confirmed down: {second:?}"
1105        );
1106        let third = next_event(&mut subscription).await?;
1107        let ClusterEvent::ShardAdopted {
1108            shards,
1109            adopted_by,
1110            from_peer,
1111            ..
1112        } = &third
1113        else {
1114            return Err(format!("third delta must be ShardAdopted: {third:?}").into());
1115        };
1116        assert_eq!(shards, &vec![1]);
1117        assert_eq!(adopted_by, "node-self@127.0.0.1");
1118        assert_eq!(from_peer, "node-1@127.0.0.1");
1119
1120        // RECOVERY: peer comes back up. The capture-before-reset must fire exactly
1121        // one PeerConnected for the now-recovered (previously adopted) peer.
1122        liveness.set(true);
1123        drop(sup.tick().await);
1124        let recovery = next_event(&mut subscription).await?;
1125        assert!(
1126            matches!(&recovery, ClusterEvent::PeerConnected { .. }),
1127            "recovery delta must be PeerConnected: {recovery:?}"
1128        );
1129
1130        // A second connected tick (already reset) must NOT re-emit recovery: the
1131        // next delta is whatever a subsequent down produces, never a duplicate
1132        // PeerConnected. Quiet tick yields nothing.
1133        let quiet = sup.tick().await;
1134        assert!(quiet.is_empty());
1135        // No further event is buffered (no spurious recovery re-emit).
1136        assert!(
1137            tokio::time::timeout(std::time::Duration::from_millis(50), subscription.next())
1138                .await
1139                .is_err(),
1140            "a steady connected peer must not re-emit PeerConnected every tick"
1141        );
1142        Ok(())
1143    }
1144
1145    async fn next_event(
1146        subscription: &mut futures::stream::BoxStream<
1147            'static,
1148            Result<aion_core::ClusterEvent, crate::cluster_publisher::ClusterStreamLagged>,
1149        >,
1150    ) -> Result<aion_core::ClusterEvent, Box<dyn std::error::Error>> {
1151        use futures::StreamExt;
1152        tokio::time::timeout(std::time::Duration::from_secs(1), subscription.next())
1153            .await?
1154            .ok_or("cluster subscription ended")?
1155            .map_err(|lag| format!("unexpected lag: {lag:?}").into())
1156    }
1157
1158    /// A record naming the DEAD peer itself (the steady-state declared owner) is
1159    /// stale and does NOT block adoption.
1160    #[tokio::test]
1161    async fn shard_published_to_the_dead_peer_itself_is_adopted() {
1162        let liveness = Arc::new(FakeLiveness::new(false));
1163        // The directory still names the (now dead) declared owner of shard 1.
1164        liveness.publish(1, "node-1@127.0.0.1", false);
1165        let adopter = Arc::new(FakeAdopter::new(false));
1166        let mut sup = supervisor(Arc::clone(&liveness), Arc::clone(&adopter), 1);
1167
1168        assert_eq!(
1169            sup.tick().await.len(),
1170            1,
1171            "a record naming the dead peer itself is stale and still adoptable"
1172        );
1173        assert_eq!(adopter.calls(), vec![vec![1]]);
1174    }
1175}