use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use std::time::Duration;
use crate::entropy::{Entropy, SeededEntropy};
pub(crate) const NET_STREAM_SALT: u64 = 0x4E37_5EED_4E37_5EED;
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NetFault {
None,
Dropped,
Partitioned,
TimedOut,
}
impl std::fmt::Display for NetFault {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Self::None => "delivered",
Self::Dropped => "dropped",
Self::Partitioned => "partitioned",
Self::TimedOut => "timed out",
})
}
}
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NetEvent {
pub seq: u64,
pub host: String,
pub latency: Duration,
pub fault: NetFault,
}
#[derive(Clone, Default)]
pub struct SimNet {
inner: Arc<Mutex<NetState>>,
}
struct NetState {
hosts: BTreeMap<String, axum::Router>,
partitioned: BTreeSet<String>,
drop_rate: f64,
latency: (Duration, Duration),
stream: Arc<dyn Entropy>,
events: Vec<NetEvent>,
}
impl Default for NetState {
fn default() -> Self {
Self {
hosts: BTreeMap::new(),
partitioned: BTreeSet::new(),
drop_rate: 0.0,
latency: (Duration::ZERO, Duration::ZERO),
stream: SeededEntropy::shared(NET_STREAM_SALT),
events: Vec::new(),
}
}
}
impl std::fmt::Debug for SimNet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let state = self.lock();
f.debug_struct("SimNet")
.field("hosts", &state.hosts.keys().collect::<Vec<_>>())
.field("partitioned", &state.partitioned)
.field("drop_rate", &state.drop_rate)
.field("latency", &state.latency)
.finish_non_exhaustive()
}
}
impl SimNet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn host(self, host: impl Into<String>, router: axum::Router) -> Self {
self.lock()
.hosts
.insert(canonical_host(&host.into()), router);
self
}
#[must_use]
pub fn latency(self, min: Duration, max: Duration) -> Self {
self.lock().latency = (min, max.max(min));
self
}
#[must_use]
pub fn drop_rate(self, rate: f64) -> Self {
self.lock().drop_rate = super::chaos::clamp_prob(rate);
self
}
pub fn partition(&self, host: impl Into<String>) {
self.lock().partitioned.insert(canonical_host(&host.into()));
}
pub fn heal(&self, host: &str) {
self.lock().partitioned.remove(&canonical_host(host));
}
#[must_use]
pub fn events(&self) -> Vec<NetEvent> {
self.lock().events.clone()
}
pub(crate) fn reseed(&self, seed: u64) {
self.lock().stream = SeededEntropy::shared(seed ^ NET_STREAM_SALT);
}
pub(crate) fn insert_host(&self, host: &str, router: axum::Router) {
self.lock().hosts.insert(canonical_host(host), router);
}
pub(crate) fn remove_host(&self, host: &str) {
self.lock().hosts.remove(&canonical_host(host));
}
pub(crate) fn service(&self, host: &str) -> Option<axum::Router> {
self.lock().hosts.get(&canonical_host(host)).cloned()
}
pub(crate) async fn transmit(
&self,
host: &str,
timeout: Option<Duration>,
) -> Result<(), NetFault> {
let (latency, fault) = {
let state = self.lock();
let latency_draw = state.stream.next_u64();
let drop_draw = state.stream.next_u64();
let (latency, fault) = if state.partitioned.contains(&canonical_host(host)) {
(Duration::ZERO, NetFault::Partitioned)
} else if super::chaos::unit_from_draw(drop_draw) < state.drop_rate {
(
sample_latency(state.latency, latency_draw),
NetFault::Dropped,
)
} else {
(sample_latency(state.latency, latency_draw), NetFault::None)
};
let (latency, fault) = match timeout {
Some(limit) if fault != NetFault::Partitioned && latency > limit => {
(limit, NetFault::TimedOut)
}
_ => (latency, fault),
};
drop(state);
(latency, fault)
};
if !latency.is_zero() {
tokio::time::sleep(latency).await;
}
{
let mut state = self.lock();
let seq = state.events.len() as u64;
state.events.push(NetEvent {
seq,
host: host.to_owned(),
latency,
fault,
});
}
match fault {
NetFault::None => Ok(()),
fault => Err(fault),
}
}
fn lock(&self) -> MutexGuard<'_, NetState> {
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
}
fn sample_latency((min, max): (Duration, Duration), draw: u64) -> Duration {
let min_ns = u64::try_from(min.as_nanos()).unwrap_or(u64::MAX);
let max_ns = u64::try_from(max.as_nanos()).unwrap_or(u64::MAX);
let span = max_ns.saturating_sub(min_ns).saturating_add(1);
Duration::from_nanos(min_ns.saturating_add(draw % span))
}
pub(crate) fn canonical_host(host: &str) -> String {
url::Host::parse(host).map_or_else(|_| host.to_ascii_lowercase(), |parsed| parsed.to_string())
}
#[cfg(test)]
mod tests {
#[test]
fn a_host_name_matches_as_a_url_host_does() {
use super::canonical_host;
assert_eq!(canonical_host("Payments"), "payments");
assert_eq!(
canonical_host("PAYMENTS.Example.COM"),
"payments.example.com"
);
assert_eq!(canonical_host("[::1]"), "[::1]");
assert_eq!(canonical_host("10.0.0.1"), "10.0.0.1");
let url = url::Url::parse("http://Bücher.Example/").unwrap();
assert_eq!(canonical_host("Bücher.Example"), url.host_str().unwrap());
}
use super::{Duration, NetFault, SimNet, sample_latency};
#[test]
fn latency_stays_in_range() {
let range = (Duration::from_millis(5), Duration::from_millis(9));
for draw in 0..100 {
let latency = sample_latency(range, draw);
assert!(latency >= range.0 && latency <= range.1, "{latency:?}");
}
let fixed = (Duration::from_millis(7), Duration::from_millis(7));
assert_eq!(sample_latency(fixed, u64::MAX), Duration::from_millis(7));
}
#[test]
fn sub_millisecond_latency_is_kept() {
let fixed = (Duration::from_micros(500), Duration::from_micros(500));
assert_eq!(sample_latency(fixed, 12_345), Duration::from_micros(500));
let range = (Duration::from_micros(1500), Duration::from_micros(2500));
for draw in [0, 1, 999, 1000, 1_000_000, u64::MAX] {
let latency = sample_latency(range, draw);
assert!(latency >= range.0 && latency <= range.1, "{latency:?}");
}
}
#[tokio::test(start_paused = true)]
async fn a_latency_past_the_timeout_is_recorded_as_a_timeout() {
let net = SimNet::new().latency(Duration::from_secs(60), Duration::from_secs(60));
let limit = Duration::from_secs(30);
let start = tokio::time::Instant::now();
assert_eq!(
net.transmit("h", Some(limit)).await,
Err(NetFault::TimedOut)
);
assert_eq!(start.elapsed(), limit, "only the timeout passes");
let events = net.events();
assert_eq!(events[0].fault, NetFault::TimedOut);
assert_eq!(events[0].latency, limit);
assert_eq!(
net.transmit("h", None).await,
Ok(()),
"no timeout, no fault"
);
}
#[test]
fn max_below_min_is_raised() {
let net = SimNet::new().latency(Duration::from_millis(9), Duration::from_millis(1));
assert_eq!(
net.lock().latency,
(Duration::from_millis(9), Duration::from_millis(9))
);
}
#[tokio::test(start_paused = true)]
async fn partition_fails_at_once_and_heal_restores() {
let net = SimNet::new();
net.partition("a");
assert_eq!(net.transmit("a", None).await, Err(NetFault::Partitioned));
net.heal("a");
assert_eq!(net.transmit("a", None).await, Ok(()));
let faults: Vec<_> = net.events().iter().map(|event| event.fault).collect();
assert_eq!(faults, vec![NetFault::Partitioned, NetFault::None]);
}
#[tokio::test(start_paused = true)]
async fn an_attempt_dropped_during_its_latency_leaves_no_event() {
let net = SimNet::new().latency(Duration::from_millis(10), Duration::from_millis(10));
let outcome = crate::sim::crash_at(0, net.transmit("h", None)).await;
assert!(outcome.is_crashed());
assert!(net.events().is_empty(), "the host was never reached");
assert_eq!(net.transmit("h", None).await, Ok(()));
assert_eq!(net.events().len(), 1);
assert_eq!(net.events()[0].seq, 0);
}
#[tokio::test(start_paused = true)]
async fn same_seed_same_decisions() {
let run = |seed| async move {
let net = SimNet::new()
.drop_rate(0.5)
.latency(Duration::ZERO, Duration::from_millis(30));
net.reseed(seed);
for _ in 0..32 {
let _ = net.transmit("h", None).await;
}
net.events()
};
assert_eq!(run(1).await, run(1).await);
assert_ne!(run(1).await, run(2).await);
}
}