diff --git a/src/baseliner.rs b/src/baseliner.rs index c0b6c69..de37b61 100644 --- a/src/baseliner.rs +++ b/src/baseliner.rs @@ -9,12 +9,14 @@ use crate::SHUTDOWN; use crate::measurement::model::{MeasurementEvent, MeasurementObservation, MeasurementStream}; use crate::metrics::{Metric, MetricsSender}; use crate::settings::Settings; -use flume::{Receiver, Sender}; +use anyhow::Context; +use flume::Sender; use log::{debug, info}; use std::collections::HashMap; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use tokio::sync::mpsc::UnboundedReceiver; const STALE_AFTER: Duration = Duration::from_secs(30); const ANOMALY_THRESHOLD_MS: f64 = 5_000.0; @@ -70,7 +72,6 @@ impl ReflectorState { #[derive(Clone, Debug)] pub struct ControlSnapshot { - pub generated_at: Instant, pub reflectors: HashMap, } @@ -78,7 +79,7 @@ pub struct Baseliner { pub settings: Settings, pub reselect_trigger: Sender, pub start_time: Instant, - pub stats_rx: Receiver, + pub stats_rx: UnboundedReceiver, pub control_tx: tokio::sync::watch::Sender>, pub selection_tx: Option>, pub baseline_metrics: MetricsSender, @@ -142,18 +143,14 @@ fn process_reply( (*state, anomaly) } -fn control_snapshot( - reflectors: &HashMap, - generated_at: Instant, -) -> ControlSnapshot { +fn control_snapshot(reflectors: &HashMap) -> ControlSnapshot { ControlSnapshot { - generated_at, reflectors: reflectors.clone(), } } impl Baseliner { - pub fn run(&self) -> anyhow::Result<()> { + pub fn run(&mut self) -> anyhow::Result<()> { let mut reflectors = HashMap::::new(); let slow_factor = ewma_factor( @@ -171,7 +168,10 @@ impl Baseliner { return Ok(()); } - let event = self.stats_rx.recv()?; + let event = self + .stats_rx + .blocking_recv() + .context("measurement channel closed")?; let reply = match event { MeasurementEvent::Observation(reply) => reply, @@ -232,7 +232,7 @@ impl Baseliner { reflector, state.recent.up, state.recent.down ); - let snapshot = control_snapshot(&reflectors, Instant::now()); + let snapshot = control_snapshot(&reflectors); if let Some(selection_tx) = &self.selection_tx { let _ = selection_tx.send(snapshot.clone()); } @@ -454,32 +454,4 @@ mod tests { assert_eq!(updated.last_receive_at, received_at); assert_close(reflectors[&stream].recent.up, 200.0); } - - #[test] - fn snapshot_contains_updated_combined_state() { - let start = Instant::now(); - let received_at = start + Duration::from_secs(1); - let generated_at = received_at + Duration::from_millis(1); - let reflector = "192.0.2.1".parse().unwrap(); - let mut reflectors = HashMap::new(); - - let (updated, anomaly) = process_reply( - &mut reflectors, - &reply(reflector, 20.0, 30.0, received_at), - start, - 0.9, - 0.5, - ); - let snapshot = control_snapshot(&reflectors, generated_at); - let stream = reply(reflector, 0.0, 0.0, received_at).stream; - - assert!(!anomaly); - assert_eq!(snapshot.generated_at, generated_at); - assert_close( - snapshot.reflectors[&stream].baseline.down, - updated.baseline.down, - ); - assert_close(snapshot.reflectors[&stream].recent.up, updated.recent.up); - assert_eq!(snapshot.reflectors[&stream].last_receive_at, received_at); - } } diff --git a/src/main.rs b/src/main.rs index d3aab15..c8235ed 100644 --- a/src/main.rs +++ b/src/main.rs @@ -351,7 +351,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { let (icmp_binding_tx, icmp_binding_rx) = flume::unbounded(); #[cfg(feature = "irtt")] let (irtt_binding_tx, irtt_binding_rx) = flume::unbounded(); - let (measurement_tx, measurement_rx) = flume::unbounded(); + let (measurement_tx, measurement_rx) = tokio::sync::mpsc::unbounded_channel(); let (reflector_tx, reflector_rx) = flume::unbounded(); let (reselect_tx, reselect_rx) = flume::bounded(1); let (selection_snapshot_tx, selection_snapshot_rx) = if reselection_enabled { @@ -440,7 +440,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { }, }); - let baseliner = Baseliner { + let mut baseliner = Baseliner { settings: settings.clone(), reselect_trigger: reselect_tx.clone(), start_time, diff --git a/src/measurement/icmp.rs b/src/measurement/icmp.rs index cbdeb59..6a60562 100644 --- a/src/measurement/icmp.rs +++ b/src/measurement/icmp.rs @@ -5,7 +5,7 @@ use std::{ time::{Duration, Instant, SystemTime}, }; -use flume::{Receiver, Sender}; +use flume::Receiver; use icmp_socket2::{ AsyncIcmpSocket, IcmpSocket4, Icmpv4Message, Icmpv4Packet, packet::{IcmpPacketBuildError, WithEchoRequest, WithTimestampRequest}, @@ -13,6 +13,7 @@ use icmp_socket2::{ }; use log::debug; use thiserror::Error; +use tokio::sync::mpsc::UnboundedSender; use crate::{ measurement::model::{ @@ -104,7 +105,7 @@ pub struct IcmpEngine { origin: Instant, binding_rx: Receiver>, - observation_tx: Sender, + observation_tx: UnboundedSender, } async fn wait_for_deadline(deadline: Option) { @@ -123,7 +124,7 @@ impl IcmpEngine { pub fn new( origin: Instant, #[allow(unused)] settings: &NetworkSettings, - observation_tx: Sender, + observation_tx: UnboundedSender, binding_rx: Receiver>, ) -> Result { #[allow(unused_mut)] @@ -420,12 +421,11 @@ impl IcmpEngine { for probe in expired { self.observation_tx - .send_async(MeasurementEvent::Loss(MeasurementLoss { + .send(MeasurementEvent::Loss(MeasurementLoss { stream: probe.stream, started_at: probe.sent_at, detected_at: now, })) - .await .map_err(|_| IcmpError::EventChannelClosed)?; } @@ -468,8 +468,7 @@ impl IcmpEngine { ) { Ok(observation) => { self.observation_tx - .send_async(MeasurementEvent::Observation(observation)) - .await + .send(MeasurementEvent::Observation(observation)) .map_err(|_| IcmpError::EventChannelClosed)?; } diff --git a/src/measurement/irtt.rs b/src/measurement/irtt.rs index ad6b818..ff87f6a 100644 --- a/src/measurement/irtt.rs +++ b/src/measurement/irtt.rs @@ -1,6 +1,6 @@ use std::{collections::HashMap, net::SocketAddr}; -use flume::{Receiver, Sender}; +use flume::Receiver; use irtt_client::{ ClientConfig, ClientEvent, NegotiationPolicy, OpenPolicy, managed::{ @@ -12,7 +12,7 @@ use irtt_client::{ }; use log::{debug, trace, warn}; use thiserror::Error; -use tokio::sync::broadcast::error::RecvError; +use tokio::sync::{broadcast::error::RecvError, mpsc::UnboundedSender}; use crate::SHUTDOWN; use crate::measurement::model::{ @@ -70,14 +70,14 @@ pub struct IrttEngine { targets: Vec, binding_rx: Receiver>, - observation_tx: Sender, + observation_tx: UnboundedSender, } impl IrttEngine { pub fn new( interval: std::time::Duration, settings: &NetworkSettings, - observation_tx: Sender, + observation_tx: UnboundedSender, binding_rx: Receiver>, ) -> Result<(ManagedClientTask, Self), IrttError> { let config = managed_client_config(interval, settings); @@ -273,8 +273,7 @@ impl IrttEngine { }; self.observation_tx - .send_async(MeasurementEvent::Observation(observation)) - .await + .send(MeasurementEvent::Observation(observation)) .map_err(|_| IrttError::EventChannelClosed)?; } @@ -294,8 +293,7 @@ impl IrttEngine { }; self.observation_tx - .send_async(MeasurementEvent::Loss(loss)) - .await + .send(MeasurementEvent::Loss(loss)) .map_err(|_| IrttError::EventChannelClosed)?; } diff --git a/src/ratecontroller.rs b/src/ratecontroller.rs index 5b6630e..f109fd5 100644 --- a/src/ratecontroller.rs +++ b/src/ratecontroller.rs @@ -515,88 +515,3 @@ impl Ratecontroller { ) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::baseliner::{EwmaStats, ReflectorState}; - use crate::measurement::model::{MeasurementSource, MeasurementStreamId}; - use std::collections::HashMap; - use std::net::IpAddr; - - fn stream(id: u64, peer: IpAddr) -> MeasurementStream { - MeasurementStream { - id: MeasurementStreamId(id), - peer, - source: MeasurementSource::IcmpTimestamp, - } - } - - #[test] - fn fixed_seed_reproduces_initial_safe_rates() { - let mut first = fastrand::Rng::with_seed(42); - let mut second = fastrand::Rng::with_seed(42); - - let first_rates = generate_initial_speeds(&mut first, 100_000.0, 100); - let second_rates = generate_initial_speeds(&mut second, 100_000.0, 100); - - assert_eq!(first_rates, second_rates); - } - - fn reflector_state( - baseline_down: f64, - baseline_up: f64, - recent_down: f64, - recent_up: f64, - last_receive_at: Instant, - ) -> ReflectorState { - ReflectorState { - baseline: EwmaStats { - down: baseline_down, - up: baseline_up, - }, - recent: EwmaStats { - down: recent_down, - up: recent_up, - }, - last_receive_at, - } - } - - fn snapshot(generated_at: Instant, marker: f64) -> ControlSnapshot { - let reflector = stream(0, "192.0.2.1".parse().unwrap()); - ControlSnapshot { - generated_at, - reflectors: HashMap::from([( - reflector, - reflector_state(marker, marker, marker, marker, generated_at), - )]), - } - } - - #[test] - fn calculates_sorted_deltas_from_active_fresh_snapshot_state() { - let now = Instant::now(); - let first = stream(0, "192.0.2.1".parse().unwrap()); - let second = stream(1, "192.0.2.2".parse().unwrap()); - let stale = stream(2, "192.0.2.3".parse().unwrap()); - let inactive = stream(3, "192.0.2.4".parse().unwrap()); - let snapshot = ControlSnapshot { - generated_at: now, - reflectors: HashMap::from([ - (first, reflector_state(10.0, 20.0, 14.0, 28.0, now)), - (second, reflector_state(10.0, 20.0, 11.0, 22.0, now)), - ( - stale, - reflector_state(10.0, 20.0, 110.0, 220.0, now - Duration::from_secs(3)), - ), - (inactive, reflector_state(10.0, 20.0, 60.0, 80.0, now)), - ]), - }; - - let (down, up) = deltas_from_snapshot(Some(&snapshot), &[first, second, stale], now, 1.0); - - assert_eq!(down, vec![1.0, 4.0]); - assert_eq!(up, vec![2.0, 8.0]); - } -} diff --git a/src/reflector_selector.rs b/src/reflector_selector.rs index 3cbbc89..31899c6 100644 --- a/src/reflector_selector.rs +++ b/src/reflector_selector.rs @@ -285,158 +285,3 @@ impl ReflectorSelector { } } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::baseliner::{EwmaStats, ReflectorState}; - use crate::measurement::model::MeasurementSource; - use std::collections::HashMap; - use std::net::IpAddr; - use std::thread; - - fn snapshot(generated_at: Instant) -> ControlSnapshot { - ControlSnapshot { - generated_at, - reflectors: HashMap::new(), - } - } - - fn snapshot_with_rtt(generated_at: Instant, peer: IpAddr, rtt: f64) -> ControlSnapshot { - let mut reflectors = HashMap::new(); - reflectors.insert( - crate::measurement::model::MeasurementStream { - id: crate::measurement::model::MeasurementStreamId(0), - peer, - source: MeasurementSource::IcmpTimestamp, - }, - ReflectorState { - baseline: EwmaStats { - down: rtt / 2.0, - up: rtt / 2.0, - }, - recent: EwmaStats { - down: rtt / 2.0, - up: rtt / 2.0, - }, - last_receive_at: generated_at, - }, - ); - ControlSnapshot { - generated_at, - reflectors, - } - } - - #[test] - fn wait_retains_the_newest_snapshot() { - let (_trigger_tx, trigger_rx) = flume::unbounded(); - let (snapshot_tx, snapshot_rx) = flume::unbounded(); - let generated_at = Instant::now(); - snapshot_tx.send(snapshot(generated_at)).unwrap(); - snapshot_tx - .send(snapshot(generated_at + Duration::from_millis(1))) - .unwrap(); - snapshot_tx - .send(snapshot(generated_at + Duration::from_millis(2))) - .unwrap(); - let mut latest_snapshot = None; - - let result = wait_for_trigger_or_deadline( - &trigger_rx, - &snapshot_rx, - &mut latest_snapshot, - Instant::now() + Duration::from_millis(20), - ); - - assert_eq!(result, WaitOutcome::DeadlineReached); - assert_eq!( - latest_snapshot.unwrap().generated_at, - generated_at + Duration::from_millis(2) - ); - } - - #[test] - fn snapshots_do_not_extend_the_deadline() { - let (_trigger_tx, trigger_rx) = flume::unbounded(); - let (snapshot_tx, snapshot_rx) = flume::unbounded(); - let producer = thread::spawn(move || { - let stop_at = Instant::now() + Duration::from_millis(500); - while Instant::now() < stop_at { - if snapshot_tx.send(snapshot(Instant::now())).is_err() { - break; - } - sleep(Duration::from_millis(1)); - } - }); - let mut latest_snapshot = None; - let started_at = Instant::now(); - - let result = wait_for_trigger_or_deadline( - &trigger_rx, - &snapshot_rx, - &mut latest_snapshot, - started_at + Duration::from_millis(30), - ); - let elapsed = started_at.elapsed(); - drop(snapshot_rx); - producer.join().unwrap(); - - assert_eq!(result, WaitOutcome::DeadlineReached); - assert!(latest_snapshot.is_some()); - assert!(elapsed < Duration::from_millis(250), "elapsed: {elapsed:?}"); - } - - #[test] - fn trigger_interrupts_the_wait() { - let (trigger_tx, trigger_rx) = flume::unbounded(); - let (_snapshot_tx, snapshot_rx) = flume::unbounded(); - trigger_tx.send(true).unwrap(); - let mut latest_snapshot = None; - let started_at = Instant::now(); - - let result = wait_for_trigger_or_deadline( - &trigger_rx, - &snapshot_rx, - &mut latest_snapshot, - started_at + Duration::from_secs(1), - ); - - assert_eq!(result, WaitOutcome::Triggered); - assert!(started_at.elapsed() < Duration::from_millis(250)); - } - - #[test] - fn baseline_wait_retains_snapshot_for_evaluation() { - let (snapshot_tx, snapshot_rx) = flume::unbounded(); - let peer: IpAddr = "192.0.2.1".parse().unwrap(); - let generated_at = Instant::now(); - snapshot_tx - .send(snapshot_with_rtt(generated_at, peer, 20.0)) - .unwrap(); - snapshot_tx - .send(snapshot_with_rtt( - generated_at + Duration::from_millis(1), - peer, - 42.0, - )) - .unwrap(); - let mut latest_snapshot = None; - - receive_snapshots_until( - &snapshot_rx, - &mut latest_snapshot, - Instant::now() + Duration::from_millis(20), - ); - - assert_eq!( - recent_rtt_for_peer( - latest_snapshot.as_ref(), - latest_snapshot.as_ref().unwrap().reflectors.keys().next(), - generated_at + Duration::from_secs(1), - Duration::from_secs(2) - ), - Some(42) - ); - } -}