Skip to main content

reddb_server/replication/
dst.rs

1//! Deterministic simulation testing helpers for the replication control plane.
2//!
3//! The simulator is intentionally in-process and single-threaded. It gives
4//! control-plane tests a transport that can inject partition, delay, reorder,
5//! and message loss without opening sockets or depending on Tokio scheduler
6//! ordering.
7
8use std::collections::BTreeSet;
9use std::time::Duration;
10
11use crate::replication::{VoteDecision, VoteRequest};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum ReplicationControlMessage {
15    ElectionVoteRequest(VoteRequest),
16    ElectionVoteDecision(VoteDecision),
17    LogicalCommit {
18        term: u64,
19        lsn: u64,
20        payload_hash: String,
21    },
22    LeaseProbe {
23        holder_id: String,
24        term: u64,
25    },
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct NetworkFaults {
30    pub loss_per_million: u32,
31    pub max_delay_ms: u64,
32    pub reorder: bool,
33}
34
35impl NetworkFaults {
36    pub fn reliable() -> Self {
37        Self {
38            loss_per_million: 0,
39            max_delay_ms: 0,
40            reorder: false,
41        }
42    }
43
44    pub fn lossy(loss_per_million: u32) -> Self {
45        Self {
46            loss_per_million,
47            max_delay_ms: 0,
48            reorder: false,
49        }
50    }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct SimulationClock {
55    now_ms: u64,
56}
57
58impl SimulationClock {
59    pub fn new() -> Self {
60        Self { now_ms: 0 }
61    }
62
63    pub fn now_ms(&self) -> u64 {
64        self.now_ms
65    }
66
67    pub fn elapsed(&self) -> Duration {
68        Duration::from_millis(self.now_ms)
69    }
70
71    pub fn advance(&mut self, by: Duration) {
72        let millis = u64::try_from(by.as_millis()).unwrap_or(u64::MAX);
73        self.now_ms = self.now_ms.saturating_add(millis);
74    }
75
76    fn advance_to(&mut self, now_ms: u64) {
77        self.now_ms = self.now_ms.max(now_ms);
78    }
79}
80
81impl Default for SimulationClock {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct Delivered<M> {
89    pub from: String,
90    pub to: String,
91    pub message: M,
92    pub delivered_at_ms: u64,
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
96pub enum SendOutcome {
97    Accepted { deliver_at_ms: u64 },
98    Dropped(DropReason),
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum DropReason {
103    Partition,
104    Loss,
105}
106
107#[derive(Debug, Clone)]
108struct Pending<M> {
109    from: String,
110    to: String,
111    message: M,
112    deliver_at_ms: u64,
113    order: u64,
114}
115
116#[derive(Debug, Clone)]
117pub struct InProcessReplicationNetwork<M> {
118    clock: SimulationClock,
119    rng: SeededRng,
120    faults: NetworkFaults,
121    partitions: BTreeSet<(String, String)>,
122    pending: Vec<Pending<M>>,
123    sequence: u64,
124}
125
126impl<M> InProcessReplicationNetwork<M> {
127    pub fn new(seed: u64, faults: NetworkFaults) -> Self {
128        Self {
129            clock: SimulationClock::new(),
130            rng: SeededRng::new(seed),
131            faults,
132            partitions: BTreeSet::new(),
133            pending: Vec::new(),
134            sequence: 0,
135        }
136    }
137
138    pub fn clock(&self) -> SimulationClock {
139        self.clock
140    }
141
142    pub fn advance(&mut self, by: Duration) {
143        self.clock.advance(by);
144    }
145
146    pub fn partition(&mut self, a: impl Into<String>, b: impl Into<String>) {
147        self.partitions.insert(partition_key(a.into(), b.into()));
148    }
149
150    pub fn heal(&mut self, a: &str, b: &str) {
151        self.partitions
152            .remove(&partition_key(a.to_string(), b.to_string()));
153    }
154
155    pub fn send(
156        &mut self,
157        from: impl Into<String>,
158        to: impl Into<String>,
159        message: M,
160    ) -> SendOutcome {
161        let from = from.into();
162        let to = to.into();
163        if self
164            .partitions
165            .contains(&partition_key(from.clone(), to.clone()))
166        {
167            return SendOutcome::Dropped(DropReason::Partition);
168        }
169        if self.faults.loss_per_million > 0
170            && self.rng.next_bounded(1_000_000) < u64::from(self.faults.loss_per_million)
171        {
172            return SendOutcome::Dropped(DropReason::Loss);
173        }
174
175        self.sequence = self.sequence.saturating_add(1);
176        let delay = if self.faults.max_delay_ms == 0 {
177            0
178        } else {
179            self.rng
180                .next_bounded(self.faults.max_delay_ms.saturating_add(1))
181        };
182        let deliver_at_ms = self.clock.now_ms().saturating_add(delay);
183        let order = if self.faults.reorder {
184            self.rng.next_u64()
185        } else {
186            self.sequence
187        };
188        self.pending.push(Pending {
189            from,
190            to,
191            message,
192            deliver_at_ms,
193            order,
194        });
195        SendOutcome::Accepted { deliver_at_ms }
196    }
197
198    pub fn advance_to_next_delivery(&mut self) -> bool {
199        let Some(next) = self.pending.iter().map(|p| p.deliver_at_ms).min() else {
200            return false;
201        };
202        self.clock.advance_to(next);
203        true
204    }
205
206    pub fn drain_ready_for(&mut self, recipient: &str) -> Vec<Delivered<M>> {
207        let now = self.clock.now_ms();
208        let mut ready = Vec::new();
209        let mut pending = Vec::with_capacity(self.pending.len());
210        for msg in self.pending.drain(..) {
211            if msg.to == recipient && msg.deliver_at_ms <= now {
212                ready.push(msg);
213            } else {
214                pending.push(msg);
215            }
216        }
217        self.pending = pending;
218        ready.sort_by_key(|msg| (msg.deliver_at_ms, msg.order));
219        ready
220            .into_iter()
221            .map(|msg| Delivered {
222                from: msg.from,
223                to: msg.to,
224                message: msg.message,
225                delivered_at_ms: msg.deliver_at_ms,
226            })
227            .collect()
228    }
229}
230
231#[derive(Debug, Clone)]
232struct SeededRng {
233    state: u64,
234}
235
236impl SeededRng {
237    fn new(seed: u64) -> Self {
238        let state = if seed == 0 {
239            0x9E37_79B9_7F4A_7C15
240        } else {
241            seed
242        };
243        Self { state }
244    }
245
246    fn next_u64(&mut self) -> u64 {
247        let mut x = self.state;
248        x ^= x << 13;
249        x ^= x >> 7;
250        x ^= x << 17;
251        self.state = x;
252        x
253    }
254
255    fn next_bounded(&mut self, upper_exclusive: u64) -> u64 {
256        if upper_exclusive == 0 {
257            0
258        } else {
259            self.next_u64() % upper_exclusive
260        }
261    }
262}
263
264fn partition_key(a: String, b: String) -> (String, String) {
265    if a <= b {
266        (a, b)
267    } else {
268        (b, a)
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use std::collections::{BTreeMap, BTreeSet};
275    use std::rc::Rc;
276    use std::sync::Arc;
277
278    use super::*;
279    use crate::replication::{
280        ElectionCoordinator, ElectionOutcome, ElectionRequest, ElectionTransport, LastVote,
281        LastVoteError, LastVoteStore, LeaseError, LeaseStore, Member, MemoryLastVoteStore,
282        RefusalReason, Voter, WriterLease,
283    };
284
285    #[test]
286    fn fault_injection_is_seed_reproducible() {
287        let trace_a = delivery_trace(0xD57, 0);
288        let trace_b = delivery_trace(0xD57, 0);
289        let trace_c = delivery_trace(0xD58, 0);
290
291        assert_eq!(trace_a, trace_b, "same seed must reproduce the trace");
292        assert_ne!(
293            trace_a, trace_c,
294            "different seed should explore a different trace"
295        );
296
297        let mut partitioned = InProcessReplicationNetwork::new(1, NetworkFaults::reliable());
298        partitioned.partition("a", "b");
299        assert_eq!(
300            partitioned.send("a", "b", 1u64),
301            SendOutcome::Dropped(DropReason::Partition)
302        );
303
304        let mut lossy = InProcessReplicationNetwork::new(1, NetworkFaults::lossy(1_000_000));
305        assert_eq!(
306            lossy.send("a", "b", 1u64),
307            SendOutcome::Dropped(DropReason::Loss)
308        );
309    }
310
311    fn delivery_trace(seed: u64, loss_per_million: u32) -> Vec<(u64, u64)> {
312        let faults = NetworkFaults {
313            loss_per_million,
314            max_delay_ms: 25,
315            reorder: true,
316        };
317        let mut network = InProcessReplicationNetwork::new(seed, faults);
318        for value in 0..12u64 {
319            let _ = network.send("a", "b", value);
320        }
321        network.advance(Duration::from_millis(25));
322        network
323            .drain_ready_for("b")
324            .into_iter()
325            .map(|msg| (msg.delivered_at_ms, msg.message))
326            .collect()
327    }
328
329    #[test]
330    fn election_safety_under_partition_has_at_most_one_leader_per_term() {
331        let members = five_voters();
332        let stores = shared_vote_stores(&members);
333        let mut network = partitioned_network(0x1358);
334        for peer in ["d", "e"] {
335            network.partition("a", peer);
336            network.partition("b", peer);
337            network.partition("c", peer);
338        }
339
340        let mut leaders = BTreeMap::new();
341        for candidate in [
342            candidate_request("a", 4, 120, 100),
343            candidate_request("d", 4, 120, 100),
344        ] {
345            let mut tx = NetworkElectionTransport::new(
346                &mut network,
347                members.clone(),
348                stores.clone(),
349                candidate.candidate.id.clone(),
350                100,
351            );
352            let outcome = ElectionCoordinator::run(&candidate, &mut tx, Duration::from_secs(60));
353            if let ElectionOutcome::Elected { term, .. } = outcome {
354                let previous = leaders.insert(term, candidate.candidate.id.clone());
355                assert_eq!(
356                    previous, None,
357                    "two leaders elected in term {term}: {previous:?} and {:?}",
358                    candidate.candidate.id
359                );
360            }
361        }
362
363        assert_eq!(leaders.get(&5), Some(&"a".to_string()));
364    }
365
366    #[test]
367    fn partitioned_elections_do_not_split_brain_or_lose_committed_writes() {
368        let committed_watermark = 100;
369        let committed_writes: BTreeSet<u64> = (1..=committed_watermark).collect();
370
371        for seed in 1..=24 {
372            let members = five_voters();
373            let stores = shared_vote_stores(&members);
374            let mut network = partitioned_network(seed);
375            let mut elected = Vec::new();
376            let candidates = if seed % 2 == 0 {
377                [("a", 120), ("d", 80)]
378            } else {
379                [("d", 80), ("a", 120)]
380            };
381
382            for (id, lsn) in candidates {
383                let req = candidate_request(id, 4, lsn, committed_watermark);
384                let mut tx = NetworkElectionTransport::new(
385                    &mut network,
386                    members.clone(),
387                    stores.clone(),
388                    id.to_string(),
389                    committed_watermark,
390                );
391                let outcome = ElectionCoordinator::run(&req, &mut tx, Duration::from_secs(60));
392                if let ElectionOutcome::Elected { term, .. } = outcome {
393                    elected.push((term, id.to_string(), lsn));
394                }
395            }
396
397            let mut leaders_by_term = BTreeSet::new();
398            for (term, id, lsn) in elected {
399                assert!(
400                    leaders_by_term.insert(term),
401                    "split-brain in term {term} under seed {seed}"
402                );
403                assert!(
404                    lsn >= committed_watermark,
405                    "leader {id} lost committed writes under seed {seed}"
406                );
407                assert!(
408                    committed_writes.iter().all(|committed| *committed <= lsn),
409                    "leader {id} does not cover all committed writes under seed {seed}"
410                );
411            }
412        }
413    }
414
415    #[test]
416    fn lease_fencing_holds_when_a_partitioned_primary_returns_stale() {
417        let members = five_voters();
418        let stores = shared_vote_stores(&members);
419        let mut network = partitioned_network(0x715);
420        for peer in ["a", "b"] {
421            network.partition("old-primary", peer);
422        }
423        let promoted = candidate_request("a", 4, 150, 100);
424        let mut tx =
425            NetworkElectionTransport::new(&mut network, members, stores, "a".to_string(), 100);
426
427        let outcome = ElectionCoordinator::run(&promoted, &mut tx, Duration::from_secs(60));
428        let ElectionOutcome::Elected { term: new_term, .. } = outcome else {
429            panic!("expected a replacement primary, got {outcome:?}");
430        };
431
432        let store = lease_store("dst-fence");
433        let lease = store
434            .try_acquire_for_term("main", "new-primary", 60_000, new_term)
435            .expect("new primary lease");
436        assert_eq!(lease.term, new_term);
437
438        let err = store
439            .try_acquire_for_term("main", "old-primary", 60_000, new_term - 1)
440            .expect_err("stale partitioned primary must be fenced");
441        assert!(
442            matches!(
443                err,
444                LeaseError::Fenced {
445                    current_term,
446                    ..
447                } if current_term == new_term
448            ),
449            "got {err:?}"
450        );
451
452        let stale_lease = WriterLease {
453            database_key: "main".to_string(),
454            holder_id: "old-primary".to_string(),
455            term: new_term - 1,
456            generation: 1,
457            acquired_at_ms: 0,
458            expires_at_ms: u64::MAX,
459        };
460        assert!(stale_lease.fenced_by_term(new_term));
461    }
462
463    #[test]
464    #[ignore = "heavy seed sweep runs nightly in CI"]
465    fn dst_seed_sweep_election_safety_no_split_brain_no_lost_committed_writes() {
466        for seed in 1..=256 {
467            let members = five_voters();
468            let stores = shared_vote_stores(&members);
469            let mut network = partitioned_network(seed);
470            let mut leaders = BTreeMap::new();
471            for (id, lsn) in [("a", 125), ("b", 130), ("d", 90), ("e", 95)] {
472                let req = candidate_request(id, 7, lsn, 100);
473                let mut tx = NetworkElectionTransport::new(
474                    &mut network,
475                    members.clone(),
476                    stores.clone(),
477                    id.to_string(),
478                    100,
479                );
480                if let ElectionOutcome::Elected { term, .. } =
481                    ElectionCoordinator::run(&req, &mut tx, Duration::from_secs(60))
482                {
483                    assert!(lsn >= 100, "seed {seed}: elected {id} below watermark");
484                    assert_eq!(
485                        leaders.insert(term, id.to_string()),
486                        None,
487                        "seed {seed}: more than one leader in term {term}"
488                    );
489                }
490            }
491        }
492    }
493
494    struct NetworkElectionTransport<'a> {
495        network: &'a mut InProcessReplicationNetwork<ReplicationControlMessage>,
496        members: Vec<Member>,
497        stores: BTreeMap<String, Rc<MemoryLastVoteStore>>,
498        candidate_id: String,
499        watermark: u64,
500        bumped_term: Option<u64>,
501        promoted_term: Option<u64>,
502    }
503
504    impl<'a> NetworkElectionTransport<'a> {
505        fn new(
506            network: &'a mut InProcessReplicationNetwork<ReplicationControlMessage>,
507            members: Vec<Member>,
508            stores: BTreeMap<String, Rc<MemoryLastVoteStore>>,
509            candidate_id: String,
510            watermark: u64,
511        ) -> Self {
512            Self {
513                network,
514                members,
515                stores,
516                candidate_id,
517                watermark,
518                bumped_term: None,
519                promoted_term: None,
520            }
521        }
522    }
523
524    impl ElectionTransport for NetworkElectionTransport<'_> {
525        fn members(&self) -> Vec<Member> {
526            self.members.clone()
527        }
528
529        fn request_vote(&mut self, peer_id: &str, req: &VoteRequest) -> VoteDecision {
530            let outcome = self.network.send(
531                self.candidate_id.clone(),
532                peer_id.to_string(),
533                ReplicationControlMessage::ElectionVoteRequest(req.clone()),
534            );
535            if !matches!(outcome, SendOutcome::Accepted { .. }) {
536                return unreachable_refusal(req);
537            }
538            if !self.network.advance_to_next_delivery() {
539                return unreachable_refusal(req);
540            }
541            let requests = self.network.drain_ready_for(peer_id);
542            let Some(request) = requests
543                .into_iter()
544                .find_map(|delivery| match delivery.message {
545                    ReplicationControlMessage::ElectionVoteRequest(request) => Some(request),
546                    _ => None,
547                })
548            else {
549                return unreachable_refusal(req);
550            };
551
552            let store = self.stores.get(peer_id).expect("known voter").clone();
553            let voter = Voter::new(peer_id, RcStore(store));
554            let decision = voter
555                .consider(&request, self.watermark)
556                .expect("memory vote store");
557            let outcome = self.network.send(
558                peer_id.to_string(),
559                self.candidate_id.clone(),
560                ReplicationControlMessage::ElectionVoteDecision(decision.clone()),
561            );
562            if !matches!(outcome, SendOutcome::Accepted { .. }) {
563                return unreachable_refusal(req);
564            }
565            if !self.network.advance_to_next_delivery() {
566                return unreachable_refusal(req);
567            }
568            self.network
569                .drain_ready_for(&self.candidate_id)
570                .into_iter()
571                .find_map(|delivery| match delivery.message {
572                    ReplicationControlMessage::ElectionVoteDecision(decision) => Some(decision),
573                    _ => None,
574                })
575                .unwrap_or_else(|| unreachable_refusal(req))
576        }
577
578        fn elapsed(&self) -> Duration {
579            self.network.clock().elapsed()
580        }
581
582        fn bump_term(&mut self, new_term: u64) {
583            self.bumped_term = Some(new_term);
584        }
585
586        fn promote(&mut self, new_term: u64) {
587            self.promoted_term = Some(new_term);
588        }
589    }
590
591    #[derive(Clone)]
592    struct RcStore(Rc<MemoryLastVoteStore>);
593
594    impl LastVoteStore for RcStore {
595        fn load(&self) -> Result<LastVote, LastVoteError> {
596            self.0.load()
597        }
598
599        fn persist(&self, vote: &LastVote) -> Result<(), LastVoteError> {
600            self.0.persist(vote)
601        }
602    }
603
604    fn unreachable_refusal(req: &VoteRequest) -> VoteDecision {
605        VoteDecision::Refused(RefusalReason::StaleTerm {
606            candidate_term: req.term,
607            voter_term: u64::MAX,
608        })
609    }
610
611    fn five_voters() -> Vec<Member> {
612        vec![
613            Member::data_voting("a"),
614            Member::data_voting("b"),
615            Member::data_voting("c"),
616            Member::data_voting("d"),
617            Member::data_voting("e"),
618        ]
619    }
620
621    fn shared_vote_stores(members: &[Member]) -> BTreeMap<String, Rc<MemoryLastVoteStore>> {
622        members
623            .iter()
624            .map(|member| (member.id.clone(), Rc::new(MemoryLastVoteStore::new())))
625            .collect()
626    }
627
628    fn partitioned_network(seed: u64) -> InProcessReplicationNetwork<ReplicationControlMessage> {
629        let mut network = InProcessReplicationNetwork::new(
630            seed,
631            NetworkFaults {
632                loss_per_million: 0,
633                max_delay_ms: 20,
634                reorder: true,
635            },
636        );
637        for left in ["a", "b", "c"] {
638            for right in ["d", "e"] {
639                network.partition(left, right);
640            }
641        }
642        network
643    }
644
645    fn candidate_request(id: &str, current_term: u64, lsn: u64, watermark: u64) -> ElectionRequest {
646        ElectionRequest {
647            candidate: Member::data_voting(id),
648            current_term,
649            last_log_lsn: lsn,
650            commit_watermark: watermark,
651        }
652    }
653
654    fn lease_store(tag: &str) -> LeaseStore {
655        use crate::storage::backend::LocalBackend;
656
657        LeaseStore::new(Arc::new(LocalBackend)).with_prefix(format!(
658            "{}/reddb-{tag}-{}",
659            std::env::temp_dir().to_string_lossy(),
660            crate::utils::now_unix_nanos(),
661        ))
662    }
663}