refactor: replace Flume channels with Tokio and std mpsc

src/ratecontroller.rs

@@ -16,13 +16,13 @@ use crate::platform::{
 use crate::settings::Settings;
 use crate::time::{ClockSample, Realtime};
 use crate::util::{ArcRwLock, RwLockExt};
-use flume::{Receiver, Sender};
 use log::{debug, info, warn};
 use std::fs::File;
 use std::io::Write;
 use std::sync::atomic::Ordering;
 use std::thread::sleep;
 use std::time::{Duration, Instant};
+use tokio::sync::mpsc;
 
 #[derive(Copy, Clone, Debug, PartialEq)]
 enum Direction {
@@ -84,20 +84,6 @@ impl<H> State<H> {
     }
 }
 
-fn drain_latest_snapshot(
-    snapshot_rx: &Receiver<ControlSnapshot>,
-    latest: &mut Option<ControlSnapshot>,
-) {
-    while let Ok(snapshot) = snapshot_rx.try_recv() {
-        if latest
-            .as_ref()
-            .is_none_or(|current| snapshot.generated_at >= current.generated_at)
-        {
-            *latest = Some(snapshot);
-        }
-    }
-}
-
 fn deltas_from_snapshot(
     snapshot: Option<&ControlSnapshot>,
     reflectors: &[MeasurementStream],
@@ -137,10 +123,9 @@ fn deltas_from_snapshot(
 
 pub struct Ratecontroller<S: InterfaceStatsProvider, T: TrafficControlBackend> {
     settings: Settings,
-    control_snapshot: Option<ControlSnapshot>,
-    control_snapshot_rx: Receiver<ControlSnapshot>,
+    control_snapshot_rx: tokio::sync::watch::Receiver<Option<ControlSnapshot>>,
     reflectors_lock: ArcRwLock<Vec<MeasurementStream>>,
-    reselect_trigger: Sender<bool>,
+    reselect_trigger: mpsc::Sender<bool>,
     state_dl: State<T::Handle>,
     state_ul: State<T::Handle>,
     rng: fastrand::Rng,
@@ -156,9 +141,9 @@ impl<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
     )]
     fn new_with_backends(
         settings: Settings,
-        control_snapshot_rx: Receiver<ControlSnapshot>,
+        control_snapshot_rx: tokio::sync::watch::Receiver<Option<ControlSnapshot>>,
         reflectors_lock: ArcRwLock<Vec<MeasurementStream>>,
-        reselect_trigger: Sender<bool>,
+        reselect_trigger: mpsc::Sender<bool>,
         metrics: MetricsSender,
         mut stats_provider: S,
         mut traffic_control: T,
@@ -183,7 +168,6 @@ impl<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
 
         Ok(Self {
             settings,
-            control_snapshot: None,
             control_snapshot_rx,
             reflectors_lock,
             reselect_trigger,
@@ -483,15 +467,17 @@ impl<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
     }
 
     fn update_deltas(&mut self) -> anyhow::Result<()> {
-        drain_latest_snapshot(&self.control_snapshot_rx, &mut self.control_snapshot);
         let now_t = Instant::now();
         let reflectors = self.reflectors_lock.read_anyhow()?;
-        let (down_deltas, up_deltas) = deltas_from_snapshot(
-            self.control_snapshot.as_ref(),
-            &reflectors,
-            now_t,
-            self.settings.advanced_settings.tick_interval,
-        );
+        let (down_deltas, up_deltas) = {
+            let snapshot = self.control_snapshot_rx.borrow();
+            deltas_from_snapshot(
+                snapshot.as_ref(),
+                &reflectors,
+                now_t,
+                self.settings.advanced_settings.tick_interval,
+            )
+        };
         self.state_dl.deltas = down_deltas;
         self.state_ul.deltas = up_deltas;
 
@@ -517,9 +503,9 @@ impl<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
 impl Ratecontroller<PlatformInterfaceStats, PlatformTrafficControl> {
     pub fn new(
         settings: Settings,
-        control_snapshot_rx: Receiver<ControlSnapshot>,
+        control_snapshot_rx: tokio::sync::watch::Receiver<Option<ControlSnapshot>>,
         reflectors_lock: ArcRwLock<Vec<MeasurementStream>>,
-        reselect_trigger: Sender<bool>,
+        reselect_trigger: mpsc::Sender<bool>,
         metrics: MetricsSender,
     ) -> anyhow::Result<Self> {
         Self::new_with_backends(
@@ -534,111 +520,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 drains_queued_snapshots_and_keeps_the_newest() {
-        let start = Instant::now();
-        let newest_at = start + Duration::from_secs(2);
-        let (tx, rx) = flume::unbounded();
-        tx.send(snapshot(start + Duration::from_secs(1), 1.0))
-            .unwrap();
-        tx.send(snapshot(newest_at, 2.0)).unwrap();
-        let mut latest = None;
-
-        drain_latest_snapshot(&rx, &mut latest);
-
-        let latest = latest.unwrap();
-        assert!(rx.is_empty());
-        assert_eq!(latest.generated_at, newest_at);
-        assert_eq!(
-            latest.reflectors[&stream(0, "192.0.2.1".parse().unwrap())]
-                .recent
-                .down,
-            2.0
-        );
-    }
-
-    #[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]);
-    }
-}