From b56c8b72cc4f48850b60fd220d90ca9d8bb5f6d0 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 14 Jul 2026 22:53:01 +0700 Subject: [PATCH] refactor(main): modularize main loop and improve resource cleanup logic --- src/main.rs | 365 +++++++++++++++++++++++++++++++--------------------- 1 file changed, 215 insertions(+), 150 deletions(-) diff --git a/src/main.rs b/src/main.rs index 4b6a83c..8b33217 100644 --- a/src/main.rs +++ b/src/main.rs @@ -22,7 +22,7 @@ mod util; use crate::baseliner::{Baseliner, ReflectorStats}; use crate::metrics::{Metric, Metrics, MetricsSender}; -use ::log::{debug, info}; +use ::log::{debug, info, warn}; use flume::RecvTimeoutError; use std::collections::HashMap; use std::net::IpAddr; @@ -30,22 +30,30 @@ use std::str::FromStr; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::thread::sleep; -use std::time::Duration; -use std::time::Instant; +use std::time::{Duration, Instant}; use std::{process, thread}; use crate::config::{Config, MeasurementType}; -use crate::netlink::Netlink; +use crate::netlink::{Netlink, Qdisc}; use crate::pinger::{InFlightProbeCache, PingListener, PingSender}; use crate::pinger_icmp::{PingerICMPEchoListener, PingerICMPEchoSender}; use crate::pinger_icmp_ts::{PingerICMPTimestampListener, PingerICMPTimestampSender}; use crate::ratecontroller::{Ratecontroller, StatsDirection}; use crate::reflector_selector::ReflectorSelector; +struct ReflectorSetup { + peers: Arc>>, + pool: Vec, + reselection_enabled: bool, + active_count: usize, +} + +type PingListenerBox = Box; +type PingSenderBox = Box; + pub static SHUTDOWN: AtomicBool = AtomicBool::new(false); extern "C" fn signal_handler(_: libc::c_int) { - println!("Signal received, shutting down sqm-autorate-rust"); SHUTDOWN.store(true, Ordering::Relaxed); } @@ -77,9 +85,47 @@ fn compute_inflight_probe_capacity( .max(INFLIGHT_CAPACITY_MIN) } -fn main() -> anyhow::Result<()> { - println!("Starting sqm-autorate-rust version {}", VERSION); +fn create_pinger(measurement_type: MeasurementType) -> (PingListenerBox, PingSenderBox) { + match measurement_type { + MeasurementType::Icmp => ( + Box::new(PingerICMPEchoListener {}), + Box::new(PingerICMPEchoSender {}), + ), + MeasurementType::IcmpTimestamps => ( + Box::new(PingerICMPTimestampListener {}), + Box::new(PingerICMPTimestampSender {}), + ), + MeasurementType::Ntp | MeasurementType::TcpTimestamps => { + todo!() + } + } +} +fn initialize_shaper(config: &Config) -> anyhow::Result<(Qdisc, Qdisc)> { + let down = Netlink::qdisc_from_ifname(&config.download_interface)?; + let up = Netlink::qdisc_from_ifname(&config.upload_interface)?; + + info!( + "Setting shaper rates to minimum (D/L): {} / {}", + config.download_min_kbits, config.upload_min_kbits, + ); + + Netlink::set_qdisc_rate(down, config.download_min_kbits as u64, config.dry_run)?; + Netlink::set_qdisc_rate(up, config.upload_min_kbits as u64, config.dry_run)?; + + let settle_time = Duration::from_secs(2); + + info!( + "Sleeping for {} seconds to give the shaper a chance to control existing bloat", + settle_time.as_secs_f64(), + ); + + sleep(settle_time); + + Ok((down, up)) +} + +fn install_signal_handlers() { unsafe { libc::signal( libc::SIGINT, @@ -90,25 +136,54 @@ fn main() -> anyhow::Result<()> { signal_handler as *const () as libc::sighandler_t, ); } +} - let config = Config::new()?; - log::init(config.log_level)?; - if config.dry_run { - info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***"); +fn interface_stats_directions(config: &Config) -> (StatsDirection, StatsDirection) { + let download = if config.download_interface.starts_with("ifb") + || config.download_interface.starts_with("veth") + { + StatsDirection::TX + } else { + StatsDirection::RX + }; + + let upload = if config.upload_interface.starts_with("ifb") + || config.upload_interface.starts_with("veth") + { + StatsDirection::RX + } else { + StatsDirection::TX + }; + + (download, upload) +} + +fn probe_identifier() -> u16 { + // The ICMP identifier is only two bytes. + (process::id() & 0xffff) as u16 +} + +fn restore_shaper(config: &Config, down: Qdisc, up: Qdisc) { + info!( + "Restoring base shaper rates (D/L): {} / {}", + config.download_base_kbits, config.upload_base_kbits, + ); + + if let Err(error) = + Netlink::set_qdisc_rate(down, config.download_base_kbits as u64, config.dry_run) + { + warn!("Failed to restore download shaper rate: {error}"); } - let mut reflectors = config.load_reflectors()?; - let start_t = Instant::now(); - // The identifier field in ICMP is only 2 bytes - // so take the last 2 bytes of the PID as the ID - let id = (process::id() & 0xFFFF) as u16; + if let Err(error) = Netlink::set_qdisc_rate(up, config.upload_base_kbits as u64, config.dry_run) + { + warn!("Failed to restore upload shaper rate: {error}"); + } +} - // Create data structures shared by different threads - let owd_baseline = Arc::new(Mutex::new(HashMap::::new())); - let owd_recent = Arc::new(Mutex::new(HashMap::::new())); - let reflector_peers_lock = Arc::new(RwLock::new(Vec::::new())); - let mut reflector_pool = Vec::::new(); - let reflector_pool_size = reflectors.len(); +fn setup_reflectors(config: &Config) -> anyhow::Result { + let configured = config.load_reflectors()?; + let configured_count = configured.len(); let default_reflectors = [ IpAddr::from_str("9.9.9.9")?, @@ -119,17 +194,57 @@ fn main() -> anyhow::Result<()> { IpAddr::from_str("94.140.14.14")?, ]; - match reflector_pool_size > config.num_reflectors as usize { - true => { - let mut peers = reflector_peers_lock.write().unwrap(); - peers.extend_from_slice(&default_reflectors); - reflector_pool.append(reflectors.as_mut()); - } - false => { - let mut peers = reflector_peers_lock.write().unwrap(); - peers.extend_from_slice(&default_reflectors); + let reselection_enabled = configured_count > config.num_reflectors as usize; + let pool = if reselection_enabled { + configured + } else { + Vec::new() + }; + + Ok(ReflectorSetup { + peers: Arc::new(RwLock::new(default_reflectors.to_vec())), + pool, + reselection_enabled, + active_count: default_reflectors.len().max(config.num_reflectors as usize), + }) +} + +fn wait_for_exit(error_rx: &flume::Receiver) -> anyhow::Result<()> { + loop { + match error_rx.recv_timeout(Duration::from_secs(1)) { + Ok(error) => { + return Err(anyhow::anyhow!("thread exited with error: {error}")); + } + Err(RecvTimeoutError::Disconnected) => { + return Ok(()); + } + Err(RecvTimeoutError::Timeout) => { + if SHUTDOWN.load(Ordering::Relaxed) { + info!("Received shutdown signal"); + return Ok(()); + } + } } } +} + +fn run(config: Config) -> anyhow::Result<()> { + if config.dry_run { + info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***"); + } + + let start_time = Instant::now(); + let probe_id = probe_identifier(); + + let ReflectorSetup { + peers: reflector_peers, + pool: reflector_pool, + reselection_enabled, + active_count: active_reflector_count, + } = setup_reflectors(&config)?; + + let owd_baseline = Arc::new(Mutex::new(HashMap::::new())); + let owd_recent = Arc::new(Mutex::new(HashMap::::new())); let (baseliner_stats_tx, baseliner_stats_rx) = flume::unbounded(); let (error_tx, error_rx) = flume::unbounded::(); @@ -148,8 +263,8 @@ fn main() -> anyhow::Result<()> { let handle = thread::Builder::new() .name("metrics".to_string()) .spawn(move || { - if let Err(e) = metrics.run() { - let _ = err_tx.send(e); + if let Err(error) = metrics.run() { + let _ = err_tx.send(error); } })?; (Some(tx), Some(handle)) @@ -181,79 +296,47 @@ fn main() -> anyhow::Result<()> { }, }); - let (mut ping_listener, mut ping_sender) = match config.measurement_type { - MeasurementType::Icmp => ( - Box::new(PingerICMPEchoListener {}) as Box, - Box::new(PingerICMPEchoSender {}) as Box, - ), - MeasurementType::IcmpTimestamps => ( - Box::new(PingerICMPTimestampListener {}) as Box, - Box::new(PingerICMPTimestampSender {}) as Box, - ), - MeasurementType::Ntp | MeasurementType::TcpTimestamps => { - todo!() - } - }; - let reselection_enabled = reflector_pool_size > config.num_reflectors as usize; - let active_reflector_count = default_reflectors.len().max(config.num_reflectors as usize); + let (mut ping_listener, mut ping_sender) = create_pinger(config.measurement_type); + let inflight_probe_capacity = compute_inflight_probe_capacity(&config, active_reflector_count, reselection_enabled); info!( "In-flight probe cache capacity: {} (active_reflectors={}, reselection_enabled={}, tick_interval_s={})", inflight_probe_capacity, active_reflector_count, reselection_enabled, config.tick_interval ); + let inflight: InFlightProbeCache = Arc::new(Mutex::new(HashMap::with_capacity(inflight_probe_capacity))); let baseliner = Baseliner { config: config.clone(), - owd_baseline: owd_baseline.clone(), - owd_recent: owd_recent.clone(), + owd_baseline: Arc::clone(&owd_baseline), + owd_recent: Arc::clone(&owd_recent), reselect_trigger: reselect_tx.clone(), - start_time: start_t, + start_time, stats_rx: baseliner_stats_rx, baseline_metrics, event_metrics: event_metrics.clone(), }; - let down_qdisc = Netlink::qdisc_from_ifname(config.download_interface.as_str())?; - let up_qdisc = Netlink::qdisc_from_ifname(config.upload_interface.as_str())?; - - /* Set initial TC values to minimum - * so there should be no initial bufferbloat to - * fool the baseliner - */ - info!( - "Setting shaper rates to minimum (D/L): {} / {}", - config.download_min_kbits, config.upload_min_kbits - ); - Netlink::set_qdisc_rate(down_qdisc, config.download_min_kbits as u64, config.dry_run)?; - Netlink::set_qdisc_rate(up_qdisc, config.upload_min_kbits as u64, config.dry_run)?; - - // Sleep for a few seconds to give the shaper a chance - // to control the queue if load is heavy - let settle_sleep_time = Duration::new(2, 0); - info!( - "Sleeping for {} to give the shaper a chance to get in control if there's bloat", - settle_sleep_time.as_secs_f64() - ); - sleep(settle_sleep_time); + let (down_qdisc, up_qdisc) = initialize_shaper(&config)?; let err_tx = error_tx.clone(); - let reflector_peers_lock_clone = reflector_peers_lock.clone(); - let inflight_listener = inflight.clone(); + let listener_peers = Arc::clone(&reflector_peers); + let listener_inflight = Arc::clone(&inflight); + let measurement_type = config.measurement_type; thread::Builder::new() .name("receiver".to_string()) .spawn(move || { - if let Err(e) = ping_listener.listen( - id, - config.measurement_type, - reflector_peers_lock_clone, - inflight_listener, + if let Err(error) = ping_listener.listen( + probe_id, + measurement_type, + listener_peers, + listener_inflight, baseliner_stats_tx, ping_metrics, ) { - let _ = err_tx.send(e); + let _ = err_tx.send(error); } })?; @@ -261,25 +344,27 @@ fn main() -> anyhow::Result<()> { thread::Builder::new() .name("baseliner".to_string()) .spawn(move || { - if let Err(e) = baseliner.run() { - let _ = err_tx.send(e); + if let Err(error) = baseliner.run() { + let _ = err_tx.send(error); } })?; let err_tx = error_tx.clone(); - let reflector_peers_lock_clone = reflector_peers_lock.clone(); - let inflight_sender = inflight.clone(); + let sender_peers = Arc::clone(&reflector_peers); + let sender_inflight = Arc::clone(&inflight); + let measurement_type = config.measurement_type; + let tick_interval = config.tick_interval; thread::Builder::new() .name("sender".to_string()) .spawn(move || { - if let Err(e) = ping_sender.send( - id, - config.measurement_type, - reflector_peers_lock_clone, - inflight_sender, - config.tick_interval, + if let Err(error) = ping_sender.send( + probe_id, + measurement_type, + sender_peers, + sender_inflight, + tick_interval, ) { - let _ = err_tx.send(e); + let _ = err_tx.send(error); } })?; @@ -288,8 +373,8 @@ fn main() -> anyhow::Result<()> { if reselection_enabled { let reflector_selector = ReflectorSelector { config: config.clone(), - owd_recent: owd_recent.clone(), - reflector_peers_lock: reflector_peers_lock.clone(), + owd_recent: Arc::clone(&owd_recent), + reflector_peers_lock: Arc::clone(&reflector_peers), reflector_pool, trigger_channel: reselect_rx, metrics: event_metrics, @@ -298,75 +383,53 @@ fn main() -> anyhow::Result<()> { thread::Builder::new() .name("reselection".to_string()) .spawn(move || { - if let Err(e) = reflector_selector.run() { - let _ = err_tx.send(e); + if let Err(error) = reflector_selector.run() { + let _ = err_tx.send(error); } })?; } - // Sleep 10 seconds before we start adjusting speeds - sleep(Duration::new(10, 0)); + // Give the baseliner time to collect initial samples before adjusting rates. + sleep(Duration::from_secs(10)); - let dl_direction = if config.download_interface.starts_with("ifb") - || config.download_interface.starts_with("veth") - { - StatsDirection::TX - } else { - StatsDirection::RX - }; - let ul_direction = if config.upload_interface.starts_with("ifb") - || config.upload_interface.starts_with("veth") - { - StatsDirection::RX - } else { - StatsDirection::TX - }; + let (dl_direction, ul_direction) = interface_stats_directions(&config); + + debug!( + "Download direction: {}:{:?}", + config.download_interface, dl_direction + ); + debug!( + "Upload direction: {}:{:?}", + config.upload_interface, ul_direction + ); let mut ratecontroller = Ratecontroller::new( config.clone(), owd_baseline, owd_recent, - reflector_peers_lock, + reflector_peers, reselect_tx, dl_direction, ul_direction, rate_metrics, )?; - debug!( - "Download direction: {}:{:?}", - config.download_interface, dl_direction - ); - - debug!( - "Upload direction: {}:{:?}", - config.upload_interface, ul_direction - ); - let err_tx = error_tx.clone(); thread::Builder::new() .name("ratecontroller".to_string()) .spawn(move || { - if let Err(e) = ratecontroller.run() { - let _ = err_tx.send(e); + if let Err(error) = ratecontroller.run() { + let _ = err_tx.send(error); } })?; - // Drop original sender so the channel disconnects if all threads exit cleanly + // Drop the original sender so the channel disconnects if all workers exit cleanly. drop(error_tx); - let result = loop { - match error_rx.recv_timeout(Duration::from_secs(1)) { - Ok(e) => break Err(anyhow::anyhow!("thread exited with error: {e}")), - Err(RecvTimeoutError::Disconnected) => break Ok(()), - Err(RecvTimeoutError::Timeout) => { - if SHUTDOWN.load(Ordering::Relaxed) { - info!("Received shutdown signal"); - break Ok(()); - } - } - } - }; + let result = wait_for_exit(&error_rx); + + // Make the error path request shutdown too, rather than leaving the other workers running. + SHUTDOWN.store(true, Ordering::Relaxed); let stopping_reason = if result.is_err() { "error" } else { "signal" }; main_event_metrics.send(Metric::Event { @@ -377,24 +440,26 @@ fn main() -> anyhow::Result<()> { }); // Drop all MetricsSender instances and the raw tx so the metrics channel - // disconnects once all threads also drop their copies, allowing the - // metrics thread to drain and exit cleanly. + // disconnects once all worker threads also drop their copies. drop(main_event_metrics); drop(metrics_tx); + if let Some(handle) = metrics_thread_handle { let _ = handle.join(); } - info!( - "Restoring base shaper rates (D/L): {} / {}", - config.download_base_kbits, config.upload_base_kbits - ); - let _ = Netlink::set_qdisc_rate( - down_qdisc, - config.download_base_kbits as u64, - config.dry_run, - ); - let _ = Netlink::set_qdisc_rate(up_qdisc, config.upload_base_kbits as u64, config.dry_run); + restore_shaper(&config, down_qdisc, up_qdisc); result } + +fn main() -> anyhow::Result<()> { + println!("Starting sqm-autorate-rust version {}", VERSION); + + install_signal_handlers(); + + let config = Config::new()?; + log::init(config.log_level)?; + + run(config) +}