refactor: use an unbounded mpsc tokio channel for measurements

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<MeasurementStream, ReflectorState>,
 }
 
@@ -78,7 +79,7 @@ pub struct Baseliner {
     pub settings: Settings,
     pub reselect_trigger: Sender<bool>,
     pub start_time: Instant,
-    pub stats_rx: Receiver<MeasurementEvent>,
+    pub stats_rx: UnboundedReceiver<MeasurementEvent>,
     pub control_tx: tokio::sync::watch::Sender<Option<ControlSnapshot>>,
     pub selection_tx: Option<Sender<ControlSnapshot>>,
     pub baseline_metrics: MetricsSender,
@@ -142,18 +143,14 @@ fn process_reply(
     (*state, anomaly)
 }
 
-fn control_snapshot(
-    reflectors: &HashMap<MeasurementStream, ReflectorState>,
-    generated_at: Instant,
-) -> ControlSnapshot {
+fn control_snapshot(reflectors: &HashMap<MeasurementStream, ReflectorState>) -> 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::<MeasurementStream, ReflectorState>::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);
-    }
 }

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,

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<Vec<IcmpBinding>>,
-    observation_tx: Sender<MeasurementEvent>,
+    observation_tx: UnboundedSender<MeasurementEvent>,
 }
 
 async fn wait_for_deadline(deadline: Option<Instant>) {
@@ -123,7 +124,7 @@ impl IcmpEngine {
     pub fn new(
         origin: Instant,
         #[allow(unused)] settings: &NetworkSettings,
-        observation_tx: Sender<MeasurementEvent>,
+        observation_tx: UnboundedSender<MeasurementEvent>,
         binding_rx: Receiver<Vec<IcmpBinding>>,
     ) -> Result<Self, IcmpError> {
         #[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)?;
                         }
 

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<ManagedTargetConfig>,
 
     binding_rx: Receiver<Vec<IrttBinding>>,
-    observation_tx: Sender<MeasurementEvent>,
+    observation_tx: UnboundedSender<MeasurementEvent>,
 }
 
 impl IrttEngine {
     pub fn new(
         interval: std::time::Duration,
         settings: &NetworkSettings,
-        observation_tx: Sender<MeasurementEvent>,
+        observation_tx: UnboundedSender<MeasurementEvent>,
         binding_rx: Receiver<Vec<IrttBinding>>,
     ) -> 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)?;
                                 }
 

src/ratecontroller.rs

@@ -515,88 +515,3 @@ impl Ratecontroller<PlatformInterfaceStats, PlatformTrafficControl> {
         )
     }
 }
-
-#[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]);
-    }
-}

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)
-        );
-    }
-}