1use crate::clock::Micros;
14use crate::rng::SimRng;
15use serde::{Deserialize, Serialize};
16use std::collections::{BTreeMap, BTreeSet, VecDeque};
17use std::fmt;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
21pub struct NodeId(pub u16);
22
23impl fmt::Display for NodeId {
24 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25 write!(f, "n{}", self.0)
26 }
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
31pub enum DropReason {
32 Partition,
34 LinkFault,
36 NodeCrashed,
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub struct LinkConfig {
44 pub min_latency: Micros,
46 pub max_latency: Micros,
49 pub drop_per_mille: u32,
51 pub duplicate_per_mille: u32,
53}
54
55impl LinkConfig {
56 pub fn new(min_latency: Micros, max_latency: Micros) -> Self {
58 assert!(
59 min_latency <= max_latency,
60 "min_latency must be <= max_latency"
61 );
62 Self {
63 min_latency,
64 max_latency,
65 drop_per_mille: 0,
66 duplicate_per_mille: 0,
67 }
68 }
69
70 pub fn with_drop_per_mille(mut self, per_mille: u32) -> Self {
72 self.drop_per_mille = per_mille;
73 self
74 }
75
76 pub fn with_duplicate_per_mille(mut self, per_mille: u32) -> Self {
78 self.duplicate_per_mille = per_mille;
79 self
80 }
81}
82
83impl Default for LinkConfig {
84 fn default() -> Self {
85 Self::new(1, 1)
86 }
87}
88
89#[derive(Debug, Clone, PartialEq, Eq)]
91pub struct Message {
92 pub from: NodeId,
94 pub to: NodeId,
96 pub seq: u64,
98 pub payload: Vec<u8>,
100 pub sent_at: Micros,
102 pub deliver_at: Micros,
104}
105
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum SendOutcome {
109 Scheduled {
111 seq: u64,
113 deliver_at: Micros,
115 duplicated: bool,
117 reordered: bool,
119 },
120 Dropped {
122 seq: u64,
124 reason: DropReason,
126 },
127}
128
129#[derive(Debug, Clone, Copy, PartialEq, Eq)]
131pub enum DeliveryOutcome {
132 Delivered,
134 Discarded(DropReason),
136}
137
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
140pub struct Delivery {
141 pub from: NodeId,
143 pub to: NodeId,
145 pub seq: u64,
147 pub outcome: DeliveryOutcome,
149}
150
151#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
153pub struct NetworkStats {
154 pub sent: u64,
156 pub delivered: u64,
158 pub dropped: u64,
160 pub duplicated: u64,
162 pub reordered: u64,
164}
165
166#[derive(Debug)]
169pub struct Network {
170 default_link: LinkConfig,
171 links: BTreeMap<(NodeId, NodeId), LinkConfig>,
172 cuts: BTreeSet<(NodeId, NodeId)>,
173 in_flight: BTreeMap<(Micros, u64), Message>,
174 inboxes: BTreeMap<NodeId, VecDeque<Message>>,
175 last_deliver_at: BTreeMap<(NodeId, NodeId), Micros>,
176 next_seq: u64,
177 stats: NetworkStats,
178}
179
180impl Network {
181 pub fn new(default_link: LinkConfig) -> Self {
184 Self {
185 default_link,
186 links: BTreeMap::new(),
187 cuts: BTreeSet::new(),
188 in_flight: BTreeMap::new(),
189 inboxes: BTreeMap::new(),
190 last_deliver_at: BTreeMap::new(),
191 next_seq: 1,
192 stats: NetworkStats::default(),
193 }
194 }
195
196 pub fn set_link_config(&mut self, from: NodeId, to: NodeId, config: LinkConfig) {
198 assert!(
199 config.min_latency <= config.max_latency,
200 "min_latency must be <= max_latency"
201 );
202 assert!(
203 config.drop_per_mille <= 1000 && config.duplicate_per_mille <= 1000,
204 "per-mille rates must be <= 1000"
205 );
206 self.links.insert((from, to), config);
207 }
208
209 pub fn partition(
212 &mut self,
213 group_a: impl IntoIterator<Item = NodeId>,
214 group_b: impl IntoIterator<Item = NodeId>,
215 ) {
216 let group_b: Vec<NodeId> = group_b.into_iter().collect();
217 for a in group_a {
218 for &b in &group_b {
219 if a != b {
220 self.cuts.insert(normalize(a, b));
221 }
222 }
223 }
224 }
225
226 pub fn heal(&mut self) {
228 self.cuts.clear();
229 }
230
231 pub fn connected(&self, a: NodeId, b: NodeId) -> bool {
233 a == b || !self.cuts.contains(&normalize(a, b))
234 }
235
236 pub fn send(
239 &mut self,
240 from: NodeId,
241 to: NodeId,
242 payload: Vec<u8>,
243 now: Micros,
244 rng: &mut SimRng,
245 ) -> SendOutcome {
246 let seq = self.next_seq;
247 self.next_seq += 1;
248 self.stats.sent += 1;
249
250 if !self.connected(from, to) {
251 self.stats.dropped += 1;
252 return SendOutcome::Dropped {
253 seq,
254 reason: DropReason::Partition,
255 };
256 }
257
258 let link = self.link_config(from, to);
259 if rng.chance(link.drop_per_mille) {
260 self.stats.dropped += 1;
261 return SendOutcome::Dropped {
262 seq,
263 reason: DropReason::LinkFault,
264 };
265 }
266
267 let latency = if link.max_latency > link.min_latency {
268 rng.range(link.min_latency, link.max_latency + 1)
269 } else {
270 link.min_latency
271 };
272 let deliver_at = now + latency;
273
274 let key = (from, to);
275 let previous = self.last_deliver_at.get(&key).copied();
276 let reordered = previous.is_some_and(|earliest_max| deliver_at < earliest_max);
277 if reordered {
278 self.stats.reordered += 1;
279 }
280 self.last_deliver_at
281 .insert(key, previous.map_or(deliver_at, |p| p.max(deliver_at)));
282
283 let message = Message {
284 from,
285 to,
286 seq,
287 payload,
288 sent_at: now,
289 deliver_at,
290 };
291 let mut duplicated = false;
292 if rng.chance(link.duplicate_per_mille) {
293 duplicated = true;
294 self.stats.duplicated += 1;
295 let dup_seq = self.next_seq;
296 self.next_seq += 1;
297 let copy = Message {
298 seq: dup_seq,
299 deliver_at: deliver_at + 1,
300 ..message.clone()
301 };
302 self.in_flight.insert((deliver_at + 1, dup_seq), copy);
303 }
304 self.in_flight.insert((deliver_at, seq), message);
305
306 SendOutcome::Scheduled {
307 seq,
308 deliver_at,
309 duplicated,
310 reordered,
311 }
312 }
313
314 pub fn deliver_due(&mut self, now: Micros, crashed: &BTreeSet<NodeId>) -> Vec<Delivery> {
317 let later = self.in_flight.split_off(&(now.saturating_add(1), 0));
318 let due = std::mem::replace(&mut self.in_flight, later);
319 let mut deliveries = Vec::new();
320 for (_, message) in due {
321 let delivery = if crashed.contains(&message.to) {
322 self.stats.dropped += 1;
323 Delivery {
324 from: message.from,
325 to: message.to,
326 seq: message.seq,
327 outcome: DeliveryOutcome::Discarded(DropReason::NodeCrashed),
328 }
329 } else {
330 self.stats.delivered += 1;
331 self.inboxes
332 .entry(message.to)
333 .or_default()
334 .push_back(message.clone());
335 Delivery {
336 from: message.from,
337 to: message.to,
338 seq: message.seq,
339 outcome: DeliveryOutcome::Delivered,
340 }
341 };
342 deliveries.push(delivery);
343 }
344 deliveries
345 }
346
347 pub fn next_delivery(&self) -> Option<Micros> {
349 self.in_flight.first_key_value().map(|(&(at, _), _)| at)
350 }
351
352 pub fn try_recv(&mut self, node: NodeId) -> Option<Message> {
354 self.inboxes.get_mut(&node).and_then(VecDeque::pop_front)
355 }
356
357 pub fn inbox_len(&self, node: NodeId) -> usize {
359 self.inboxes.get(&node).map_or(0, VecDeque::len)
360 }
361
362 pub fn clear_inbox(&mut self, node: NodeId) {
364 self.inboxes.remove(&node);
365 }
366
367 pub fn stats(&self) -> NetworkStats {
369 self.stats
370 }
371
372 fn link_config(&self, from: NodeId, to: NodeId) -> LinkConfig {
373 self.links
374 .get(&(from, to))
375 .copied()
376 .unwrap_or(self.default_link)
377 }
378}
379
380fn normalize(a: NodeId, b: NodeId) -> (NodeId, NodeId) {
381 if a <= b {
382 (a, b)
383 } else {
384 (b, a)
385 }
386}
387
388#[cfg(test)]
389mod tests {
390 use super::*;
391 use crate::rng::Seed;
392
393 const A: NodeId = NodeId(1);
394 const B: NodeId = NodeId(2);
395
396 fn no_crash() -> BTreeSet<NodeId> {
397 BTreeSet::new()
398 }
399
400 #[test]
401 fn messages_deliver_in_latency_order() {
402 let mut net = Network::new(LinkConfig::default());
403 net.set_link_config(A, B, LinkConfig::new(10, 10));
404 let mut rng = SimRng::from_seed(Seed::new(1));
405 net.send(A, B, b"x".to_vec(), 0, &mut rng);
406 net.send(A, B, b"y".to_vec(), 5, &mut rng);
407
408 assert_eq!(net.next_delivery(), Some(10));
409 assert!(net.deliver_due(9, &no_crash()).is_empty());
410 let due = net.deliver_due(10, &no_crash());
411 assert_eq!(due.len(), 1);
412 assert_eq!(net.try_recv(B).unwrap().payload, b"x");
413 assert!(net.try_recv(B).is_none());
414
415 let due = net.deliver_due(15, &no_crash());
416 assert_eq!(due.len(), 1);
417 assert_eq!(net.try_recv(B).unwrap().payload, b"y");
418 }
419
420 #[test]
421 fn full_drop_rate_drops_everything() {
422 let mut net = Network::new(LinkConfig::default());
423 net.set_link_config(A, B, LinkConfig::new(1, 1).with_drop_per_mille(1000));
424 let mut rng = SimRng::from_seed(Seed::new(2));
425 for i in 0..10u8 {
426 let outcome = net.send(A, B, vec![i], 0, &mut rng);
427 assert!(matches!(
428 outcome,
429 SendOutcome::Dropped {
430 reason: DropReason::LinkFault,
431 ..
432 }
433 ));
434 }
435 assert_eq!(net.stats().sent, 10);
436 assert_eq!(net.stats().dropped, 10);
437 assert_eq!(net.stats().delivered, 0);
438 }
439
440 #[test]
441 fn full_duplicate_rate_doubles_delivery() {
442 let mut net = Network::new(LinkConfig::default());
443 net.set_link_config(A, B, LinkConfig::new(1, 1).with_duplicate_per_mille(1000));
444 let mut rng = SimRng::from_seed(Seed::new(3));
445 for i in 0..5u8 {
446 let outcome = net.send(A, B, vec![i], 0, &mut rng);
447 assert!(matches!(
448 outcome,
449 SendOutcome::Scheduled {
450 duplicated: true,
451 ..
452 }
453 ));
454 }
455 assert_eq!(net.deliver_due(100, &no_crash()).len(), 10);
456 assert_eq!(net.stats().duplicated, 5);
457 assert_eq!(net.stats().delivered, 10);
458 let mut drained = 0;
459 while net.try_recv(B).is_some() {
460 drained += 1;
461 }
462 assert_eq!(drained, 10);
463 }
464
465 #[test]
466 fn partition_cuts_both_directions_until_heal() {
467 let mut net = Network::new(LinkConfig::default());
468 let mut rng = SimRng::from_seed(Seed::new(4));
469 net.send(A, B, b"before".to_vec(), 0, &mut rng);
470
471 net.partition([A], [B]);
472 assert!(!net.connected(A, B));
473 assert!(!net.connected(B, A));
474 for (from, to) in [(A, B), (B, A)] {
475 let outcome = net.send(from, to, b"x".to_vec(), 1, &mut rng);
476 assert!(matches!(
477 outcome,
478 SendOutcome::Dropped {
479 reason: DropReason::Partition,
480 ..
481 }
482 ));
483 }
484 assert_eq!(net.deliver_due(100, &no_crash()).len(), 1);
486 assert_eq!(net.try_recv(B).unwrap().payload, b"before");
487
488 net.heal();
489 assert!(net.connected(A, B));
490 assert!(matches!(
491 net.send(A, B, b"after".to_vec(), 2, &mut rng),
492 SendOutcome::Scheduled { .. }
493 ));
494 }
495
496 #[test]
497 fn delivery_to_crashed_node_is_discarded() {
498 let mut net = Network::new(LinkConfig::default());
499 let mut rng = SimRng::from_seed(Seed::new(5));
500 net.send(A, B, b"x".to_vec(), 0, &mut rng);
501
502 let crashed = BTreeSet::from([B]);
503 let due = net.deliver_due(100, &crashed);
504 assert_eq!(
505 due,
506 vec![Delivery {
507 from: A,
508 to: B,
509 seq: 1,
510 outcome: DeliveryOutcome::Discarded(DropReason::NodeCrashed),
511 }]
512 );
513 assert_eq!(net.inbox_len(B), 0);
514 assert_eq!(net.stats().dropped, 1);
515 assert_eq!(net.stats().delivered, 0);
516 }
517}