316
crates/sq-sim/src/network.rs
Normal file
316
crates/sq-sim/src/network.rs
Normal file
@@ -0,0 +1,316 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user