use std::collections::HashSet;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use bytes::Bytes;
use crate::errors::CoreError;
use crate::transport::session_transport::{FramePhase, SessionTransport};
fn splitmix64_next(state: &mut u64) -> u64 {
*state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut z = *state;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
fn prob_to_threshold(prob: f64) -> u64 {
let p = prob.clamp(0.0, 1.0);
if p <= 0.0 {
0
} else if p >= 1.0 {
u64::MAX
} else {
(p * (u64::MAX as f64)) as u64
}
}
struct StochasticConfig {
loss_threshold: u64,
dup_threshold: u64,
reorder_threshold: u64,
}
struct FaultState {
send_index: AtomicU64,
arm_drop: AtomicU64,
drop_indices: HashSet<u64>,
dup_indices: HashSet<u64>,
reorder_indices: HashSet<u64>,
pending_reorder: Mutex<Option<Vec<u8>>>,
delay_ms: AtomicU64,
stochastic: Option<StochasticConfig>,
stochastic_armed: AtomicBool,
prng: Mutex<u64>,
}
enum SendFault {
Forward,
Drop,
Duplicate,
Reorder,
}
fn lock_recover<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
m.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[derive(Clone)]
pub struct FaultControl {
state: Arc<FaultState>,
}
impl FaultControl {
fn build(
drop_indices: HashSet<u64>,
dup_indices: HashSet<u64>,
reorder_indices: HashSet<u64>,
delay_ms: u64,
stochastic: Option<StochasticConfig>,
seed: u64,
) -> Self {
Self {
state: Arc::new(FaultState {
send_index: AtomicU64::new(0),
arm_drop: AtomicU64::new(0),
drop_indices,
dup_indices,
reorder_indices,
pending_reorder: Mutex::new(None),
delay_ms: AtomicU64::new(delay_ms),
stochastic_armed: AtomicBool::new(stochastic.is_some()),
stochastic,
prng: Mutex::new(seed),
}),
}
}
pub fn new() -> Self {
Self::build(HashSet::new(), HashSet::new(), HashSet::new(), 0, None, 0)
}
pub fn with_drop_indices(indices: &[u64]) -> Self {
Self::build(
indices.iter().copied().collect(),
HashSet::new(),
HashSet::new(),
0,
None,
0,
)
}
pub fn with_dup_indices(indices: &[u64]) -> Self {
Self::build(
HashSet::new(),
indices.iter().copied().collect(),
HashSet::new(),
0,
None,
0,
)
}
pub fn with_reorder_indices(indices: &[u64]) -> Self {
Self::build(
HashSet::new(),
HashSet::new(),
indices.iter().copied().collect(),
0,
None,
0,
)
}
pub fn with_delay(delay: Duration) -> Self {
Self::build(
HashSet::new(),
HashSet::new(),
HashSet::new(),
delay.as_millis() as u64,
None,
0,
)
}
pub fn with_seed(
seed: u64,
loss_prob: f64,
dup_prob: f64,
reorder_prob: f64,
delay_ms: u64,
) -> Self {
let stochastic = StochasticConfig {
loss_threshold: prob_to_threshold(loss_prob),
dup_threshold: prob_to_threshold(dup_prob),
reorder_threshold: prob_to_threshold(reorder_prob),
};
Self::build(
HashSet::new(),
HashSet::new(),
HashSet::new(),
delay_ms,
Some(stochastic),
seed,
)
}
pub fn arm_drop_next(&self, n: u64) {
self.state.arm_drop.store(n, Ordering::Relaxed);
}
pub fn set_delay(&self, delay: Duration) {
self.state
.delay_ms
.store(delay.as_millis() as u64, Ordering::Relaxed);
}
pub fn arm_stochastic(&self, armed: bool) {
self.state.stochastic_armed.store(armed, Ordering::Relaxed);
}
fn draw(&self) -> u64 {
let mut state = lock_recover(&self.state.prng);
splitmix64_next(&mut state)
}
fn classify(&self) -> SendFault {
let index = self.state.send_index.fetch_add(1, Ordering::Relaxed);
if self.state.drop_indices.contains(&index) || self.consume_armed_drop() {
return SendFault::Drop;
}
if self.state.dup_indices.contains(&index) {
return SendFault::Duplicate;
}
if self.state.reorder_indices.contains(&index) {
return SendFault::Reorder;
}
if self.state.stochastic_armed.load(Ordering::Relaxed) {
if let Some(cfg) = self.state.stochastic.as_ref() {
if cfg.loss_threshold > 0 && self.draw() < cfg.loss_threshold {
return SendFault::Drop;
}
if cfg.dup_threshold > 0 && self.draw() < cfg.dup_threshold {
return SendFault::Duplicate;
}
if cfg.reorder_threshold > 0 && self.draw() < cfg.reorder_threshold {
return SendFault::Reorder;
}
}
}
SendFault::Forward
}
fn delay(&self) -> Option<Duration> {
match self.state.delay_ms.load(Ordering::Relaxed) {
0 => None,
ms => Some(Duration::from_millis(ms)),
}
}
fn take_pending_reorder(&self) -> Option<Vec<u8>> {
lock_recover(&self.state.pending_reorder).take()
}
fn hold_for_reorder(&self, data: Vec<u8>) -> Option<Vec<u8>> {
lock_recover(&self.state.pending_reorder).replace(data)
}
fn consume_armed_drop(&self) -> bool {
loop {
let pending = self.state.arm_drop.load(Ordering::Relaxed);
if pending == 0 {
return false;
}
if self
.state
.arm_drop
.compare_exchange_weak(pending, pending - 1, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
return true;
}
}
}
}
impl Default for FaultControl {
fn default() -> Self {
Self::new()
}
}
pub struct LossyTransport<T> {
inner: T,
control: FaultControl,
}
impl<T> LossyTransport<T> {
pub fn new(inner: T, control: FaultControl) -> Self {
Self { inner, control }
}
pub fn drop_sends(inner: T, indices: &[u64]) -> Self {
Self::new(inner, FaultControl::with_drop_indices(indices))
}
pub fn dup_sends(inner: T, indices: &[u64]) -> Self {
Self::new(inner, FaultControl::with_dup_indices(indices))
}
pub fn reorder_sends(inner: T, indices: &[u64]) -> Self {
Self::new(inner, FaultControl::with_reorder_indices(indices))
}
}
impl<T: SessionTransport> LossyTransport<T> {
pub async fn flush(&self) -> Result<(), CoreError> {
if let Some(held) = self.control.take_pending_reorder() {
self.inner.send_bytes(&held).await?;
}
Ok(())
}
}
impl<T: SessionTransport> SessionTransport for LossyTransport<T> {
async fn send_bytes(&self, data: &[u8]) -> Result<(), CoreError> {
if let Some(delay) = self.control.delay() {
tokio::time::sleep(delay).await;
}
match self.control.classify() {
SendFault::Drop => {
Ok(())
}
SendFault::Reorder => {
if let Some(prev) = self.control.hold_for_reorder(data.to_vec()) {
self.inner.send_bytes(&prev).await?;
}
Ok(())
}
SendFault::Duplicate => {
self.inner.send_bytes(data).await?;
self.inner.send_bytes(data).await?;
self.release_pending_reorder().await
}
SendFault::Forward => {
self.inner.send_bytes(data).await?;
self.release_pending_reorder().await
}
}
}
async fn recv_bytes(&self) -> Result<Bytes, CoreError> {
self.inner.recv_bytes().await
}
fn set_frame_phase(&self, phase: FramePhase) {
self.inner.set_frame_phase(phase);
}
fn set_outbound_cid(&self, cid: [u8; 8]) {
self.inner.set_outbound_cid(cid);
}
fn has_migration_candidate(&self) -> bool {
self.inner.has_migration_candidate()
}
async fn send_to_candidate(&self, data: &[u8]) -> Result<bool, CoreError> {
self.inner.send_to_candidate(data).await
}
fn confirm_authenticated_source(&self) {
self.inner.confirm_authenticated_source();
}
fn promote_candidate(&self) -> bool {
self.inner.promote_candidate()
}
async fn migrate(&self, local_addr: String) -> Result<(), CoreError> {
self.inner.migrate(local_addr).await
}
async fn migrate_server(&self, local_addr: String) -> Result<(), CoreError> {
self.inner.migrate_server(local_addr).await
}
}
impl<T: SessionTransport> LossyTransport<T> {
async fn release_pending_reorder(&self) -> Result<(), CoreError> {
if let Some(held) = self.control.take_pending_reorder() {
self.inner.send_bytes(&held).await?;
}
Ok(())
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
use super::*;
use std::sync::Mutex;
struct RecordingTransport {
forwarded: Arc<Mutex<Vec<Vec<u8>>>>,
}
impl SessionTransport for RecordingTransport {
async fn send_bytes(&self, data: &[u8]) -> Result<(), CoreError> {
self.forwarded.lock().expect("poisoned").push(data.to_vec());
Ok(())
}
async fn recv_bytes(&self) -> Result<Bytes, CoreError> {
Err(CoreError::NetworkError("recv unused in this test".into()))
}
}
#[tokio::test]
async fn lossy_transport_drops_configured_send_indices() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let lossy = LossyTransport::drop_sends(inner, &[1]);
lossy.send_bytes(b"f0").await.expect("send f0");
lossy.send_bytes(b"f1").await.expect("send f1"); lossy.send_bytes(b"f2").await.expect("send f2");
let got = forwarded.lock().expect("poisoned");
assert_eq!(
&*got,
&[b"f0".to_vec(), b"f2".to_vec()],
"frame at index 1 must be dropped, not forwarded to the inner transport"
);
}
#[tokio::test]
async fn arm_drop_next_drops_exactly_the_next_send() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let control = FaultControl::new();
let lossy = LossyTransport::new(inner, control.clone());
control.arm_drop_next(1); lossy.send_bytes(b"d0").await.expect("send d0"); lossy.send_bytes(b"d1").await.expect("send d1");
lossy.send_bytes(b"d2").await.expect("send d2");
let got = forwarded.lock().expect("poisoned");
assert_eq!(
&*got,
&[b"d1".to_vec(), b"d2".to_vec()],
"an armed drop of 1 must skip exactly the next send, then forward the rest"
);
}
#[tokio::test]
async fn dup_sends_forwards_configured_indices_twice() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let lossy = LossyTransport::dup_sends(inner, &[1]);
lossy.send_bytes(b"u0").await.expect("send u0");
lossy.send_bytes(b"u1").await.expect("send u1"); lossy.send_bytes(b"u2").await.expect("send u2");
let got = forwarded.lock().expect("poisoned");
assert_eq!(
&*got,
&[
b"u0".to_vec(),
b"u1".to_vec(),
b"u1".to_vec(),
b"u2".to_vec()
],
"the duplicated frame must reach the inner transport twice, in place"
);
}
#[tokio::test]
async fn reorder_sends_swaps_with_the_following_frame() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let lossy = LossyTransport::reorder_sends(inner, &[1]);
lossy.send_bytes(b"r0").await.expect("send r0");
lossy.send_bytes(b"r1").await.expect("send r1"); lossy.send_bytes(b"r2").await.expect("send r2");
let got = forwarded.lock().expect("poisoned");
assert_eq!(
&*got,
&[b"r0".to_vec(), b"r2".to_vec(), b"r1".to_vec()],
"the reordered frame must land after the following frame"
);
}
#[tokio::test]
async fn reorder_of_final_frame_is_released_by_flush() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let lossy = LossyTransport::reorder_sends(inner, &[1]);
lossy.send_bytes(b"r0").await.expect("send r0");
lossy.send_bytes(b"r1").await.expect("send r1");
assert_eq!(&*forwarded.lock().expect("poisoned"), &[b"r0".to_vec()]);
lossy.flush().await.expect("flush the held frame");
assert_eq!(
&*forwarded.lock().expect("poisoned"),
&[b"r0".to_vec(), b"r1".to_vec()],
"flush must release a reorder held past the final send"
);
}
#[tokio::test]
async fn delay_holds_each_send_for_at_least_the_configured_duration() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let control = FaultControl::with_delay(Duration::from_millis(25));
let lossy = LossyTransport::new(inner, control);
let start = tokio::time::Instant::now();
lossy.send_bytes(b"slow").await.expect("send slow");
assert!(
start.elapsed() >= Duration::from_millis(25),
"delay must hold the send for at least the configured duration"
);
assert_eq!(&*forwarded.lock().expect("poisoned"), &[b"slow".to_vec()]);
}
async fn dropped_indices_for(seed: u64, loss_prob: f64, n: u64) -> Vec<u64> {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let control = FaultControl::with_seed(seed, loss_prob, 0.0, 0.0, 0);
let lossy = LossyTransport::new(inner, control);
for i in 0..n {
let tag = i.to_le_bytes();
lossy.send_bytes(&tag).await.expect("send");
}
let forwarded = forwarded.lock().expect("poisoned");
let forwarded_set: HashSet<u64> = forwarded
.iter()
.map(|f| {
let mut b = [0u8; 8];
b.copy_from_slice(&f[..8]);
u64::from_le_bytes(b)
})
.collect();
(0..n).filter(|i| !forwarded_set.contains(i)).collect()
}
#[tokio::test]
async fn stochastic_loss_is_deterministic_for_a_fixed_seed() {
const N: u64 = 2_000;
let a = dropped_indices_for(0xDEAD_BEEF_CAFE_F00D, 0.1, N).await;
let b = dropped_indices_for(0xDEAD_BEEF_CAFE_F00D, 0.1, N).await;
assert_eq!(
a, b,
"same seed + config + send count must drop the exact same indices"
);
assert!(
!a.is_empty(),
"10% loss over 2000 sends must drop at least one frame"
);
let c = dropped_indices_for(0x0123_4567_89AB_CDEF, 0.1, N).await;
assert_ne!(
a, c,
"a different seed should produce a different drop pattern"
);
}
#[tokio::test]
async fn stochastic_loss_rate_matches_the_configured_probability() {
const N: u64 = 10_000;
const P: f64 = 0.1;
let dropped = dropped_indices_for(0x5EED_1234_5678_9ABC, P, N).await;
let frac = dropped.len() as f64 / N as f64;
assert!(
(frac - P).abs() <= 0.02,
"observed drop fraction {frac:.4} must be within 0.1 ± 0.02 over {N} sends \
({} dropped)",
dropped.len()
);
}
#[tokio::test]
async fn non_stochastic_controls_drop_nothing() {
async fn forwards_all(control: FaultControl, n: u64) -> bool {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let lossy = LossyTransport::new(inner, control);
for i in 0..n {
lossy.send_bytes(&i.to_le_bytes()).await.expect("send");
}
let count = forwarded.lock().expect("poisoned").len() as u64;
count == n
}
assert!(
forwards_all(FaultControl::new(), 1_000).await,
"FaultControl::new() must forward every send (no stochastic drops)"
);
assert!(
forwards_all(FaultControl::with_dup_indices(&[]), 1_000).await,
"with_dup_indices(&[]) must forward every send 1:1"
);
assert!(
forwards_all(FaultControl::with_drop_indices(&[]), 1_000).await,
"with_drop_indices(&[]) must drop nothing"
);
}
#[test]
fn splitmix64_matches_reference_vector() {
let mut state = 0u64;
let got = [
splitmix64_next(&mut state),
splitmix64_next(&mut state),
splitmix64_next(&mut state),
];
assert_eq!(
got,
[
0xE220_A839_7B1D_CDAF,
0x6E78_9E6A_A1B9_65F4,
0x06C4_5D18_8009_454F
],
"SplitMix64 must match the canonical reference stream for seed 0"
);
}
#[tokio::test]
async fn stochastic_arm_gate_controls_when_loss_applies() {
let forwarded = Arc::new(Mutex::new(Vec::new()));
let inner = RecordingTransport {
forwarded: forwarded.clone(),
};
let control = FaultControl::with_seed(7, 1.0, 0.0, 0.0, 0);
control.arm_stochastic(false);
let lossy = LossyTransport::new(inner, control.clone());
for i in 0..5u64 {
lossy.send_bytes(&i.to_le_bytes()).await.expect("send");
}
assert_eq!(
forwarded.lock().expect("poisoned").len(),
5,
"a disarmed stochastic control must forward every send"
);
control.arm_stochastic(true);
for i in 5..10u64 {
lossy.send_bytes(&i.to_le_bytes()).await.expect("send");
}
assert_eq!(
forwarded.lock().expect("poisoned").len(),
5,
"after arming, a 100%-loss control must drop every further send"
);
}
#[tokio::test]
async fn stochastic_endpoints_are_all_or_nothing() {
let all = dropped_indices_for(42, 1.0, 256).await;
assert_eq!(all.len(), 256, "loss_prob = 1.0 must drop every send");
let none = dropped_indices_for(42, 0.0, 256).await;
assert!(none.is_empty(), "loss_prob = 0.0 must drop nothing");
}
}