use std::collections::{HashMap, VecDeque};
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct FaultSeed(pub u64);
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct MessageId(pub u64);
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct NodeId(pub u64);
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ConnectionId {
pub from: NodeId,
pub to: NodeId,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LinkState {
Connected,
Partitioned,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum NetworkEvent {
Delay { message_id: MessageId, ticks: u64 },
Drop { message_id: MessageId },
Duplicate { message_id: MessageId },
Partition { from: NodeId, to: NodeId },
Heal { from: NodeId, to: NodeId },
Reorder { connection: ConnectionId },
}
#[derive(Clone, Debug)]
pub struct NetworkSchedule {
pub seed: FaultSeed,
pub events: Vec<NetworkEvent>,
}
impl NetworkSchedule {
pub fn new(seed: FaultSeed) -> Self {
Self {
seed,
events: Vec::new(),
}
}
pub fn add_event(&mut self, event: NetworkEvent) {
self.events.push(event);
}
pub fn get_seed(&self) -> FaultSeed {
self.seed
}
}
#[derive(Clone, Debug)]
struct InFlightMessage {
id: MessageId,
from: NodeId,
to: NodeId,
payload: Vec<u8>,
delivery_tick: u64,
}
#[derive(Clone, Debug)]
pub struct ConnectivityMatrix {
links: HashMap<ConnectionId, LinkState>,
}
impl ConnectivityMatrix {
pub fn new() -> Self {
Self {
links: HashMap::new(),
}
}
pub fn can_reach(&self, from: NodeId, to: NodeId) -> bool {
let conn = ConnectionId { from, to };
matches!(self.links.get(&conn), Some(LinkState::Connected) | None)
}
pub fn partition(&mut self, from: NodeId, to: NodeId) {
self.links
.insert(ConnectionId { from, to }, LinkState::Partitioned);
}
pub fn heal(&mut self, from: NodeId, to: NodeId) {
self.links
.insert(ConnectionId { from, to }, LinkState::Connected);
}
pub fn get_state(&self, from: NodeId, to: NodeId) -> LinkState {
let conn = ConnectionId { from, to };
*self.links.get(&conn).unwrap_or(&LinkState::Connected)
}
}
impl Default for ConnectivityMatrix {
fn default() -> Self {
Self::new()
}
}
pub trait ReplicationTransport: Send + Sync {
fn send(&mut self, from: NodeId, to: NodeId, message: Vec<u8>) -> MessageId;
fn tick(&mut self);
fn receive(&mut self, node: NodeId) -> Vec<(NodeId, Vec<u8>)>;
fn apply_fault(&mut self, event: NetworkEvent);
fn current_tick(&self) -> u64;
fn get_connectivity(&self, from: NodeId, to: NodeId) -> LinkState;
}
pub struct InMemoryTransport {
current_tick: u64,
message_counter: u64,
queues: HashMap<NodeId, VecDeque<(NodeId, Vec<u8>)>>,
}
impl InMemoryTransport {
pub fn new() -> Self {
Self {
current_tick: 0,
message_counter: 0,
queues: HashMap::new(),
}
}
}
impl Default for InMemoryTransport {
fn default() -> Self {
Self::new()
}
}
impl ReplicationTransport for InMemoryTransport {
fn send(&mut self, from: NodeId, to: NodeId, message: Vec<u8>) -> MessageId {
let id = MessageId(self.message_counter);
self.message_counter += 1;
self.queues
.entry(to)
.or_insert_with(VecDeque::new)
.push_back((from, message));
id
}
fn tick(&mut self) {
self.current_tick += 1;
}
fn receive(&mut self, node: NodeId) -> Vec<(NodeId, Vec<u8>)> {
self.queues
.remove(&node)
.map(|q| q.into_iter().collect())
.unwrap_or_default()
}
fn apply_fault(&mut self, _event: NetworkEvent) {
}
fn current_tick(&self) -> u64 {
self.current_tick
}
fn get_connectivity(&self, _from: NodeId, _to: NodeId) -> LinkState {
LinkState::Connected
}
}
pub struct AdversarialTransport {
current_tick: u64,
message_counter: u64,
in_flight: Vec<InFlightMessage>,
delivered_queues: HashMap<NodeId, VecDeque<(NodeId, Vec<u8>)>>,
connectivity: ConnectivityMatrix,
dropped_messages: std::collections::HashSet<MessageId>,
}
impl AdversarialTransport {
pub fn new() -> Self {
Self {
current_tick: 0,
message_counter: 0,
in_flight: Vec::new(),
delivered_queues: HashMap::new(),
connectivity: ConnectivityMatrix::new(),
dropped_messages: std::collections::HashSet::new(),
}
}
pub fn get_in_flight_count(&self) -> usize {
self.in_flight.len()
}
pub fn get_delivered_count(&self, node: NodeId) -> usize {
self.delivered_queues
.get(&node)
.map(|q| q.len())
.unwrap_or(0)
}
}
impl Default for AdversarialTransport {
fn default() -> Self {
Self::new()
}
}
impl ReplicationTransport for AdversarialTransport {
fn send(&mut self, from: NodeId, to: NodeId, message: Vec<u8>) -> MessageId {
let id = MessageId(self.message_counter);
self.message_counter += 1;
if !self.connectivity.can_reach(from, to) {
self.dropped_messages.insert(id);
return id;
}
self.in_flight.push(InFlightMessage {
id,
from,
to,
payload: message,
delivery_tick: self.current_tick,
});
id
}
fn tick(&mut self) {
self.current_tick += 1;
let mut to_deliver = Vec::new();
let mut remaining = Vec::new();
for msg in self.in_flight.drain(..) {
if msg.delivery_tick <= self.current_tick {
to_deliver.push(msg);
} else {
remaining.push(msg);
}
}
self.in_flight = remaining;
for msg in to_deliver {
if !self.dropped_messages.contains(&msg.id) {
self.delivered_queues
.entry(msg.to)
.or_insert_with(VecDeque::new)
.push_back((msg.from, msg.payload));
}
}
}
fn receive(&mut self, node: NodeId) -> Vec<(NodeId, Vec<u8>)> {
self.delivered_queues
.remove(&node)
.map(|q| q.into_iter().collect())
.unwrap_or_default()
}
fn apply_fault(&mut self, event: NetworkEvent) {
match event {
NetworkEvent::Delay { message_id, ticks } => {
for msg in &mut self.in_flight {
if msg.id == message_id {
msg.delivery_tick += ticks;
break;
}
}
}
NetworkEvent::Drop { message_id } => {
self.dropped_messages.insert(message_id);
}
NetworkEvent::Duplicate { message_id } => {
if let Some(orig) = self.in_flight.iter().find(|m| m.id == message_id) {
let dup = InFlightMessage {
id: MessageId(self.message_counter),
from: orig.from,
to: orig.to,
payload: orig.payload.clone(),
delivery_tick: orig.delivery_tick,
};
self.message_counter += 1;
self.in_flight.push(dup);
}
}
NetworkEvent::Partition { from, to } => {
self.connectivity.partition(from, to);
}
NetworkEvent::Heal { from, to } => {
self.connectivity.heal(from, to);
}
NetworkEvent::Reorder { connection } => {
let mut msgs_on_conn: Vec<_> = self
.in_flight
.iter()
.enumerate()
.filter(|(_, m)| m.from == connection.from && m.to == connection.to)
.collect();
if msgs_on_conn.len() > 1 {
msgs_on_conn.reverse();
let indices: Vec<_> = msgs_on_conn.iter().map(|(i, _)| *i).collect();
for (new_pos, &old_pos) in indices.iter().enumerate() {
if new_pos != old_pos {
self.in_flight.swap(new_pos, old_pos);
}
}
}
}
}
}
fn current_tick(&self) -> u64 {
self.current_tick
}
fn get_connectivity(&self, from: NodeId, to: NodeId) -> LinkState {
self.connectivity.get_state(from, to)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_in_memory_transport_basic() {
let mut transport = InMemoryTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
transport.send(n1, n2, vec![1, 2, 3]);
let messages = transport.receive(n2);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].0, n1);
assert_eq!(messages[0].1, vec![1, 2, 3]);
}
#[test]
fn test_adversarial_transport_basic() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
let msg_id = transport.send(n1, n2, vec![1, 2, 3]);
assert!(msg_id.0 >= 0);
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].0, n1);
}
#[test]
fn test_partition() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
transport.apply_fault(NetworkEvent::Partition { from: n1, to: n2 });
let _msg_id = transport.send(n1, n2, vec![1, 2, 3]);
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 0);
}
#[test]
fn test_delay() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
let msg_id = transport.send(n1, n2, vec![1, 2, 3]);
transport.apply_fault(NetworkEvent::Delay {
message_id: msg_id,
ticks: 5,
});
for _ in 0..4 {
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 0);
}
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 1);
}
#[test]
fn test_duplicate() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
let msg_id = transport.send(n1, n2, vec![1, 2, 3]);
transport.apply_fault(NetworkEvent::Duplicate { message_id: msg_id });
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].1, vec![1, 2, 3]);
assert_eq!(messages[1].1, vec![1, 2, 3]);
}
#[test]
fn test_drop() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
let msg_id = transport.send(n1, n2, vec![1, 2, 3]);
transport.apply_fault(NetworkEvent::Drop { message_id: msg_id });
transport.tick();
let messages = transport.receive(n2);
assert_eq!(messages.len(), 0);
}
#[test]
fn test_asymmetric_partition() {
let mut transport = AdversarialTransport::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
transport.apply_fault(NetworkEvent::Partition { from: n1, to: n2 });
let _msg_id = transport.send(n1, n2, vec![1, 2, 3]);
transport.tick();
assert_eq!(transport.receive(n2).len(), 0);
let msg_id = transport.send(n2, n1, vec![4, 5, 6]);
assert!(msg_id.0 >= 0);
transport.tick();
assert_eq!(transport.receive(n1).len(), 1);
}
#[test]
fn test_connectivity_matrix() {
let mut connectivity = ConnectivityMatrix::new();
let n1 = NodeId(1);
let n2 = NodeId(2);
assert!(connectivity.can_reach(n1, n2));
connectivity.partition(n1, n2);
assert!(!connectivity.can_reach(n1, n2));
connectivity.heal(n1, n2);
assert!(connectivity.can_reach(n1, n2));
}
}