1use 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}