Skip to main content

nodedb_cluster/raft_loop/
loop_core.rs

1// SPDX-License-Identifier: BUSL-1.1
2
3//! `RaftLoop` struct, constructors, top-level run loop, and thin wrappers
4//! over `MultiRaft` proposal APIs. The tick body lives in
5//! [`super::tick`]; the inbound-RPC handler lives in
6//! [`super::handle_rpc`]; the async join orchestration lives in
7//! [`super::join`].
8
9use std::pin::Pin;
10use std::sync::{Arc, Mutex, RwLock};
11use std::time::{Duration, Instant};
12
13use tracing::debug;
14
15use nodedb_raft::message::LogEntry;
16
17use crate::applied_watcher::GroupAppliedWatchers;
18use crate::catalog::ClusterCatalog;
19use crate::error::Result;
20use crate::forward::{NoopPlanExecutor, PlanExecutor};
21use crate::loop_metrics::LoopMetrics;
22use crate::metadata_group::applier::{MetadataApplier, NoopMetadataApplier};
23use crate::multi_raft::MultiRaft;
24use crate::topology::ClusterTopology;
25use crate::transport::NexarTransport;
26
27use super::hooks::{
28    AssignRemoteSurrogate, CalvinSubmit, CalvinSubmitInbox, ReleaseReservation, ReserveRead,
29    ShuffleAggregator, ShuffleConsumer, ShuffleProducer, ShuffleReceiver, SnapshotApplier,
30    SnapshotBuilder, SnapshotQuarantineHook,
31};
32
33/// Default tick interval (10ms — fast enough for sub-second elections).
34///
35/// Matches `ClusterTransportTuning::raft_tick_interval_ms` default.
36pub(super) const DEFAULT_TICK_INTERVAL: Duration = Duration::from_millis(10);
37
38/// Callback for applying committed Raft log entries to the state machine.
39///
40/// Called synchronously during the tick loop. Implementations should be fast
41/// (enqueue to SPSC, not perform I/O directly).
42pub trait CommitApplier: Send + Sync + 'static {
43    /// Apply committed entries for a Raft group.
44    ///
45    /// Returns the index of the last successfully applied entry.
46    fn apply_committed(&self, group_id: u64, entries: &[LogEntry]) -> u64;
47}
48
49/// Type-erased async handler for incoming `VShardEnvelope` messages.
50///
51/// Receives raw envelope bytes, returns response bytes. Set by the main binary
52/// to dispatch to the appropriate engine handler (Event Plane, timeseries, etc.).
53pub type VShardEnvelopeHandler = Arc<
54    dyn Fn(Vec<u8>) -> Pin<Box<dyn std::future::Future<Output = Result<Vec<u8>>> + Send>>
55        + Send
56        + Sync,
57>;
58
59/// Raft event loop coordinator.
60///
61/// Owns the MultiRaft state (behind `Arc<Mutex>`) and drives it via periodic
62/// ticks. Implements [`crate::transport::RaftRpcHandler`] (in
63/// [`super::handle_rpc`]) so it can be passed directly to
64/// [`NexarTransport::serve`] for incoming RPC dispatch.
65///
66/// The `F: RequestForwarder` generic parameter was removed in C-δ.6 when the
67/// SQL-string forwarding path was retired. Cross-node SQL routing now goes
68/// through `gateway.execute / ExecuteRequest` (C-β path).
69pub struct RaftLoop<A: CommitApplier, P: PlanExecutor = NoopPlanExecutor> {
70    pub(super) node_id: u64,
71    pub(super) multi_raft: Arc<Mutex<MultiRaft>>,
72    pub(super) transport: Arc<NexarTransport>,
73    pub(super) topology: Arc<RwLock<ClusterTopology>>,
74    pub(super) applier: A,
75    /// Applies committed entries from the metadata Raft group (group 0).
76    pub(super) metadata_applier: Arc<dyn MetadataApplier>,
77    /// Executes incoming `ExecuteRequest` RPCs without SQL re-planning.
78    pub(super) plan_executor: Arc<P>,
79    pub(super) tick_interval: Duration,
80    /// Optional handler for incoming VShardEnvelope messages.
81    /// Set when the Event Plane or other subsystems need cross-node messaging.
82    pub(super) vshard_handler: Option<VShardEnvelopeHandler>,
83    /// Optional catalog handle for persisting topology/routing updates
84    /// from the join flow. When `None`, persistence is skipped — useful
85    /// for unit tests that don't care about durability.
86    pub(super) catalog: Option<Arc<ClusterCatalog>>,
87    /// Cooperative shutdown signal observed by every detached
88    /// `tokio::spawn` task in [`super::tick`]. `run()` flips it on
89    /// its own shutdown, and [`Self::begin_shutdown`] provides a
90    /// direct entry point for test harnesses that abort the run /
91    /// serve handles and need the spawned tasks to drop their
92    /// `Arc<Mutex<MultiRaft>>` clones immediately so the per-group
93    /// redb log files can release their in-process locks.
94    ///
95    /// Using `watch::Sender` here rather than a raw `AtomicBool` +
96    /// `Notify` pair gives us two properties at once: the latest
97    /// value is visible to every newly-subscribed receiver (no
98    /// missed-notification race when a new detached task is
99    /// spawned just after `begin_shutdown`), and awaiting
100    /// `receiver.changed()` is cancellable inside `tokio::select!`.
101    pub(super) shutdown_watch: tokio::sync::watch::Sender<bool>,
102    /// Standardized loop observations. Updated inside `run()` after
103    /// every `do_tick`. Register via
104    /// [`Self::loop_metrics`] with the cluster registry.
105    pub(super) loop_metrics: Arc<LoopMetrics>,
106    /// Boot-time readiness signal. Flipped to `true` by
107    /// [`super::tick::do_tick`] after the first tick completes phase
108    /// 4 (apply committed entries) — i.e. once Raft has driven at
109    /// least one no-op or replayed entries past the persisted
110    /// applied watermark on this node.
111    ///
112    /// The host crate's `start_raft` returns `subscribe_ready()` to
113    /// `main.rs`, which awaits it before binding any client-facing
114    /// listener. This guarantees the first SQL DDL the operator
115    /// runs after process start cannot race against an
116    /// uninitialized metadata raft group, which previously surfaced
117    /// as `metadata propose: not leader` under fast restart loops.
118    pub(super) ready_watch: tokio::sync::watch::Sender<bool>,
119    /// Per-Raft-group apply watermark watchers. Bumped from
120    /// [`super::tick::do_tick`] after applies and from
121    /// [`super::handle_rpc`] after snapshot installs. The host crate
122    /// shares this Arc with `SharedState` so proposers, lease
123    /// renewals, and consistent reads can wait on the apply
124    /// watermark of the *specific* group whose proposal they made.
125    pub(super) group_watchers: Arc<GroupAppliedWatchers>,
126    /// Tracks whether this node was the metadata-group leader on the
127    /// previous tick. Used to detect false→true edges so the cluster
128    /// epoch (see [`crate::cluster_epoch`]) can be bumped exactly once
129    /// per leadership acquisition. `AtomicBool` because [`super::tick::do_tick`]
130    /// runs against `&self`.
131    pub(super) prev_metadata_leader: std::sync::atomic::AtomicBool,
132
133    /// Optional quarantine hook for the snapshot receive path.
134    ///
135    /// When set (by the `nodedb` binary via `with_snapshot_quarantine_hook`),
136    /// the `InstallSnapshotRequest` handler checks whether the incoming chunk
137    /// is already quarantined, records successes to clear transient strikes, and
138    /// records consecutive failures to quarantine persistently.
139    ///
140    /// Cluster-only tests leave this as `None`, which disables all quarantine
141    /// accounting for snapshot chunks.
142    pub(super) snapshot_quarantine_hook: Option<Arc<dyn SnapshotQuarantineHook>>,
143
144    /// Optional cross-node streaming-shuffle receiver (E1).
145    ///
146    /// When set (by the `nodedb` binary via `with_shuffle_receiver`), the
147    /// `ShufflePush` transport read-loop deposits chunks and records barrier
148    /// completion through it. Cluster-only tests leave this `None`, which makes
149    /// any incoming `ShufflePush` stream return a typed "not configured" error.
150    pub(super) shuffle_receiver: Option<Arc<dyn ShuffleReceiver>>,
151
152    /// Optional cross-node shuffle PRODUCER (E4a).
153    ///
154    /// When set (by the `nodedb` binary via `with_shuffle_producer`), an incoming
155    /// `ShuffleProduceRequest` runs the local scan + hash-partition fan-out
156    /// through it. Cluster-only tests leave this `None`, which makes any incoming
157    /// `ShuffleProduce` request return a typed "not configured" error.
158    pub(super) shuffle_producer: Option<Arc<dyn ShuffleProducer>>,
159
160    /// Optional cross-node shuffle CONSUMER (E4b).
161    ///
162    /// When set (by the `nodedb` binary via `with_shuffle_consumer`), an incoming
163    /// `ShuffleConsumeRequest` runs the node-local grace join over the part's
164    /// staged sides through it. Cluster-only tests leave this `None`, which makes
165    /// any incoming `ShuffleConsume` request return a typed "not configured"
166    /// error.
167    pub(super) shuffle_consumer: Option<Arc<dyn ShuffleConsumer>>,
168
169    /// Optional cross-node distributed GROUP BY shuffle CONSUMER (E5b).
170    ///
171    /// SINGLE-SIDED aggregate sibling of [`shuffle_consumer`](Self::shuffle_consumer).
172    /// When set (by the `nodedb` binary via `with_shuffle_aggregator`), an
173    /// incoming `ShuffleAggregateConsumeRequest` runs the node-local partial-state
174    /// merge + finalize over the part's single staged side through it.
175    /// Cluster-only tests leave this `None`, which makes any incoming
176    /// `ShuffleAggregateConsume` request return a typed "not configured" error.
177    pub(super) shuffle_aggregator: Option<Arc<dyn ShuffleAggregator>>,
178
179    /// Optional routed-surrogate-exchange assigner (F1b).
180    ///
181    /// When set (by the `nodedb` binary via `with_assign_remote_surrogate`), an
182    /// incoming `AssignSurrogateRequest` runs a LOCAL `SurrogateAssigner::assign`
183    /// for the endpoint key through it (this node IS the home vShard leader, so a
184    /// local assign yields the authoritative value). Cluster-only tests leave
185    /// this `None`, which makes any incoming `AssignSurrogate` request return a
186    /// typed "not configured" error.
187    pub(super) assign_remote_surrogate: Option<Arc<dyn AssignRemoteSurrogate>>,
188
189    /// Optional routed Calvin-submit hook (Cv1).
190    ///
191    /// When set (by the `nodedb` binary via `with_calvin_submit`), an incoming
192    /// `SubmitCalvinTxnRequest` submits the carried `TxClass` to THIS node's
193    /// Calvin sequencer inbox and awaits completion through it (this node IS the
194    /// sequencer-group leader, so the submit is actually sequenced and acked
195    /// here). Cluster-only tests leave this `None`, which makes any incoming
196    /// `SubmitCalvinTxn` request return a typed "not configured" error.
197    pub(super) calvin_submit: Option<Arc<dyn CalvinSubmit>>,
198
199    /// Optional routed Calvin-INBOX submit hook (Cv1).
200    ///
201    /// OLLP dependent sibling of [`calvin_submit`](Self::calvin_submit). When set
202    /// (by the `nodedb` binary via `with_calvin_submit_inbox`), an incoming
203    /// `SubmitCalvinInboxRequest` submits the carried `TxClass` to THIS node's
204    /// Calvin sequencer inbox and awaits only the ASSIGNMENT (not completion)
205    /// through it. Cluster-only tests leave this `None`, which makes any incoming
206    /// `SubmitCalvinInbox` request return a typed "not configured" error.
207    pub(super) calvin_submit_inbox: Option<Arc<dyn CalvinSubmitInbox>>,
208
209    /// Optional routed reserve-read hook (Calvin OLLP).
210    ///
211    /// When set (by the `nodedb` binary via `with_reserve_read`), an incoming
212    /// `ReserveReadRequest` assign-only reserves the read lock for the carried
213    /// `LockKey` through it (this node IS the sequencer-group leader, so the
214    /// reserve is enforced against the authoritative lock table). Cluster-only
215    /// tests leave this `None`, which makes any incoming `ReserveRead` request
216    /// return a typed "not configured" error.
217    pub(super) reserve_read: Option<Arc<dyn ReserveRead>>,
218
219    /// Optional routed release-reservation hook (Calvin OLLP).
220    ///
221    /// Ack-only sibling of [`reserve_read`](Self::reserve_read). When set (by
222    /// the `nodedb` binary via `with_release_reservation`), an incoming
223    /// `ReleaseReservationRequest` releases the reservation held by the
224    /// carried owner through it. Cluster-only tests leave this `None`, which
225    /// makes any incoming `ReleaseReservation` request return a typed "not
226    /// configured" error.
227    pub(super) release_reservation: Option<Arc<dyn ReleaseReservation>>,
228
229    /// Optional per-group snapshot builder for the SEND path.
230    ///
231    /// When set (by the `nodedb` binary via `with_snapshot_builder`), the
232    /// install-snapshot dispatch in [`super::tick`] calls it to produce the real
233    /// serialized engine state for the lagging follower's group vshards before
234    /// framing the chunked `InstallSnapshot` RPC. Cluster-only tests leave this
235    /// `None`, which makes the sender fall back to the stub (empty) chunk.
236    pub(super) snapshot_builder: Option<Arc<dyn SnapshotBuilder>>,
237
238    /// Optional per-group snapshot applier for the RECEIVE path.
239    ///
240    /// When set (by the `nodedb` binary via `with_snapshot_applier`), the
241    /// install-snapshot finalize path applies the received per-group snapshot to
242    /// the local Data-Plane state machine AFTER the atomic `.partial`→`.snap`
243    /// rename and BEFORE advancing Raft. Cluster-only tests leave this `None`,
244    /// which makes the follower advance Raft without restoring engine state
245    /// (correct for the empty bootstrap stub shipped by those tests).
246    pub(super) snapshot_applier: Option<Arc<dyn SnapshotApplier>>,
247
248    /// In-progress partial snapshot receives, keyed by `group_id`.
249    ///
250    /// Each entry tracks the `.partial` file, running CRC, and expected next
251    /// byte offset for a follower that is currently receiving a chunked snapshot.
252    /// Entries are created on `offset == 0`, updated on each chunk, and removed
253    /// on finalization or offset regression.
254    pub(super) partial_snapshots: Arc<crate::install_snapshot::PartialSnapshotMap>,
255
256    /// Data directory for persistent partial-snapshot files and the
257    /// `recv_snapshots/` subdirectory. `None` in tests that don't exercise
258    /// the disk path.
259    pub(super) data_dir: Option<std::path::PathBuf>,
260
261    /// Snapshot chunk size for the sender path (bytes).
262    pub(super) snapshot_chunk_bytes: u64,
263
264    /// Orphan partial-snapshot max age for the GC sweeper (seconds).
265    pub(super) orphan_partial_max_age_secs: u64,
266
267    /// Cluster replication factor (target voters per group), loaded once from
268    /// `ClusterSettings` at startup; immutable for the loop's lifetime. Used
269    /// to cap voter promotion at min(RF, N). Defaults to 1 (single-node, no
270    /// extra replicas) when not overridden via `with_replication_factor`.
271    pub(super) replication_factor: u32,
272
273    /// Monotonic tick counter; throttles periodic maintenance (placement
274    /// reconcile) to a coarse cadence off the 10ms tick. `AtomicU64` because
275    /// [`super::tick::do_tick`] runs against `&self`.
276    pub(super) tick_count: std::sync::atomic::AtomicU64,
277
278    /// Notification channel for kicking placement reconcile immediately
279    /// when a node joins, instead of waiting up to ~1 s for the throttled
280    /// tick. Fired from [`super::join`] on the success path; consumed by
281    /// an extra `select!` arm in [`Self::run`].
282    pub(super) reconcile_notify: tokio::sync::Notify,
283}
284
285impl<A: CommitApplier> RaftLoop<A> {
286    pub fn new(
287        multi_raft: MultiRaft,
288        transport: Arc<NexarTransport>,
289        topology: Arc<RwLock<ClusterTopology>>,
290        applier: A,
291    ) -> Self {
292        let node_id = multi_raft.node_id();
293        let (shutdown_watch, _) = tokio::sync::watch::channel(false);
294        let (ready_watch, _) = tokio::sync::watch::channel(false);
295        Self {
296            node_id,
297            multi_raft: Arc::new(Mutex::new(multi_raft)),
298            transport,
299            topology,
300            applier,
301            metadata_applier: Arc::new(NoopMetadataApplier),
302            plan_executor: Arc::new(NoopPlanExecutor),
303            tick_interval: DEFAULT_TICK_INTERVAL,
304            vshard_handler: None,
305            catalog: None,
306            shutdown_watch,
307            ready_watch,
308            loop_metrics: LoopMetrics::new("raft_tick_loop"),
309            group_watchers: Arc::new(GroupAppliedWatchers::new()),
310            prev_metadata_leader: std::sync::atomic::AtomicBool::new(false),
311            snapshot_quarantine_hook: None,
312            shuffle_receiver: None,
313            shuffle_producer: None,
314            shuffle_consumer: None,
315            shuffle_aggregator: None,
316            assign_remote_surrogate: None,
317            calvin_submit: None,
318            calvin_submit_inbox: None,
319            reserve_read: None,
320            release_reservation: None,
321            snapshot_builder: None,
322            snapshot_applier: None,
323            partial_snapshots: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
324            data_dir: None,
325            snapshot_chunk_bytes: 4 * 1024 * 1024,
326            orphan_partial_max_age_secs: 300,
327            replication_factor: 1,
328            tick_count: std::sync::atomic::AtomicU64::new(0),
329            reconcile_notify: tokio::sync::Notify::new(),
330        }
331    }
332}
333
334impl<A: CommitApplier, P: PlanExecutor> RaftLoop<A, P> {
335    /// Install the snapshot quarantine hook (mutable setter variant).
336    ///
337    /// Prefer `with_snapshot_quarantine_hook` on the builder chain unless you
338    /// need to set the hook after construction.
339    pub fn set_snapshot_quarantine_hook(&mut self, hook: Arc<dyn SnapshotQuarantineHook>) {
340        self.snapshot_quarantine_hook = Some(hook);
341    }
342
343    /// Shared handle to the per-group apply watcher registry.
344    pub fn group_watchers(&self) -> Arc<GroupAppliedWatchers> {
345        Arc::clone(&self.group_watchers)
346    }
347
348    /// Shared handle to this loop's standardized metrics.
349    pub fn loop_metrics(&self) -> Arc<LoopMetrics> {
350        Arc::clone(&self.loop_metrics)
351    }
352
353    /// Count of Raft groups currently mounted on this node — used to
354    /// render the `raft_tick_loop_pending_groups` gauge.
355    pub fn pending_groups(&self) -> usize {
356        let mr = self.multi_raft.lock().unwrap_or_else(|p| p.into_inner());
357        mr.group_count()
358    }
359
360    /// Signal cooperative shutdown to every detached task spawned
361    /// inside [`super::tick::do_tick`].
362    ///
363    /// This is the entry point for test harnesses that want to
364    /// tear down a `RaftLoop` without waiting for the external
365    /// `run()` shutdown watch channel to propagate. In production
366    /// the same signal is emitted automatically by `run()` when
367    /// its external shutdown receiver fires.
368    ///
369    /// Idempotent: calling this multiple times is a no-op after
370    /// the first.
371    pub fn begin_shutdown(&self) {
372        let _ = self.shutdown_watch.send(true);
373    }
374
375    /// Subscribe to the boot-time readiness signal.
376    ///
377    /// The returned receiver starts at `false` and flips to `true`
378    /// exactly once, after the first [`super::tick::do_tick`]
379    /// completes phase 4 (apply committed entries). Used by the
380    /// host crate to gate client-facing listener startup until the
381    /// metadata raft group has produced its first applied entry.
382    pub fn subscribe_ready(&self) -> tokio::sync::watch::Receiver<bool> {
383        self.ready_watch.subscribe()
384    }
385
386    /// This node's id (exposed for handlers and tests).
387    pub fn node_id(&self) -> u64 {
388        self.node_id
389    }
390
391    /// Target replication factor for this cluster (voters per group).
392    ///
393    /// Loaded once from `ClusterSettings` at startup and immutable for the
394    /// loop's lifetime. Returns `1` when no override was set (single-node
395    /// default).
396    pub fn replication_factor(&self) -> u32 {
397        self.replication_factor
398    }
399
400    /// Run the event loop until shutdown.
401    ///
402    /// This drives Raft elections, heartbeats, and message dispatch.
403    /// Call [`NexarTransport::serve`] separately with `Arc<Self>` as the handler.
404    ///
405    /// When the externally-supplied `shutdown` receiver fires,
406    /// the loop also propagates the signal to the internal
407    /// cooperative-shutdown channel so every detached task
408    /// spawned inside `do_tick` exits promptly and drops its
409    /// `Arc<Mutex<MultiRaft>>` clone.
410    pub async fn run(&self, mut shutdown: tokio::sync::watch::Receiver<bool>) {
411        let mut interval = tokio::time::interval(self.tick_interval);
412        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
413        self.loop_metrics.set_up(true);
414
415        // Startup GC sweep: remove orphaned partial-snapshot files from
416        // previous runs that did not complete.
417        if let Some(ref dir) = self.data_dir {
418            match crate::install_snapshot::gc::sweep_orphans(dir, self.orphan_partial_max_age_secs)
419            {
420                Ok((removed, errs)) => {
421                    if removed > 0 {
422                        tracing::info!(removed, "startup: removed orphaned partial snapshot files");
423                    }
424                    for e in errs {
425                        tracing::warn!(error = %e, "startup: partial snapshot GC error");
426                    }
427                }
428                Err(e) => {
429                    tracing::warn!(error = %e, "startup: failed to sweep partial snapshot directory");
430                }
431            }
432        }
433
434        loop {
435            tokio::select! {
436                _ = interval.tick() => {
437                    let started = Instant::now();
438                    self.do_tick();
439                    self.loop_metrics.observe(started.elapsed());
440                }
441                _ = self.reconcile_notify.notified() => {
442                    if *shutdown.borrow() {
443                        return;
444                    }
445                    self.reconcile_placement();
446                }
447                _ = shutdown.changed() => {
448                    if *shutdown.borrow() {
449                        debug!("raft loop shutting down");
450                        self.begin_shutdown();
451                        break;
452                    }
453                }
454            }
455        }
456        self.loop_metrics.set_up(false);
457    }
458
459    /// Returns the inner multi-raft handle. Exposed for tests and for
460    /// the host crate's metadata proposer so it can hold a second
461    /// reference to the same underlying mutex without pulling the
462    /// whole raft loop into the caller's lifetime.
463    pub fn multi_raft_handle(&self) -> Arc<Mutex<crate::multi_raft::MultiRaft>> {
464        self.multi_raft.clone()
465    }
466
467    /// Snapshot all Raft group states for observability (SHOW RAFT GROUPS).
468    pub fn group_statuses(&self) -> Vec<crate::multi_raft::GroupStatus> {
469        let mr = self.multi_raft.lock().unwrap_or_else(|p| p.into_inner());
470        mr.group_statuses()
471    }
472}
473
474#[cfg(test)]
475mod tests {
476    use super::*;
477    use crate::routing::RoutingTable;
478    use nodedb_types::config::tuning::ClusterTransportTuning;
479    use std::sync::atomic::{AtomicU64, Ordering};
480    use std::time::Instant;
481
482    /// Test applier that counts applied entries across both data and
483    /// metadata groups. The metadata-group variant ([`CountingMetadataApplier`])
484    /// increments the same counter so tests that propose against group 0
485    /// (the metadata group) still see the count move.
486    pub(crate) struct CountingApplier {
487        applied: Arc<AtomicU64>,
488    }
489
490    impl CountingApplier {
491        pub(crate) fn new() -> Self {
492            Self {
493                applied: Arc::new(AtomicU64::new(0)),
494            }
495        }
496
497        pub(crate) fn count(&self) -> u64 {
498            self.applied.load(Ordering::Relaxed)
499        }
500
501        pub(crate) fn metadata_applier(&self) -> Arc<CountingMetadataApplier> {
502            Arc::new(CountingMetadataApplier {
503                applied: self.applied.clone(),
504            })
505        }
506    }
507
508    impl CommitApplier for CountingApplier {
509        fn apply_committed(&self, _group_id: u64, entries: &[LogEntry]) -> u64 {
510            self.applied
511                .fetch_add(entries.len() as u64, Ordering::Relaxed);
512            entries.last().map(|e| e.index).unwrap_or(0)
513        }
514    }
515
516    pub(crate) struct CountingMetadataApplier {
517        applied: Arc<AtomicU64>,
518    }
519
520    impl MetadataApplier for CountingMetadataApplier {
521        fn apply(&self, entries: &[(u64, Vec<u8>)]) -> u64 {
522            self.applied
523                .fetch_add(entries.len() as u64, Ordering::Relaxed);
524            entries.last().map(|(idx, _)| *idx).unwrap_or(0)
525        }
526    }
527
528    /// Helper: create a transport on an ephemeral port.
529    fn make_transport(node_id: u64) -> Arc<NexarTransport> {
530        Arc::new(
531            NexarTransport::new(
532                node_id,
533                "127.0.0.1:0".parse().unwrap(),
534                crate::transport::credentials::TransportCredentials::Insecure,
535            )
536            .unwrap(),
537        )
538    }
539
540    /// Verify that `with_replication_factor` stores the supplied value and that
541    /// the default (no builder call) is 1.
542    #[tokio::test]
543    async fn replication_factor_accessor() {
544        let dir = tempfile::tempdir().unwrap();
545        let transport = make_transport(1);
546        let rt = RoutingTable::uniform(1, &[1], 1);
547        let mut mr = MultiRaft::new(1, rt, dir.path().to_path_buf());
548        mr.add_group(0, vec![]).unwrap();
549        mr.add_group(1, vec![]).unwrap();
550
551        let topo = Arc::new(RwLock::new(ClusterTopology::new()));
552
553        // Default: 1 (single-node sentinel, no builder call).
554        let loop_default = RaftLoop::new(
555            MultiRaft::new(
556                1,
557                RoutingTable::uniform(1, &[1], 1),
558                tempfile::tempdir().unwrap().path().to_path_buf(),
559            ),
560            make_transport(1),
561            topo.clone(),
562            CountingApplier::new(),
563        );
564        assert_eq!(loop_default.replication_factor(), 1);
565
566        // Explicit override via builder.
567        let loop_rf3 =
568            RaftLoop::new(mr, transport, topo, CountingApplier::new()).with_replication_factor(3);
569        assert_eq!(loop_rf3.replication_factor(), 3);
570    }
571
572    #[tokio::test]
573    async fn single_node_raft_loop_commits() {
574        let dir = tempfile::tempdir().unwrap();
575        let transport = make_transport(1);
576        // uniform(1, ...) creates metadata group 0 + data group 1.
577        let rt = RoutingTable::uniform(1, &[1], 1);
578        let mut mr = MultiRaft::new(1, rt, dir.path().to_path_buf());
579        // Add both the metadata group (0) and the data group (1).
580        mr.add_group(0, vec![]).unwrap();
581        mr.add_group(1, vec![]).unwrap();
582
583        for node in mr.groups_mut().values_mut() {
584            node.election_deadline_override(Instant::now() - Duration::from_millis(1));
585        }
586
587        let applier = CountingApplier::new();
588        let meta = applier.metadata_applier();
589        let topo = Arc::new(RwLock::new(ClusterTopology::new()));
590        let raft_loop =
591            Arc::new(RaftLoop::new(mr, transport, topo, applier).with_metadata_applier(meta));
592
593        let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);
594
595        let rl = raft_loop.clone();
596        let run_handle = tokio::spawn(async move {
597            rl.run(shutdown_rx).await;
598        });
599
600        tokio::time::sleep(Duration::from_millis(50)).await;
601
602        assert!(
603            raft_loop.applier.count() >= 1,
604            "expected at least 1 applied entry (no-op), got {}",
605            raft_loop.applier.count()
606        );
607
608        // vshard 0 maps to data group 1 (not metadata group 0).
609        let (_gid, idx) = raft_loop.propose(0, b"hello".to_vec()).unwrap();
610        assert!(idx >= 2);
611
612        tokio::time::sleep(Duration::from_millis(50)).await;
613
614        assert!(
615            raft_loop.applier.count() >= 2,
616            "expected at least 2 applied entries, got {}",
617            raft_loop.applier.count()
618        );
619
620        shutdown_tx.send(true).unwrap();
621        run_handle.abort();
622    }
623
624    #[tokio::test]
625    async fn three_node_election_over_quic() {
626        let t1 = make_transport(1);
627        let t2 = make_transport(2);
628        let t3 = make_transport(3);
629
630        t1.register_peer(2, t2.local_addr());
631        t1.register_peer(3, t3.local_addr());
632        t2.register_peer(1, t1.local_addr());
633        t2.register_peer(3, t3.local_addr());
634        t3.register_peer(1, t1.local_addr());
635        t3.register_peer(2, t2.local_addr());
636
637        // uniform(1, ...) creates metadata group 0 + data group 1.
638        // Both are added to every MultiRaft so vshard proposals (group 1)
639        // and metadata proposals (group 0) both work.
640        let rt = RoutingTable::uniform(1, &[1, 2, 3], 3);
641
642        let dir1 = tempfile::tempdir().unwrap();
643        let mut mr1 = MultiRaft::new(1, rt.clone(), dir1.path().to_path_buf());
644        mr1.add_group(0, vec![2, 3]).unwrap();
645        mr1.add_group(1, vec![2, 3]).unwrap();
646        for node in mr1.groups_mut().values_mut() {
647            node.election_deadline_override(Instant::now() - Duration::from_millis(1));
648        }
649
650        let transport_tuning = ClusterTransportTuning::default();
651        let election_timeout_min =
652            Duration::from_millis(transport_tuning.effective_election_timeout_min_ms());
653        let election_timeout_max =
654            Duration::from_millis(transport_tuning.effective_election_timeout_max_ms());
655
656        let dir2 = tempfile::tempdir().unwrap();
657        let mut mr2 = MultiRaft::new(2, rt.clone(), dir2.path().to_path_buf())
658            .with_election_timeout(election_timeout_min, election_timeout_max);
659        mr2.add_group(0, vec![1, 3]).unwrap();
660        mr2.add_group(1, vec![1, 3]).unwrap();
661
662        let dir3 = tempfile::tempdir().unwrap();
663        let mut mr3 = MultiRaft::new(3, rt.clone(), dir3.path().to_path_buf())
664            .with_election_timeout(election_timeout_min, election_timeout_max);
665        mr3.add_group(0, vec![1, 2]).unwrap();
666        mr3.add_group(1, vec![1, 2]).unwrap();
667
668        let a1 = CountingApplier::new();
669        let m1 = a1.metadata_applier();
670        let a2 = CountingApplier::new();
671        let m2 = a2.metadata_applier();
672        let a3 = CountingApplier::new();
673        let m3 = a3.metadata_applier();
674
675        let topo1 = Arc::new(RwLock::new(ClusterTopology::new()));
676        let topo2 = Arc::new(RwLock::new(ClusterTopology::new()));
677        let topo3 = Arc::new(RwLock::new(ClusterTopology::new()));
678
679        let rl1 = Arc::new(RaftLoop::new(mr1, t1.clone(), topo1, a1).with_metadata_applier(m1));
680        let rl2 = Arc::new(RaftLoop::new(mr2, t2.clone(), topo2, a2).with_metadata_applier(m2));
681        let rl3 = Arc::new(RaftLoop::new(mr3, t3.clone(), topo3, a3).with_metadata_applier(m3));
682
683        let (shutdown_tx, _) = tokio::sync::watch::channel(false);
684
685        let rl2_h = rl2.clone();
686        let sr2 = shutdown_tx.subscribe();
687        tokio::spawn(async move { t2.serve(rl2_h, sr2).await });
688
689        let rl3_h = rl3.clone();
690        let sr3 = shutdown_tx.subscribe();
691        tokio::spawn(async move { t3.serve(rl3_h, sr3).await });
692
693        let rl1_r = rl1.clone();
694        let sr1 = shutdown_tx.subscribe();
695        tokio::spawn(async move { rl1_r.run(sr1).await });
696
697        let rl2_r = rl2.clone();
698        let sr2r = shutdown_tx.subscribe();
699        tokio::spawn(async move { rl2_r.run(sr2r).await });
700
701        let rl3_r = rl3.clone();
702        let sr3r = shutdown_tx.subscribe();
703        tokio::spawn(async move { rl3_r.run(sr3r).await });
704
705        let rl1_h = rl1.clone();
706        let sr1h = shutdown_tx.subscribe();
707        tokio::spawn(async move { t1.serve(rl1_h, sr1h).await });
708
709        // Poll until node 1 commits at least the no-op (election done).
710        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
711        loop {
712            if rl1.applier.count() >= 1 {
713                break;
714            }
715            assert!(
716                tokio::time::Instant::now() < deadline,
717                "node 1 should have committed at least the no-op, got {}",
718                rl1.applier.count()
719            );
720            tokio::time::sleep(Duration::from_millis(20)).await;
721        }
722
723        let (_gid, idx) = rl1.propose(0, b"distributed-cmd".to_vec()).unwrap();
724        assert!(idx >= 2);
725
726        // Poll until all nodes replicate the proposed command.
727        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
728        loop {
729            if rl1.applier.count() >= 2 && rl2.applier.count() >= 1 && rl3.applier.count() >= 1 {
730                break;
731            }
732            assert!(
733                tokio::time::Instant::now() < deadline,
734                "replication timed out: n1={}, n2={}, n3={}",
735                rl1.applier.count(),
736                rl2.applier.count(),
737                rl3.applier.count()
738            );
739            tokio::time::sleep(Duration::from_millis(20)).await;
740        }
741
742        shutdown_tx.send(true).unwrap();
743    }
744}