refactor(measurement): introduce Controller for dynamic ICMP probe configuration

src/controller.rs

@@ -0,0 +1,102 @@
+use std::{collections::HashMap, net::IpAddr, time::Duration};
+
+use flume::{Receiver, SendError, Sender};
+use thiserror::Error;
+
+use crate::measurement::{
+    icmp::IcmpBinding,
+    model::{MeasurementSource, MeasurementStream, MeasurementStreamId, ProbeSchedule},
+};
+
+#[derive(Debug, Error)]
+pub enum ControllerError {
+    #[error("ICMP channel error")]
+    IcmpChannelError(#[from] SendError<Vec<IcmpBinding>>),
+}
+
+pub struct Controller {
+    reflector_rx: Receiver<Vec<IpAddr>>,
+    icmp_binding_tx: Sender<Vec<IcmpBinding>>,
+
+    source: MeasurementSource,
+    period: Duration,
+
+    next_stream_id: u64,
+    stream_ids: HashMap<(IpAddr, MeasurementSource), MeasurementStreamId>,
+}
+
+impl Controller {
+    pub fn new(
+        reflector_rx: Receiver<Vec<IpAddr>>,
+        icmp_binding_tx: Sender<Vec<IcmpBinding>>,
+    ) -> Self {
+        Self {
+            reflector_rx,
+            icmp_binding_tx,
+            source: MeasurementSource::IcmpTimestamp,
+            period: Duration::from_millis(500),
+            next_stream_id: 0,
+            stream_ids: HashMap::new(),
+        }
+    }
+
+    pub async fn run(mut self, initial_reflectors: Vec<IpAddr>) -> Result<(), ControllerError> {
+        self.update_reflectors(initial_reflectors).await?;
+
+        while let Ok(reflectors) = self.reflector_rx.recv_async().await {
+            self.update_reflectors(reflectors).await?;
+        }
+
+        Ok(())
+    }
+
+    fn stream_for(&mut self, peer: IpAddr) -> MeasurementStream {
+        let key = (peer, self.source);
+
+        let id = *self.stream_ids.entry(key).or_insert_with(|| {
+            let id = MeasurementStreamId(self.next_stream_id);
+            self.next_stream_id += 1;
+            id
+        });
+
+        MeasurementStream {
+            id,
+            peer,
+            source: self.source,
+        }
+    }
+
+    async fn update_reflectors(&mut self, reflectors: Vec<IpAddr>) -> Result<(), ControllerError> {
+        let count = reflectors.len();
+
+        let bindings = reflectors
+            .into_iter()
+            .enumerate()
+            .map(|(index, peer)| {
+                let stream = self.stream_for(peer);
+
+                IcmpBinding::new(
+                    stream,
+                    ProbeSchedule {
+                        period: self.period,
+                        offset: evenly_spaced_offset(self.period, index, count),
+                    },
+                )
+            })
+            .collect();
+
+        self.icmp_binding_tx
+            .send_async(bindings)
+            .await
+            .map_err(|e| ControllerError::IcmpChannelError(e))
+    }
+}
+
+fn evenly_spaced_offset(period: Duration, index: usize, count: usize) -> Duration {
+    let nanos = period.as_nanos() * index as u128 / count as u128;
+
+    Duration::new(
+        (nanos / 1_000_000_000) as u64,
+        (nanos % 1_000_000_000) as u32,
+    )
+}

src/main.rs

@@ -8,6 +8,7 @@
 extern crate core;
 
 mod baseliner;
+mod controller;
 mod log;
 mod measurement;
 mod metrics;
@@ -22,6 +23,7 @@ mod time;
 mod util;
 
 use crate::baseliner::Baseliner;
+use crate::controller::Controller;
 use crate::measurement::icmp::{IcmpBinding, IcmpEngine};
 use crate::measurement::model::{
     MeasurementSource, MeasurementStream, MeasurementStreamId, ProbeSchedule,
@@ -35,6 +37,7 @@ use crate::ratecontroller::Ratecontroller;
 use crate::reflector_selector::ReflectorSelector;
 use crate::settings::MeasurementType;
 use crate::settings::Settings;
+use crate::util::RwLockExt;
 use ::log::{info, warn};
 use flume::RecvTimeoutError;
 use std::collections::HashMap;
@@ -42,7 +45,7 @@ use std::net::{IpAddr, Ipv4Addr};
 use std::str::FromStr;
 use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
-use std::thread::sleep;
+use std::thread::{JoinHandle, sleep};
 use std::time::{Duration, Instant};
 use std::{process, thread};
 
@@ -223,6 +226,26 @@ fn setup_reflectors(settings: &Settings) -> anyhow::Result<ReflectorSetup> {
     })
 }
 
+fn spawn_worker<F>(
+    name: &str,
+    error_tx: &flume::Sender<anyhow::Error>,
+    worker: F,
+) -> anyhow::Result<JoinHandle<()>>
+where
+    F: FnOnce() -> anyhow::Result<()> + Send + 'static,
+{
+    let error_tx = error_tx.clone();
+
+    thread::Builder::new()
+        .name(name.to_owned())
+        .spawn(move || {
+            if let Err(error) = worker() {
+                let _ = error_tx.send(error);
+            }
+        })
+        .map_err(|e| anyhow::anyhow!(e))
+}
+
 fn wait_for_exit(error_rx: &flume::Receiver<anyhow::Error>) -> anyhow::Result<()> {
     loop {
         match error_rx.recv_timeout(Duration::from_secs(1)) {
@@ -243,7 +266,7 @@ fn wait_for_exit(error_rx: &flume::Receiver<anyhow::Error>) -> anyhow::Result<()
 }
 
 fn run(settings: &Settings) -> anyhow::Result<()> {
-    let rt = tokio::runtime::Runtime::new().unwrap();
+    let rt = tokio::runtime::Runtime::new()?;
 
     if settings.advanced_settings.dry_run {
         info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***");
@@ -252,26 +275,6 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
     let start_time = Instant::now();
     let probe_id = probe_identifier();
 
-    let (ie_txt, _) = flume::unbounded();
-    let (binding_tx, binding_rx) = flume::unbounded();
-    let engine = {
-        let _guard = rt.enter();
-        IcmpEngine::new(start_time, &settings.network, ie_txt, binding_rx)?
-    };
-    binding_tx.send(vec![IcmpBinding::new(
-        MeasurementStream {
-            id: MeasurementStreamId(0),
-            peer: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
-            source: MeasurementSource::IcmpEcho,
-        },
-        ProbeSchedule {
-            period: Duration::from_secs(1),
-            offset: Duration::from_secs(0),
-        },
-    )])?;
-
-    rt.spawn(engine.run());
-
     let ReflectorSetup {
         peers: reflector_peers,
         pool: reflector_pool,
@@ -281,14 +284,32 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
 
     let (baseliner_stats_tx, baseliner_stats_rx) = flume::unbounded();
     let (control_snapshot_tx, control_snapshot_rx) = flume::unbounded();
+    let (error_tx, error_rx) = flume::unbounded::<anyhow::Error>();
+    let (icmp_binding_tx, icmp_binding_rx) = flume::unbounded();
+    let (measurement_tx, measurement_rx) = flume::unbounded();
+    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 {
         let (tx, rx) = flume::unbounded();
         (Some(tx), Some(rx))
     } else {
         (None, None)
     };
-    let (error_tx, error_rx) = flume::unbounded::<anyhow::Error>();
+
-    let (reselect_tx, reselect_rx) = flume::bounded(1);
+    let engine = {
+        let _guard = rt.enter();
+        IcmpEngine::new(
+            start_time,
+            &settings.network,
+            measurement_tx,
+            icmp_binding_rx,
+        )?
+    };
+
+    rt.spawn(engine.run());
+
+    let controller = Controller::new(reflector_rx, icmp_binding_tx);
+    rt.spawn(controller.run(reflector_peers.read_anyhow()?.clone()));
 
     let dropped = Arc::new(AtomicU32::new(0));
 
@@ -299,14 +320,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
             metrics_rx: rx,
             metrics_dropped: Arc::clone(&dropped),
         };
-        let err_tx = error_tx.clone();
+        let handle = spawn_worker("metrics", &error_tx, move || metrics.run())?;
-        let handle = thread::Builder::new()
-            .name("metrics".to_string())
-            .spawn(move || {
-                if let Err(error) = metrics.run() {
-                    let _ = err_tx.send(error);
-                }
-            })?;
         (Some(tx), Some(handle))
     } else {
         (None, None)
@@ -366,52 +380,35 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
     let mut main_traffic_control = traffic_control_backend();
     let (down_shaper, up_shaper) = initialize_shaper(settings, &mut main_traffic_control)?;
 
-    let err_tx = error_tx.clone();
     let listener_peers = Arc::clone(&reflector_peers);
     let listener_inflight = Arc::clone(&inflight);
     let measurement_type = settings.advanced_settings.measurement_type;
-    thread::Builder::new()
+    spawn_worker("receiver", &error_tx, move || {
-        .name("receiver".to_string())
+        ping_listener.listen(
-        .spawn(move || {
+            probe_id,
-            if let Err(error) = ping_listener.listen(
+            measurement_type,
-                probe_id,
+            listener_peers,
-                measurement_type,
+            listener_inflight,
-                listener_peers,
+            baseliner_stats_tx,
-                listener_inflight,
+            ping_metrics,
-                baseliner_stats_tx,
+        )
-                ping_metrics,
+    })?;
-            ) {
-                let _ = err_tx.send(error);
-            }
-        })?;
 
-    let err_tx = error_tx.clone();
+    spawn_worker("baseliner", &error_tx, move || baseliner.run())?;
-    thread::Builder::new()
-        .name("baseliner".to_string())
-        .spawn(move || {
-            if let Err(error) = baseliner.run() {
-                let _ = err_tx.send(error);
-            }
-        })?;
 
-    let err_tx = error_tx.clone();
     let sender_peers = Arc::clone(&reflector_peers);
     let sender_inflight = Arc::clone(&inflight);
     let measurement_type = settings.advanced_settings.measurement_type;
     let tick_interval = settings.advanced_settings.tick_interval;
-    thread::Builder::new()
+    spawn_worker("sender", &error_tx, move || {
-        .name("sender".to_string())
+        ping_sender.send(
-        .spawn(move || {
+            probe_id,
-            if let Err(error) = ping_sender.send(
+            measurement_type,
-                probe_id,
+            sender_peers,
-                measurement_type,
+            sender_inflight,
-                sender_peers,
+            tick_interval,
-                sender_inflight,
+        )
-                tick_interval,
+    })?;
-            ) {
-                let _ = err_tx.send(error);
-            }
-        })?;
 
     let main_event_metrics = event_metrics.clone();
 
@@ -424,15 +421,9 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
             reflector_pool,
             trigger_channel: reselect_rx,
             metrics: event_metrics,
+            reflector_tx,
         };
-        let err_tx = error_tx.clone();
+        spawn_worker("reselection", &error_tx, move || reflector_selector.run())?;
-        thread::Builder::new()
-            .name("reselection".to_string())
-            .spawn(move || {
-                if let Err(error) = reflector_selector.run() {
-                    let _ = err_tx.send(error);
-                }
-            })?;
     }
 
     // Give the baseliner time to collect initial samples before adjusting rates.
@@ -446,14 +437,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
         rate_metrics,
     )?;
 
-    let err_tx = error_tx.clone();
+    spawn_worker("ratecontroller", &error_tx, move || ratecontroller.run())?;
-    thread::Builder::new()
-        .name("ratecontroller".to_string())
-        .spawn(move || {
-            if let Err(error) = ratecontroller.run() {
-                let _ = err_tx.send(error);
-            }
-        })?;
 
     // Drop the original sender so the channel disconnects if all workers exit cleanly.
     drop(error_tx);

src/reflector_selector.rs

@@ -10,7 +10,7 @@ use crate::baseliner::ControlSnapshot;
 use crate::metrics::{Metric, MetricsSender};
 use crate::settings::Settings;
 use crate::util::RwLockExt;
-use flume::{Receiver, RecvError, RecvTimeoutError, Selector};
+use flume::{Receiver, RecvError, RecvTimeoutError, Selector, Sender};
 use log::{debug, info};
 use std::net::IpAddr;
 use std::sync::atomic::Ordering;
@@ -108,6 +108,7 @@ pub struct ReflectorSelector {
     pub settings: Settings,
     pub snapshot_rx: Receiver<ControlSnapshot>,
     pub reflector_peers_lock: Arc<RwLock<Vec<IpAddr>>>,
+    pub reflector_tx: Sender<Vec<IpAddr>>,
     pub reflector_pool: Vec<IpAddr>,
     pub trigger_channel: Receiver<bool>,
     pub metrics: MetricsSender,
@@ -186,6 +187,7 @@ impl ReflectorSelector {
             // Clone next_peers because we need it again after the baseline sleep
             // to iterate over candidates for RTT measurement.
             *reflectors_peers = next_peers.clone();
+            self.reflector_tx.send(next_peers.clone())?;
 
             // Drop the MutexGuard explicitly, as Rust won't unlock the mutex by default
             // until the guard goes out of scope
@@ -251,6 +253,7 @@ impl ReflectorSelector {
                 });
             }
 
+            self.reflector_tx.send(new_peers.clone())?;
             *reflectors_peers = new_peers;
         }
     }