use std::collections::HashMap;
use std::time::Duration;
use elara_core::NodeId;
use crate::chaos::{ChaosConfig, ChaosNetwork};
pub struct NetworkLink {
pub from: NodeId,
pub to: NodeId,
pub network: ChaosNetwork,
}
pub struct NetworkSimulator {
links: HashMap<(NodeId, NodeId), ChaosNetwork>,
default_config: ChaosConfig,
current_time: Duration,
seed_counter: u64,
}
impl NetworkSimulator {
pub fn new(default_config: ChaosConfig) -> Self {
NetworkSimulator {
links: HashMap::new(),
default_config,
current_time: Duration::ZERO,
seed_counter: 0,
}
}
pub fn good() -> Self {
Self::new(ChaosConfig::good())
}
pub fn poor() -> Self {
Self::new(ChaosConfig::poor())
}
pub fn hostile() -> Self {
Self::new(ChaosConfig::hostile())
}
fn get_or_create_link(&mut self, from: NodeId, to: NodeId) -> &mut ChaosNetwork {
let seed = self.seed_counter;
self.seed_counter += 1;
self.links
.entry((from, to))
.or_insert_with(|| ChaosNetwork::with_seed(self.default_config.clone(), seed))
}
pub fn send(&mut self, from: NodeId, to: NodeId, data: Vec<u8>) {
let link = self.get_or_create_link(from, to);
link.send(data);
}
pub fn tick(&mut self, dt: Duration) -> Vec<(NodeId, NodeId, Vec<u8>)> {
self.current_time += dt;
let mut delivered = Vec::new();
for ((from, to), link) in &mut self.links {
let packets = link.tick(dt);
for data in packets {
delivered.push((*from, *to, data));
}
}
delivered
}
pub fn current_time(&self) -> Duration {
self.current_time
}
pub fn set_link_config(&mut self, from: NodeId, to: NodeId, config: ChaosConfig) {
let seed = self.seed_counter;
self.seed_counter += 1;
self.links
.insert((from, to), ChaosNetwork::with_seed(config, seed));
}
pub fn link_stats(&self, from: NodeId, to: NodeId) -> Option<&crate::chaos::ChaosStats> {
self.links.get(&(from, to)).map(|l| l.stats())
}
pub fn all_stats(&self) -> Vec<((NodeId, NodeId), &crate::chaos::ChaosStats)> {
self.links.iter().map(|(k, v)| (*k, v.stats())).collect()
}
}
pub struct ScenarioBuilder {
nodes: Vec<NodeId>,
config: ChaosConfig,
duration: Duration,
tick_interval: Duration,
}
impl ScenarioBuilder {
pub fn new() -> Self {
ScenarioBuilder {
nodes: Vec::new(),
config: ChaosConfig::default(),
duration: Duration::from_secs(60),
tick_interval: Duration::from_millis(10),
}
}
pub fn with_nodes(mut self, count: usize) -> Self {
self.nodes = (0..count).map(|i| NodeId::new(i as u64)).collect();
self
}
pub fn with_config(mut self, config: ChaosConfig) -> Self {
self.config = config;
self
}
pub fn with_duration(mut self, duration: Duration) -> Self {
self.duration = duration;
self
}
pub fn with_tick_interval(mut self, interval: Duration) -> Self {
self.tick_interval = interval;
self
}
pub fn build(self) -> (NetworkSimulator, Vec<NodeId>) {
(NetworkSimulator::new(self.config), self.nodes)
}
pub fn duration(&self) -> Duration {
self.duration
}
pub fn tick_interval(&self) -> Duration {
self.tick_interval
}
}
impl Default for ScenarioBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_network_simulator_basic() {
let mut sim = NetworkSimulator::good();
let node1 = NodeId::new(1);
let node2 = NodeId::new(2);
for i in 0..10 {
sim.send(node1, node2, vec![i]);
sim.send(node2, node1, vec![i + 100]);
}
let mut total_delivered = 0;
for _ in 0..100 {
let delivered = sim.tick(Duration::from_millis(10));
total_delivered += delivered.len();
}
assert!(total_delivered >= 18);
}
#[test]
fn test_scenario_builder() {
let (mut sim, nodes) = ScenarioBuilder::new()
.with_nodes(5)
.with_config(ChaosConfig::poor())
.with_duration(Duration::from_secs(10))
.build();
assert_eq!(nodes.len(), 5);
for from in &nodes {
for to in &nodes {
if from != to {
sim.send(*from, *to, vec![1, 2, 3]);
}
}
}
for _ in 0..200 {
sim.tick(Duration::from_millis(10));
}
let stats = sim.all_stats();
assert!(!stats.is_empty());
}
}