317 lines
8.9 KiB
Rust
317 lines
8.9 KiB
Rust
use std::collections::{HashMap, HashSet, VecDeque};
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
/// Identifier for a node in the virtual network.
|
|
pub type NodeId = String;
|
|
|
|
/// A pending message in the virtual network.
|
|
#[derive(Debug, Clone)]
|
|
struct PendingMessage {
|
|
from: NodeId,
|
|
to: NodeId,
|
|
data: Vec<u8>,
|
|
}
|
|
|
|
/// Virtual network for simulation testing.
|
|
/// Supports partition, latency injection, and random packet drop.
|
|
pub struct VirtualNetwork {
|
|
/// Delivered message queues: node_id -> received messages.
|
|
inbox: Arc<Mutex<HashMap<NodeId, VecDeque<(NodeId, Vec<u8>)>>>>,
|
|
/// Pending messages not yet delivered (used for latency simulation).
|
|
pending: Arc<Mutex<VecDeque<PendingMessage>>>,
|
|
/// Partitioned links: (from, to) pairs that are blocked.
|
|
partitions: Arc<Mutex<HashSet<(NodeId, NodeId)>>>,
|
|
/// Drop probability (0.0 to 1.0).
|
|
drop_probability: Arc<Mutex<f64>>,
|
|
}
|
|
|
|
impl VirtualNetwork {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
inbox: Arc::new(Mutex::new(HashMap::new())),
|
|
pending: Arc::new(Mutex::new(VecDeque::new())),
|
|
partitions: Arc::new(Mutex::new(HashSet::new())),
|
|
drop_probability: Arc::new(Mutex::new(0.0)),
|
|
}
|
|
}
|
|
|
|
/// Partition the network between two nodes (bidirectional).
|
|
pub fn partition(&self, a: &str, b: &str) {
|
|
let mut parts = self.partitions.lock().unwrap();
|
|
parts.insert((a.to_string(), b.to_string()));
|
|
parts.insert((b.to_string(), a.to_string()));
|
|
}
|
|
|
|
/// Heal the partition between two nodes (bidirectional).
|
|
pub fn heal(&self, a: &str, b: &str) {
|
|
let mut parts = self.partitions.lock().unwrap();
|
|
parts.remove(&(a.to_string(), b.to_string()));
|
|
parts.remove(&(b.to_string(), a.to_string()));
|
|
}
|
|
|
|
/// Heal all partitions.
|
|
pub fn heal_all(&self) {
|
|
self.partitions.lock().unwrap().clear();
|
|
}
|
|
|
|
/// Set the probability that a message will be dropped (0.0 = no drops, 1.0 = all dropped).
|
|
pub fn set_drop_probability(&self, prob: f64) {
|
|
*self.drop_probability.lock().unwrap() = prob.clamp(0.0, 1.0);
|
|
}
|
|
|
|
/// Send a message from one node to another.
|
|
/// If the link is partitioned, the message is silently dropped.
|
|
pub fn send(&self, from: &str, to: &str, data: Vec<u8>) -> Result<(), NetworkError> {
|
|
// Check for partition.
|
|
{
|
|
let parts = self.partitions.lock().unwrap();
|
|
if parts.contains(&(from.to_string(), to.to_string())) {
|
|
return Ok(()); // Silently dropped.
|
|
}
|
|
}
|
|
|
|
// Check for random drop.
|
|
{
|
|
let prob = *self.drop_probability.lock().unwrap();
|
|
if prob > 0.0 {
|
|
let random: f64 = simple_random();
|
|
if random < prob {
|
|
return Ok(()); // Randomly dropped.
|
|
}
|
|
}
|
|
}
|
|
|
|
// Queue the message for delivery.
|
|
let mut pending = self.pending.lock().unwrap();
|
|
pending.push_back(PendingMessage {
|
|
from: from.to_string(),
|
|
to: to.to_string(),
|
|
data,
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Deliver all pending messages to their inboxes.
|
|
/// Call this to simulate message delivery (allows controlling when messages arrive).
|
|
pub fn deliver_pending(&self) {
|
|
let messages: Vec<PendingMessage> = {
|
|
let mut pending = self.pending.lock().unwrap();
|
|
pending.drain(..).collect()
|
|
};
|
|
|
|
let mut inbox = self.inbox.lock().unwrap();
|
|
for msg in messages {
|
|
inbox
|
|
.entry(msg.to.clone())
|
|
.or_default()
|
|
.push_back((msg.from, msg.data));
|
|
}
|
|
}
|
|
|
|
/// Receive a message for a given node. Returns None if no messages are available.
|
|
pub fn recv(&self, node: &str) -> Option<(NodeId, Vec<u8>)> {
|
|
let mut inbox = self.inbox.lock().unwrap();
|
|
inbox.get_mut(node).and_then(|q| q.pop_front())
|
|
}
|
|
|
|
/// Get the number of pending (undelivered) messages.
|
|
pub fn pending_count(&self) -> usize {
|
|
self.pending.lock().unwrap().len()
|
|
}
|
|
|
|
/// Get the number of messages in a node's inbox.
|
|
pub fn inbox_count(&self, node: &str) -> usize {
|
|
self.inbox
|
|
.lock()
|
|
.unwrap()
|
|
.get(node)
|
|
.map(|q| q.len())
|
|
.unwrap_or(0)
|
|
}
|
|
}
|
|
|
|
impl Default for VirtualNetwork {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Simple deterministic pseudo-random based on thread-local state.
|
|
fn simple_random() -> f64 {
|
|
use std::cell::Cell;
|
|
thread_local! {
|
|
static STATE: Cell<u64> = const { Cell::new(12345) };
|
|
}
|
|
STATE.with(|s| {
|
|
let mut state = s.get();
|
|
state ^= state << 13;
|
|
state ^= state >> 7;
|
|
state ^= state << 17;
|
|
s.set(state);
|
|
(state % 10000) as f64 / 10000.0
|
|
})
|
|
}
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum NetworkError {
|
|
#[error("node '{0}' not reachable")]
|
|
Unreachable(String),
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_send_and_receive() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.send("node-1", "node-2", b"hello".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
let (from, data) = net.recv("node-2").unwrap();
|
|
assert_eq!(from, "node-1");
|
|
assert_eq!(data, b"hello");
|
|
}
|
|
|
|
#[test]
|
|
fn test_no_messages_returns_none() {
|
|
let net = VirtualNetwork::new();
|
|
assert!(net.recv("node-1").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_partition_drops_messages() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.partition("node-1", "node-2");
|
|
|
|
net.send("node-1", "node-2", b"hello".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
assert!(net.recv("node-2").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_partition_is_bidirectional() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.partition("node-1", "node-2");
|
|
|
|
net.send("node-1", "node-2", b"a->b".to_vec()).unwrap();
|
|
net.send("node-2", "node-1", b"b->a".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
assert!(net.recv("node-2").is_none());
|
|
assert!(net.recv("node-1").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_heal_restores_communication() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.partition("node-1", "node-2");
|
|
net.send("node-1", "node-2", b"before".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
assert!(net.recv("node-2").is_none());
|
|
|
|
net.heal("node-1", "node-2");
|
|
net.send("node-1", "node-2", b"after".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
let (_, data) = net.recv("node-2").unwrap();
|
|
assert_eq!(data, b"after");
|
|
}
|
|
|
|
#[test]
|
|
fn test_heal_all() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.partition("a", "b");
|
|
net.partition("a", "c");
|
|
net.heal_all();
|
|
|
|
net.send("a", "b", b"msg".to_vec()).unwrap();
|
|
net.send("a", "c", b"msg".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
assert!(net.recv("b").is_some());
|
|
assert!(net.recv("c").is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_messages_ordered() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
for i in 0..5 {
|
|
net.send("a", "b", format!("msg-{i}").into_bytes())
|
|
.unwrap();
|
|
}
|
|
net.deliver_pending();
|
|
|
|
for i in 0..5 {
|
|
let (_, data) = net.recv("b").unwrap();
|
|
assert_eq!(data, format!("msg-{i}").as_bytes());
|
|
}
|
|
assert!(net.recv("b").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn test_pending_and_inbox_counts() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.send("a", "b", b"1".to_vec()).unwrap();
|
|
net.send("a", "b", b"2".to_vec()).unwrap();
|
|
|
|
assert_eq!(net.pending_count(), 2);
|
|
assert_eq!(net.inbox_count("b"), 0);
|
|
|
|
net.deliver_pending();
|
|
|
|
assert_eq!(net.pending_count(), 0);
|
|
assert_eq!(net.inbox_count("b"), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_partition_does_not_affect_other_links() {
|
|
let net = VirtualNetwork::new();
|
|
|
|
net.partition("a", "b");
|
|
|
|
// a -> c should still work.
|
|
net.send("a", "c", b"hello".to_vec()).unwrap();
|
|
net.deliver_pending();
|
|
|
|
assert!(net.recv("c").is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn test_drop_probability_all() {
|
|
let net = VirtualNetwork::new();
|
|
net.set_drop_probability(1.0);
|
|
|
|
for _ in 0..10 {
|
|
net.send("a", "b", b"msg".to_vec()).unwrap();
|
|
}
|
|
net.deliver_pending();
|
|
|
|
// All messages should be dropped.
|
|
assert_eq!(net.inbox_count("b"), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_drop_probability_none() {
|
|
let net = VirtualNetwork::new();
|
|
net.set_drop_probability(0.0);
|
|
|
|
for _ in 0..10 {
|
|
net.send("a", "b", b"msg".to_vec()).unwrap();
|
|
}
|
|
net.deliver_pending();
|
|
|
|
// No messages should be dropped.
|
|
assert_eq!(net.inbox_count("b"), 10);
|
|
}
|
|
}
|