From f9b0d9f7abae1c7110b26f2547456d4b62e014e7 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 6 Oct 2026 01:14:32 +0700 Subject: [PATCH] refactor(measurement): introduce Controller for dynamic ICMP probe configuration --- src/controller.rs | 102 +++++++++++++++++++++++++ src/main.rs | 154 +++++++++++++++++--------------------- src/reflector_selector.rs | 5 +- 3 files changed, 175 insertions(+), 86 deletions(-) create mode 100644 src/controller.rs diff --git a/src/controller.rs b/src/controller.rs new file mode 100644 index 0000000..11408a8 --- /dev/null +++ b/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>), +} + +pub struct Controller { + reflector_rx: Receiver>, + icmp_binding_tx: Sender>, + + source: MeasurementSource, + period: Duration, + + next_stream_id: u64, + stream_ids: HashMap<(IpAddr, MeasurementSource), MeasurementStreamId>, +} + +impl Controller { + pub fn new( + reflector_rx: Receiver>, + icmp_binding_tx: Sender>, + ) -> 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) -> 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) -> 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, + ) +} diff --git a/src/main.rs b/src/main.rs index 5bc51c7..fdf34b0 100644 --- a/src/main.rs +++ b/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 { }) } +fn spawn_worker( + name: &str, + error_tx: &flume::Sender, + worker: F, +) -> anyhow::Result> +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::Result<()> { loop { match error_rx.recv_timeout(Duration::from_secs(1)) { @@ -243,7 +266,7 @@ fn wait_for_exit(error_rx: &flume::Receiver) -> 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::(); + 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::(); - 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 = thread::Builder::new() - .name("metrics".to_string()) - .spawn(move || { - if let Err(error) = metrics.run() { - let _ = err_tx.send(error); - } - })?; + let handle = spawn_worker("metrics", &error_tx, move || metrics.run())?; (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() - .name("receiver".to_string()) - .spawn(move || { - if let Err(error) = ping_listener.listen( - probe_id, - measurement_type, - listener_peers, - listener_inflight, - baseliner_stats_tx, - ping_metrics, - ) { - let _ = err_tx.send(error); - } - })?; + spawn_worker("receiver", &error_tx, move || { + ping_listener.listen( + probe_id, + measurement_type, + listener_peers, + listener_inflight, + baseliner_stats_tx, + ping_metrics, + ) + })?; - let err_tx = error_tx.clone(); - thread::Builder::new() - .name("baseliner".to_string()) - .spawn(move || { - if let Err(error) = baseliner.run() { - let _ = err_tx.send(error); - } - })?; + spawn_worker("baseliner", &error_tx, move || baseliner.run())?; - 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() - .name("sender".to_string()) - .spawn(move || { - if let Err(error) = ping_sender.send( - probe_id, - measurement_type, - sender_peers, - sender_inflight, - tick_interval, - ) { - let _ = err_tx.send(error); - } - })?; + spawn_worker("sender", &error_tx, move || { + ping_sender.send( + probe_id, + measurement_type, + sender_peers, + sender_inflight, + tick_interval, + ) + })?; 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(); - thread::Builder::new() - .name("reselection".to_string()) - .spawn(move || { - if let Err(error) = reflector_selector.run() { - let _ = err_tx.send(error); - } - })?; + spawn_worker("reselection", &error_tx, move || reflector_selector.run())?; } // 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(); - thread::Builder::new() - .name("ratecontroller".to_string()) - .spawn(move || { - if let Err(error) = ratecontroller.run() { - let _ = err_tx.send(error); - } - })?; + spawn_worker("ratecontroller", &error_tx, move || ratecontroller.run())?; // Drop the original sender so the channel disconnects if all workers exit cleanly. drop(error_tx); diff --git a/src/reflector_selector.rs b/src/reflector_selector.rs index ab78889..92c1227 100644 --- a/src/reflector_selector.rs +++ b/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, pub reflector_peers_lock: Arc>>, + pub reflector_tx: Sender>, pub reflector_pool: Vec, pub trigger_channel: Receiver, 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; } }