diff --git a/src/baseliner.rs b/src/baseliner.rs deleted file mode 100644 index cb49820..0000000 --- a/src/baseliner.rs +++ /dev/null @@ -1,452 +0,0 @@ -// SPDX-FileCopyrightText: 2022-Present Charles Corrigan mailto:chas-iot@runegate.org (github @chas-iot) -// SPDX-FileCopyrightText: 2022-Present Daniel Lakeland mailto:dlakelan@street-artists.org (github @dlakelan) -// SPDX-FileCopyrightText: 2022-Present Mark Baker mailto:mark@vpost.net (github @Fail-Safe) -// SPDX-FileCopyrightText: 2022-Present Nils Andreas Svee mailto:contact@lochnair.net (github @Lochnair) -// -// SPDX-License-Identifier: MPL-2.0 - -use crate::SHUTDOWN; -use crate::measurement::model::{MeasurementEvent, MeasurementObservation, MeasurementStream}; -use crate::metrics::{Metric, MetricsSender}; -use crate::settings::Settings; -use anyhow::Context; -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::{self, UnboundedReceiver}; - -const STALE_AFTER: Duration = Duration::from_secs(30); -const ANOMALY_THRESHOLD_MS: f64 = 5_000.0; -const BASELINE_HALF_LIFE_SECS: f64 = 135.0; -const RECENT_HALF_LIFE_SECS: f64 = 0.4; - -#[derive(Clone, Copy, Debug)] -pub struct EwmaStats { - pub down: f64, - pub up: f64, -} - -impl EwmaStats { - fn from_reply(reply: &MeasurementObservation) -> Self { - Self { - down: reply - .one_way - .map(|o| o.downlink.as_nanos() as f64 / 1000000.0) - .unwrap_or(reply.rtt_ms() / 2.0), - up: reply - .one_way - .map(|o| o.uplink.as_nanos() as f64 / 1000000.0) - .unwrap_or(reply.rtt_ms() / 2.0), - } - } - - fn update(&mut self, sample: Self, factor: f64) { - self.down = self.down * factor + sample.down * (1.0 - factor); - self.up = self.up * factor + sample.up * (1.0 - factor); - } -} - -#[derive(Clone, Copy, Debug)] -pub struct ReflectorState { - pub baseline: EwmaStats, - pub recent: EwmaStats, - pub last_receive_at: Instant, -} - -impl ReflectorState { - fn new(sample: EwmaStats, received_at: Instant) -> Self { - Self { - baseline: sample, - recent: sample, - last_receive_at: received_at, - } - } - - fn reset(&mut self, sample: EwmaStats, received_at: Instant) { - *self = Self::new(sample, received_at); - } -} - -#[derive(Clone, Debug)] -pub struct ControlSnapshot { - pub reflectors: HashMap, -} - -pub struct Baseliner { - pub settings: Settings, - pub reselect_trigger: mpsc::Sender, - pub start_time: Instant, - pub stats_rx: UnboundedReceiver, - pub control_tx: tokio::sync::watch::Sender>, - pub baseline_metrics: MetricsSender, - pub event_metrics: MetricsSender, - pub ping_metrics: MetricsSender, -} - -fn ewma_factor(tick: f64, dur: f64) -> f64 { - ((0.5_f64).ln() / (dur / tick)).exp() -} - -fn update_reflector_state( - state: &mut ReflectorState, - sample: EwmaStats, - received_at: Instant, - start_time: Instant, - slow_factor: f64, - fast_factor: f64, -) -> bool { - if received_at.saturating_duration_since(state.last_receive_at) > STALE_AFTER { - state.reset(sample, received_at); - } - - if sample.up > state.baseline.up + ANOMALY_THRESHOLD_MS - || sample.down > state.baseline.down + ANOMALY_THRESHOLD_MS - { - state.last_receive_at = start_time; - return true; - } - - state.baseline.update(sample, slow_factor); - state.recent.update(sample, fast_factor); - - state.baseline.down = state.baseline.down.min(state.recent.down); - state.baseline.up = state.baseline.up.min(state.recent.up); - state.last_receive_at = received_at; - - false -} - -fn process_reply( - reflectors: &mut HashMap, - reply: &MeasurementObservation, - start_time: Instant, - slow_factor: f64, - fast_factor: f64, -) -> (ReflectorState, bool) { - let sample = EwmaStats::from_reply(reply); - let state = reflectors - .entry(reply.stream) - .or_insert_with(|| ReflectorState::new(sample, reply.observed_at)); - let anomaly = update_reflector_state( - state, - sample, - reply.observed_at, - start_time, - slow_factor, - fast_factor, - ); - - (*state, anomaly) -} - -fn control_snapshot(reflectors: &HashMap) -> ControlSnapshot { - ControlSnapshot { - reflectors: reflectors.clone(), - } -} - -impl Baseliner { - pub fn run(&mut self) -> anyhow::Result<()> { - let mut reflectors = HashMap::::new(); - - let slow_factor = ewma_factor( - self.settings.advanced_settings.tick_interval, - BASELINE_HALF_LIFE_SECS, - ); - let fast_factor = ewma_factor( - self.settings.advanced_settings.tick_interval, - RECENT_HALF_LIFE_SECS, - ); - - loop { - if SHUTDOWN.load(Ordering::Relaxed) { - info!("Baseliner shutting down"); - return Ok(()); - } - - let event = self - .stats_rx - .blocking_recv() - .context("measurement channel closed")?; - - let reply = match event { - MeasurementEvent::Observation(reply) => reply, - MeasurementEvent::Loss(_) => continue, - }; - - let reflector = reply.stream.peer; - let (state, anomaly) = process_reply( - &mut reflectors, - &reply, - self.start_time, - slow_factor, - fast_factor, - ); - - if anomaly { - info!( - "Reflector {} has OWD > 5 seconds more than baseline, triggering reselection", - reflector - ); - self.event_metrics.send(Metric::Event { - name: "reselection", - reason: "anomaly", - reflector: Some(reflector), - tags: &[], - }); - let _ = self.reselect_trigger.try_send(true); - } - - self.ping_metrics.send(Metric::Ping { - reflector, - measurement_type: reply.stream.source, - rtt: reply.rtt_ms(), - up_time: reply - .one_way - .map(|o| o.uplink.as_nanos() as f64 / 1000000.0) - .unwrap_or(reply.rtt_ms() / 2.0), - down_time: reply - .one_way - .map(|o| o.downlink.as_nanos() as f64 / 1000000.0) - .unwrap_or(reply.rtt_ms() / 2.0), - }); - - self.baseline_metrics.send(Metric::Baseline { - reflector, - baseline_up_ewma: state.baseline.up, - baseline_down_ewma: state.baseline.down, - recent_up_ewma: state.recent.up, - recent_down_ewma: state.recent.down, - }); - - debug!( - "Reflector {} up baseline = {} down baseline = {}", - reflector, state.baseline.up, state.baseline.down - ); - debug!( - "Reflector {} up recent = {} down recent = {}", - reflector, state.recent.up, state.recent.down - ); - - let snapshot = control_snapshot(&reflectors); - self.control_tx.send(Some(snapshot))?; - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::measurement::model::{ - MeasurementSource, MeasurementStream, MeasurementStreamId, OneWayClock, OneWayLatency, - SignedDuration, - }; - use std::net::IpAddr; - use std::time::Duration; - - fn 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 sample(down: f64, up: f64) -> EwmaStats { - EwmaStats { down, up } - } - - fn reply( - reflector: IpAddr, - down_time: f64, - up_time: f64, - received_at: Instant, - ) -> MeasurementObservation { - MeasurementObservation { - stream: MeasurementStream { - id: MeasurementStreamId(0), - peer: reflector, - source: MeasurementSource::IcmpEcho, - }, - started_at: received_at, - observed_at: received_at, - round_trip: None, - one_way: Some(OneWayLatency { - uplink: SignedDuration::from_millis(up_time as i64), - downlink: SignedDuration::from_millis(down_time as i64), - clock: OneWayClock::Uncalibrated, - }), - } - } - - fn assert_close(actual: f64, expected: f64) { - assert!( - (actual - expected).abs() < 1e-9, - "expected {expected}, got {actual}" - ); - } - - #[test] - fn updates_ewmas() { - let start = Instant::now(); - let receive_time = start + Duration::from_secs(1); - let mut state = state(100.0, 100.0, 100.0, 100.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(200.0, 300.0), - receive_time, - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_close(state.baseline.down, 110.0); - assert_close(state.baseline.up, 120.0); - assert_close(state.recent.down, 150.0); - assert_close(state.recent.up, 200.0); - assert_eq!(state.last_receive_at, receive_time); - } - - #[test] - fn keeps_baseline_at_or_below_recent() { - let start = Instant::now(); - let mut state = state(200.0, 200.0, 100.0, 100.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(100.0, 100.0), - start + Duration::from_secs(1), - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_close(state.baseline.down, 100.0); - assert_close(state.baseline.up, 100.0); - assert_close(state.recent.down, 100.0); - assert_close(state.recent.up, 100.0); - } - - #[test] - fn resets_stats_after_a_long_gap() { - let start = Instant::now(); - let receive_time = start + Duration::from_secs(31); - let mut state = state(100.0, 200.0, 300.0, 400.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(500.0, 600.0), - receive_time, - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_close(state.baseline.down, 500.0); - assert_close(state.baseline.up, 600.0); - assert_close(state.recent.down, 500.0); - assert_close(state.recent.up, 600.0); - assert_eq!(state.last_receive_at, receive_time); - } - - #[test] - fn does_not_reset_at_exactly_thirty_seconds() { - let start = Instant::now(); - let receive_time = start + STALE_AFTER; - let mut state = state(100.0, 200.0, 100.0, 200.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(200.0, 300.0), - receive_time, - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_close(state.baseline.down, 110.0); - assert_close(state.baseline.up, 210.0); - assert_close(state.recent.down, 150.0); - assert_close(state.recent.up, 250.0); - } - - #[test] - fn marks_anomalies_for_reselection() { - let start = Instant::now(); - let mut state = state(100.0, 200.0, 300.0, 400.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(5_101.0, 200.0), - start + Duration::from_secs(1), - start, - 0.9, - 0.5, - ); - - assert!(anomaly); - assert_close(state.baseline.down, 100.0); - assert_close(state.baseline.up, 200.0); - assert_close(state.recent.down, 300.0); - assert_close(state.recent.up, 400.0); - assert_eq!(state.last_receive_at, start); - } - - #[test] - fn does_not_mark_exact_anomaly_threshold() { - let start = Instant::now(); - let receive_time = start + Duration::from_secs(1); - let mut state = state(100.0, 200.0, 100.0, 200.0, start); - let anomaly = update_reflector_state( - &mut state, - sample(5_100.0, 5_200.0), - receive_time, - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_eq!(state.last_receive_at, receive_time); - } - - #[test] - fn processed_reply_updates_combined_reflector_state() { - let start = Instant::now(); - let received_at = start + Duration::from_secs(1); - let reflector = "192.0.2.1".parse().unwrap(); - let mut reflectors = HashMap::new(); - let stream = reply(reflector, 0.0, 0.0, start).stream; - reflectors.insert(stream, state(100.0, 100.0, 100.0, 100.0, start)); - - let (updated, anomaly) = process_reply( - &mut reflectors, - &reply(reflector, 200.0, 300.0, received_at), - start, - 0.9, - 0.5, - ); - - assert!(!anomaly); - assert_close(updated.baseline.down, 110.0); - assert_close(updated.baseline.up, 120.0); - assert_close(updated.recent.down, 150.0); - assert_close(updated.recent.up, 200.0); - assert_eq!(updated.last_receive_at, received_at); - assert_close(reflectors[&stream].recent.up, 200.0); - } -} diff --git a/src/controller.rs b/src/controller.rs index 7bc07bf..4af719c 100644 --- a/src/controller.rs +++ b/src/controller.rs @@ -74,7 +74,7 @@ impl Controller { } } - pub async fn run(mut self, initial_reflectors: Vec) -> Result<(), ControllerError> { + pub async fn run(mut self, initial_reflectors: Vec) -> anyhow::Result<()> { self.update_reflectors(initial_reflectors).await?; while let Some(reflectors) = self.reflector_rx.recv().await { self.update_reflectors(reflectors).await?; diff --git a/src/main.rs b/src/main.rs index 2f05eb9..2fc83e5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,37 +7,36 @@ extern crate core; -mod baseliner; mod controller; +mod law; mod log; mod measurement; mod metrics; mod platform; -mod ratecontroller; mod reflector_selector; +mod runtime; mod settings; mod time; mod util; -use crate::baseliner::Baseliner; use crate::controller::Controller; +use crate::law::{Law, LawContext, SqmaLaw}; use crate::measurement::icmp::IcmpEngine; #[cfg(feature = "irtt")] use crate::measurement::irtt::IrttEngine; use crate::measurement::model::{MeasurementSource, MeasurementStream}; use crate::metrics::{Metric, Metrics, MetricsSender}; use crate::platform::{Platform, warn_platform_limitations}; -use crate::ratecontroller::Ratecontroller; use crate::reflector_selector::ReflectorSelector; +use crate::runtime::Runtime; use crate::settings::{ReflectorDefinition, ReflectorProtocol, Settings}; use crate::util::RwLockExt; use ::log::{info, warn}; use std::collections::HashMap; use std::net::{IpAddr, SocketAddr, ToSocketAddrs}; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::atomic::AtomicU32; use std::sync::{Arc, RwLock}; -use std::thread; -use std::thread::{JoinHandle, sleep}; +use std::thread::sleep; use std::time::{Duration, Instant}; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio::sync::watch; @@ -50,15 +49,10 @@ struct ReflectorSetup { reflectors: HashMap, reselection_enabled: bool, } -pub static SHUTDOWN: AtomicBool = AtomicBool::new(false); - -extern "C" fn signal_handler(_: libc::c_int) { - SHUTDOWN.store(true, Ordering::Relaxed); -} const VERSION: &str = env!("CARGO_PKG_VERSION"); -fn initialize_shaper(settings: &Settings, platform: &Platform) -> anyhow::Result<()> { +async fn initialize_shaper(settings: &Settings, platform: &Platform) -> anyhow::Result<()> { platform.validate_shaper(&settings.network.download_interface)?; platform.validate_shaper(&settings.network.upload_interface)?; @@ -81,47 +75,11 @@ fn initialize_shaper(settings: &Settings, platform: &Platform) -> anyhow::Result info!("Sleeping for 2 seconds to give the shaper a chance to control existing bloat"); - sleep(Duration::from_secs(2)); + tokio::time::sleep(Duration::from_secs(2)).await; Ok(()) } -fn install_signal_handlers() { - unsafe { - libc::signal( - libc::SIGINT, - signal_handler as *const () as libc::sighandler_t, - ); - libc::signal( - libc::SIGTERM, - signal_handler as *const () as libc::sighandler_t, - ); - } -} - -fn restore_shaper(settings: &Settings, platform: &Platform) { - info!( - "Requesting restoration to base shaper rates (D/U): {} / {}", - settings.network.download_base_kbits, settings.network.upload_base_kbits, - ); - - if let Err(error) = platform.set_rate( - &settings.network.download_interface, - settings.network.download_base_kbits as u64, - settings.advanced_settings.dry_run, - ) { - warn!("Failed to restore download shaper rate: {error}"); - } - - if let Err(error) = platform.set_rate( - &settings.network.upload_interface, - settings.network.upload_base_kbits as u64, - settings.advanced_settings.dry_run, - ) { - warn!("Failed to restore upload shaper rate: {error}"); - } -} - fn resolve_reflectors( definitions: &[ReflectorDefinition], protocol: ReflectorProtocol, @@ -246,71 +204,15 @@ fn setup_reflectors(settings: &Settings) -> anyhow::Result { }) } -fn spawn_task( - rt: &tokio::runtime::Runtime, - error_tx: &UnboundedSender, - task: F, -) -> tokio::task::JoinHandle<()> -where - F: Future> + Send + 'static, - E: Into + Send + 'static, -{ - let error_tx = error_tx.clone(); +async fn run(settings: Settings) -> anyhow::Result<()> { + let mut workers = tokio::task::JoinSet::>::new(); - rt.spawn(async move { - if let Err(error) = task.await { - let _ = error_tx.send(error.into()); - } - }) -} - -fn spawn_worker( - name: &str, - error_tx: &UnboundedSender, - 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)) -} - -async fn wait_for_exit(error_rx: &mut UnboundedReceiver) -> anyhow::Result<()> { - loop { - match tokio::time::timeout(Duration::from_secs(1), error_rx.recv()).await { - Ok(Some(error)) => { - return Err(anyhow::anyhow!("worker exited with error: {error}")); - } - Ok(None) => { - return Ok(()); - } - Err(_) => { - if SHUTDOWN.load(Ordering::Relaxed) { - info!("Received shutdown signal"); - return Ok(()); - } - } - } - } -} - -fn run(settings: &Settings) -> anyhow::Result<()> { #[cfg(not(feature = "irtt"))] if settings.advanced_settings.measurement_type == ReflectorProtocol::Irtt { anyhow::bail!( "IRTT measurement selected, but this binary was built without the `irtt` feature" ); } - let rt = tokio::runtime::Runtime::new()?; if settings.advanced_settings.dry_run { info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***"); @@ -325,16 +227,22 @@ fn run(settings: &Settings) -> anyhow::Result<()> { pool: reflector_pool, reflectors, reselection_enabled, - } = setup_reflectors(settings)?; + } = setup_reflectors(&settings)?; let (control_snapshot_tx, control_snapshot_rx) = watch::channel(None); - let (error_tx, mut error_rx) = tokio::sync::mpsc::unbounded_channel::(); let (icmp_binding_tx, icmp_binding_rx) = watch::channel(Vec::new()); #[cfg(feature = "irtt")] let (irtt_binding_tx, irtt_binding_rx) = watch::channel(Vec::new()); let (measurement_tx, measurement_rx) = tokio::sync::mpsc::unbounded_channel(); let (reflector_tx, reflector_rx) = tokio::sync::mpsc::unbounded_channel(); let (reselect_tx, reselect_rx) = tokio::sync::mpsc::channel(1); + let (traffic_tx, traffic_rx) = watch::channel(None); + + let platform = Arc::new(Platform::default()); + initialize_shaper(&settings, &platform).await?; + + let runtime = Arc::new(Runtime::new(platform, settings.clone(), traffic_tx)); + workers.spawn(runtime.clone().run()); let source = match settings.advanced_settings.measurement_type { ReflectorProtocol::IcmpEcho => MeasurementSource::IcmpEcho, @@ -352,43 +260,38 @@ fn run(settings: &Settings) -> anyhow::Result<()> { source, Duration::from_secs_f64(settings.advanced_settings.tick_interval), ); - let icmp_engine = { - let _guard = rt.enter(); - IcmpEngine::new( - start_time, - &settings.network, - measurement_tx.clone(), - icmp_binding_rx, - )? - }; + let icmp_engine = IcmpEngine::new( + start_time, + &settings.network, + measurement_tx.clone(), + icmp_binding_rx, + )?; #[cfg(feature = "irtt")] - let (irtt_task, irtt_engine) = { - let _guard = rt.enter(); - IrttEngine::new( - Duration::from_secs_f64(settings.advanced_settings.tick_interval), - &settings.network, - measurement_tx, - irtt_binding_rx, - )? - }; + let (irtt_task, irtt_engine) = IrttEngine::new( + Duration::from_secs_f64(settings.advanced_settings.tick_interval), + &settings.network, + measurement_tx, + irtt_binding_rx, + )?; + + workers.spawn(icmp_engine.run()); - spawn_task(&rt, &error_tx, icmp_engine.run()); #[cfg(feature = "irtt")] - spawn_task(&rt, &error_tx, irtt_engine.run(irtt_task)); + workers.spawn(irtt_engine.run(irtt_task)); let dropped = Arc::new(AtomicU32::new(0)); - let (metrics_tx, metrics_thread_handle) = if settings.observability.enabled { + let metrics_tx = if settings.observability.enabled { let (tx, rx) = tokio::sync::mpsc::channel(1000); let metrics = Metrics { settings: settings.clone(), metrics_rx: rx, metrics_dropped: Arc::clone(&dropped), }; - let handle = spawn_task(&rt, &error_tx, metrics.run()); - (Some(tx), Some(handle)) + workers.spawn(metrics.run()); + Some(tx) } else { - (None, None) + None }; let make_sender = |enabled: bool| -> MetricsSender { @@ -415,26 +318,18 @@ fn run(settings: &Settings) -> anyhow::Result<()> { }, }); - let mut baseliner = Baseliner { - settings: settings.clone(), - reselect_trigger: reselect_tx.clone(), - start_time, - stats_rx: measurement_rx, - control_tx: control_snapshot_tx, - baseline_metrics, - event_metrics: event_metrics.clone(), - ping_metrics, - }; - - let platform = Arc::new(Platform::default()); - initialize_shaper(settings, &platform)?; - - spawn_worker("baseliner", &error_tx, move || baseliner.run())?; - spawn_task( - &rt, - &error_tx, - controller.run(reflector_peers.read_anyhow()?.clone()), + let law = Box::new(SqmaLaw::new(settings.clone(), control_snapshot_tx)); + let law_ctx = LawContext::new( + measurement_rx, + rate_metrics, + reselect_tx, + Arc::clone(&runtime), + Arc::clone(&active_streams), + traffic_rx, ); + workers.spawn(law.run(law_ctx)); + + workers.spawn(controller.run(reflector_peers.read_anyhow()?.clone())); let main_event_metrics = event_metrics.clone(); @@ -449,30 +344,25 @@ fn run(settings: &Settings) -> anyhow::Result<()> { metrics: event_metrics, reflector_tx, }; - spawn_task(&rt, &error_tx, reflector_selector.run()); + workers.spawn(reflector_selector.run()); } - // Give the baseliner time to collect initial samples before adjusting rates. - sleep(Duration::from_secs(10)); + let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?; - let mut ratecontroller = Ratecontroller::new( - settings.clone(), - control_snapshot_rx, - active_streams, - reselect_tx, - rate_metrics, - Arc::clone(&platform), - )?; + let result: anyhow::Result<()> = tokio::select! { + signal = tokio::signal::ctrl_c() => signal.map_err(Into::into), - spawn_worker("ratecontroller", &error_tx, move || ratecontroller.run())?; + _ = sigterm.recv() => Ok(()), - // Drop the original sender so the channel disconnects if all workers exit cleanly. - drop(error_tx); - - let result = rt.block_on(wait_for_exit(&mut error_rx)); - - // Make the error path request shutdown too, rather than leaving the other workers running. - SHUTDOWN.store(true, Ordering::Relaxed); + completed = workers.join_next() => match completed { + Some(Ok(Err(error))) => Err(error), + Some(Ok(Ok(()))) => { + Err(anyhow::anyhow!("Worker exited unexpectedly")) + } + Some(Err(error)) => Err(error.into()), + None => Err(anyhow::anyhow!("No workers remaining")), + }, + }; let stopping_reason = if result.is_err() { "error" } else { "signal" }; main_event_metrics.send(Metric::Event { @@ -482,30 +372,27 @@ fn run(settings: &Settings) -> anyhow::Result<()> { tags: &[], }); + workers.shutdown().await; + // Drop all MetricsSender instances and the raw tx so the metrics channel // disconnects once all worker threads also drop their copies. drop(main_event_metrics); drop(metrics_tx); - if let Some(handle) = metrics_thread_handle { - let _ = rt.block_on(handle); - } - - restore_shaper(settings, &platform); + runtime.restore_shaper()?; result } -fn main() -> anyhow::Result<()> { +#[tokio::main] +async fn main() -> anyhow::Result<()> { println!("Starting sqm-autorate-rust version {}", VERSION); - install_signal_handlers(); - let settings = Settings::load()?; settings.validate()?; log::init(settings.output.log_level)?; warn_platform_limitations(); - run(&settings) + run(settings).await } diff --git a/src/measurement/icmp.rs b/src/measurement/icmp.rs index fa980b8..c474ed1 100644 --- a/src/measurement/icmp.rs +++ b/src/measurement/icmp.rs @@ -432,7 +432,7 @@ impl IcmpEngine { Ok(()) } - pub async fn run(mut self) -> Result<(), IcmpError> { + pub async fn run(mut self) -> anyhow::Result<()> { const HOUSEKEEPING_INTERVAL: Duration = Duration::from_millis(500); let mut housekeeping = tokio::time::interval_at( @@ -480,7 +480,7 @@ impl IcmpEngine { // Not ours, malformed, or late. } - Err(err) => return Err(err), + Err(err) => return Err(err.into()) } } diff --git a/src/measurement/irtt.rs b/src/measurement/irtt.rs index 51957e7..bccbf8e 100644 --- a/src/measurement/irtt.rs +++ b/src/measurement/irtt.rs @@ -13,7 +13,6 @@ use log::{debug, trace, warn}; use thiserror::Error; use tokio::sync::{broadcast::error::RecvError, mpsc::UnboundedSender, watch}; -use crate::SHUTDOWN; use crate::measurement::model::{ MeasurementEvent, MeasurementLoss, MeasurementObservation, MeasurementStream, MeasurementStreamId, OneWayClock, OneWayLatency, SignedDuration, @@ -123,7 +122,7 @@ impl IrttEngine { Ok(()) } - pub async fn run(mut self, task: ManagedClientTask) -> Result<(), IrttError> { + pub async fn run(mut self, task: ManagedClientTask) -> anyhow::Result<()> { let mut task = Box::pin(task); let client = self.client.clone(); @@ -133,7 +132,7 @@ impl IrttEngine { outcome = task.as_mut() => { return match outcome.end_reason { ManagedEndReason::StopRequested => Ok(()), - end_reason => Err(IrttError::ManagedFailure(end_reason)), + end_reason => Err(IrttError::ManagedFailure(end_reason).into()), }; } result = self.run_events() => result, @@ -145,7 +144,7 @@ impl IrttEngine { result?; match outcome.end_reason { ManagedEndReason::StopRequested => Ok(()), - end_reason => Err(IrttError::ManagedFailure(end_reason)), + end_reason => Err(IrttError::ManagedFailure(end_reason).into()), } } @@ -166,9 +165,6 @@ impl IrttEngine { } _ = retry.tick() => { - if SHUTDOWN.load(std::sync::atomic::Ordering::Relaxed) { - return Ok(()); - } if self.client.status().targets.iter().any(|target| { target.desired && target.lifecycle == ManagedTargetLifecycle::Terminal }) { @@ -361,12 +357,6 @@ impl IrttEngine { } } } - - _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => { - if SHUTDOWN.load(std::sync::atomic::Ordering::Relaxed) { - return Ok(()); - } - } } } } diff --git a/src/ratecontroller.rs b/src/ratecontroller.rs deleted file mode 100644 index dbef7e0..0000000 --- a/src/ratecontroller.rs +++ /dev/null @@ -1,486 +0,0 @@ -// SPDX-FileCopyrightText: 2022-Present Charles Corrigan mailto:chas-iot@runegate.org (github @chas-iot) -// SPDX-FileCopyrightText: 2022-Present Daniel Lakeland mailto:dlakelan@street-artists.org (github @dlakelan) -// SPDX-FileCopyrightText: 2022-Present Mark Baker mailto:mark@vpost.net (github @Fail-Safe) -// SPDX-FileCopyrightText: 2022-Present Nils Andreas Svee mailto:contact@lochnair.net (github @Lochnair) -// -// SPDX-License-Identifier: MPL-2.0 - -use crate::SHUTDOWN; -use crate::baseliner::ControlSnapshot; -use crate::measurement::model::MeasurementStream; -use crate::metrics::{Metric, MetricsSender}; -use crate::platform::Platform; -use crate::settings::Settings; -use crate::time::{ClockSample, Realtime}; -use crate::util::{ArcRwLock, RwLockExt}; -use log::{debug, info, warn}; -use std::fs::File; -use std::io::Write; -use std::sync::Arc; -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 { - Down, - Up, -} - -fn generate_initial_speeds(rng: &mut fastrand::Rng, base_speed: f64, size: u32) -> Vec { - let mut rates = Vec::new(); - - for _ in 0..size { - rates.push((rng.f64() * 0.2 + 0.75) * base_speed); - } - - rates -} - -fn get_interface_stats( - platform: &Platform, - upload_interface: &str, -) -> anyhow::Result<(i128, i128)> { - let stats = platform.read_stats(upload_interface)?; - - Ok((stats.rx_bytes.into(), stats.tx_bytes.into())) -} - -#[derive(Clone, Debug)] -struct State { - current_bytes: i128, - current_rate: f64, - delta_stat: f64, - deltas: Vec, - load: f64, - next_rate: f64, - nrate: usize, - previous_bytes: i128, - prev_t: Instant, - safe_rates: Vec, - utilisation: f64, -} - -impl State { - fn new(previous_bytes: i128, safe_rates: Vec) -> Self { - State { - current_bytes: 0, - current_rate: 0.0, - delta_stat: 0.0, - deltas: Vec::new(), - load: 0.0, - next_rate: 0.0, - nrate: 0, - previous_bytes, - prev_t: Instant::now(), - safe_rates, - utilisation: 0.0, - } - } -} - -fn deltas_from_snapshot( - snapshot: Option<&ControlSnapshot>, - reflectors: &[MeasurementStream], - now: Instant, - tick_interval: f64, -) -> (Vec, Vec) { - let mut down = Vec::new(); - let mut up = Vec::new(); - - let Some(snapshot) = snapshot else { - return (down, up); - }; - - for reflector in reflectors { - let Some(state) = snapshot.reflectors.get(reflector) else { - continue; - }; - - if now.duration_since(state.last_receive_at).as_secs_f64() < tick_interval * 2.0 { - let down_delta = state.recent.down - state.baseline.down; - let up_delta = state.recent.up - state.baseline.up; - down.push(down_delta); - up.push(up_delta); - - debug!( - "Reflector: {} ({:?}) down_delay: {} up_delay: {}", - reflector.peer, reflector.source, down_delta, up_delta - ); - } - } - - down.sort_by(|a, b| a.total_cmp(b)); - up.sort_by(|a, b| a.total_cmp(b)); - - (down, up) -} - -pub struct Ratecontroller { - settings: Settings, - control_snapshot_rx: tokio::sync::watch::Receiver>, - reflectors_lock: ArcRwLock>, - reselect_trigger: mpsc::Sender, - state_dl: State, - state_ul: State, - rng: fastrand::Rng, - platform: Arc, - metrics: MetricsSender, -} - -impl Ratecontroller { - pub fn new( - settings: Settings, - control_snapshot_rx: tokio::sync::watch::Receiver>, - reflectors_lock: ArcRwLock>, - reselect_trigger: mpsc::Sender, - metrics: MetricsSender, - platform: Arc, - ) -> anyhow::Result { - platform.validate_shaper(&settings.network.download_interface)?; - platform.validate_shaper(&settings.network.upload_interface)?; - let mut rng = fastrand::Rng::new(); - let dl_safe_rates = generate_initial_speeds( - &mut rng, - settings.network.download_base_kbits, - settings.advanced_settings.speed_hist_size, - ); - let ul_safe_rates = generate_initial_speeds( - &mut rng, - settings.network.upload_base_kbits, - settings.advanced_settings.speed_hist_size, - ); - - let (cur_rx, cur_tx) = get_interface_stats(&platform, &settings.network.upload_interface)?; - - Ok(Self { - settings, - control_snapshot_rx, - reflectors_lock, - reselect_trigger, - state_dl: State::new(cur_rx, dl_safe_rates), - state_ul: State::new(cur_tx, ul_safe_rates), - rng, - platform, - metrics, - }) - } - - fn request_initial_rates(&mut self) -> anyhow::Result<()> { - // Set rates to 60% of base rate to make sure we start with sane baselines. - self.state_dl.current_rate = self.settings.network.download_base_kbits * 0.6; - self.state_ul.current_rate = self.settings.network.upload_base_kbits * 0.6; - - self.platform.set_rate( - &self.settings.network.download_interface, - self.state_dl.current_rate.round() as u64, - self.settings.advanced_settings.dry_run, - )?; - self.platform.set_rate( - &self.settings.network.upload_interface, - self.state_ul.current_rate.round() as u64, - self.settings.advanced_settings.dry_run, - )?; - - Ok(()) - } - - fn calculate_rate(&mut self, direction: Direction) -> anyhow::Result<()> { - let (base_rate, delay_ms, min_rate, state) = if direction == Direction::Down { - ( - self.settings.network.download_base_kbits, - self.settings.advanced_settings.download_delay_ms, - self.settings.network.download_min_kbits(), - &mut self.state_dl, - ) - } else { - ( - self.settings.network.upload_base_kbits, - self.settings.advanced_settings.upload_delay_ms, - self.settings.network.upload_min_kbits(), - &mut self.state_ul, - ) - }; - - let now_t = Instant::now(); - let dur = now_t.duration_since(state.prev_t); - - if !state.deltas.is_empty() { - state.next_rate = state.current_rate; - - state.delta_stat = if state.deltas.len() >= 3 { - state.deltas[2] - } else { - state.deltas[0] - }; - - if state.delta_stat > 0.0 { - /* - * TODO - find where the (8 / 1000) comes from and - * i. convert to a pre-computed factor - * ii. ideally, see if it can be defined in terms of constants, eg ticks per second and number of active reflectors - */ - state.utilisation = (8.0 / 1000.0) - * (state.current_bytes as f64 - state.previous_bytes as f64) - / dur.as_secs_f64(); - state.load = state.utilisation / state.current_rate; - - if state.delta_stat < delay_ms - && state.load > self.settings.advanced_settings.high_load_level - { - state.safe_rates[state.nrate] = (state.current_rate * state.load).floor(); - let max_rate = state - .safe_rates - .iter() - .max_by(|a, b| a.total_cmp(b)) - .unwrap(); - state.next_rate = state.current_rate - * (1.0 + 0.1 * (1.0_f64 - state.current_rate / max_rate).max(0.0)) - + (base_rate * 0.03); - state.nrate += 1; - state.nrate %= self.settings.advanced_settings.speed_hist_size as usize; - } - - if state.delta_stat > delay_ms { - match state - .safe_rates - .get(self.rng.usize(..state.safe_rates.len())) - { - Some(rnd_rate) => { - state.next_rate = rnd_rate.min(0.9 * state.current_rate * state.load); - } - None => { - state.next_rate = 0.9 * state.current_rate * state.load; - } - } - } - } - } - - state.next_rate = state.next_rate.max(min_rate).floor(); - state.previous_bytes = state.current_bytes; - state.prev_t = now_t; - - Ok(()) - } - - pub fn run(&mut self) -> anyhow::Result<()> { - let sleep_time = - Duration::from_secs_f64(self.settings.advanced_settings.min_change_interval); - - let mut lastchg_t = Instant::now(); - let mut lastdump_t = Instant::now(); - - self.request_initial_rates()?; - - let mut speed_hist_fd: Option = None; - let mut speed_hist_fd_inner: File; - let mut stats_fd: Option = None; - let mut stats_fd_inner: File; - - if !self.settings.output.suppress_statistics { - speed_hist_fd_inner = File::options() - .create(true) - .write(true) - .truncate(true) - .open(self.settings.output.speed_hist_file.as_str())?; - - speed_hist_fd_inner.write_all("time,counter,upspeed,downspeed\n".as_bytes())?; - speed_hist_fd_inner.flush()?; - - speed_hist_fd = Some(speed_hist_fd_inner); - - stats_fd_inner = File::options() - .create(true) - .write(true) - .truncate(true) - .open(self.settings.output.stats_file.as_str())?; - - stats_fd_inner.write_all( - "times,timens,rxload,txload,deltadelaydown,deltadelayup,dlrate,uprate\n".as_bytes(), - )?; - stats_fd_inner.flush()?; - - stats_fd = Some(stats_fd_inner); - } - - loop { - if SHUTDOWN.load(Ordering::Relaxed) { - info!("Rate controller shutting down"); - return Ok(()); - } - sleep(sleep_time); - let now_t = Instant::now(); - - if now_t.duration_since(lastchg_t).as_secs_f64() - > self.settings.advanced_settings.min_change_interval - { - // if it's been long enough, and the stats indicate needing to change speeds - // change speeds here - - (self.state_dl.current_bytes, self.state_ul.current_bytes) = - get_interface_stats(&self.platform, &self.settings.network.upload_interface)?; - self.update_deltas()?; - - if self.state_dl.deltas.is_empty() || self.state_ul.deltas.is_empty() { - warn!("No reflector data available, dropping to minimum rates"); - self.metrics.send(Metric::Event { - name: "reflector_unavailable", - reason: "", - reflector: None, - tags: &[], - }); - self.state_dl.next_rate = self.settings.network.download_min_kbits(); - self.state_ul.next_rate = self.settings.network.upload_min_kbits(); - - self.platform.set_rate( - &self.settings.network.download_interface, - self.state_dl.next_rate as u64, - self.settings.advanced_settings.dry_run, - )?; - self.platform.set_rate( - &self.settings.network.upload_interface, - self.state_ul.next_rate as u64, - self.settings.advanced_settings.dry_run, - )?; - - self.state_dl.current_rate = self.state_dl.next_rate; - self.state_ul.current_rate = self.state_ul.next_rate; - continue; - } - - self.calculate_rate(Direction::Down)?; - self.calculate_rate(Direction::Up)?; - - if self.state_dl.next_rate != self.state_dl.current_rate - || self.state_ul.next_rate != self.state_ul.current_rate - { - info!( - "Requesting rates (D/U): {} / {} kbit/s", - self.state_dl.next_rate as u64, self.state_ul.next_rate as u64 - ); - } - - if self.state_dl.next_rate != self.state_dl.current_rate { - self.platform.set_rate( - &self.settings.network.download_interface, - self.state_dl.next_rate as u64, - self.settings.advanced_settings.dry_run, - )?; - } - - if self.state_ul.next_rate != self.state_ul.current_rate { - self.platform.set_rate( - &self.settings.network.upload_interface, - self.state_ul.next_rate as u64, - self.settings.advanced_settings.dry_run, - )?; - } - - self.state_dl.current_rate = self.state_dl.next_rate; - self.state_ul.current_rate = self.state_ul.next_rate; - - let stats_time = ClockSample::::now(); - debug!( - "{},{},{},{},{},{},{},{}", - stats_time.secs(), - stats_time.nsecs(), - self.state_dl.load, - self.state_ul.load, - self.state_dl.delta_stat, - self.state_ul.delta_stat, - self.state_dl.current_rate as u64, - self.state_ul.current_rate as u64 - ); - - self.metrics.send(Metric::Rate { - dl_rate: self.state_dl.current_rate, - ul_rate: self.state_ul.current_rate, - rx_load: self.state_dl.load, - tx_load: self.state_ul.load, - delta_delay_down: self.state_dl.delta_stat, - delta_delay_up: self.state_ul.delta_stat, - }); - - let stats_write_error = stats_fd.as_mut().and_then(|fd| { - fd.write_all( - format!( - "{},{},{},{},{},{},{},{}\n", - stats_time.secs(), - stats_time.nsecs(), - self.state_dl.load, - self.state_ul.load, - self.state_dl.delta_stat, - self.state_ul.delta_stat, - self.state_dl.current_rate as u64, - self.state_ul.current_rate as u64 - ) - .as_bytes(), - ) - .err() - }); - if let Some(e) = stats_write_error { - warn!("Failed to write statistics: {}", e); - } - - lastchg_t = now_t; - } - - if let Some(ref mut fd) = speed_hist_fd - && now_t.duration_since(lastdump_t).as_secs_f64() > 300.0 - { - for i in 0..self.settings.advanced_settings.speed_hist_size as usize { - let hist_time = ClockSample::::now(); - if let Err(e) = fd.write_all( - format!( - "{},{},{},{}\n", - hist_time.as_secs_f64(), - i, - self.state_ul.safe_rates[i], - self.state_dl.safe_rates[i] - ) - .as_bytes(), - ) { - warn!("Failed to write speed history file: {}", e); - } - } - - lastdump_t = now_t; - } - } - } - - fn update_deltas(&mut self) -> anyhow::Result<()> { - let now_t = Instant::now(); - let reflectors = self.reflectors_lock.read_anyhow()?; - 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; - - let minimum = usize::from(self.settings.reflectors.num_reflectors).min(reflectors.len()); - if reflectors.is_empty() - || self.state_dl.deltas.len() < minimum - || self.state_ul.deltas.len() < minimum - { - // trigger reselection - warn!( - "Not enough delta values (D: {}, U: {}, need {}), triggering reselection", - self.state_dl.deltas.len(), - self.state_ul.deltas.len(), - minimum - ); - let _ = self.reselect_trigger.try_send(true); - } - - Ok(()) - } -} diff --git a/src/reflector_selector.rs b/src/reflector_selector.rs index b173358..c03fc60 100644 --- a/src/reflector_selector.rs +++ b/src/reflector_selector.rs @@ -5,8 +5,7 @@ // // SPDX-License-Identifier: MPL-2.0 -use crate::SHUTDOWN; -use crate::baseliner::ControlSnapshot; +use crate::law::ControlSnapshot; use crate::measurement::model::MeasurementStream; use crate::metrics::{Metric, MetricsSender}; use crate::settings::Settings; @@ -53,11 +52,6 @@ impl ReflectorSelector { tokio::time::sleep(baseline_sleep_time).await; loop { - if SHUTDOWN.load(Ordering::Relaxed) { - info!("Reflector selector shutting down"); - return Ok(()); - } - let triggered = tokio::select! { trigger = self.trigger_channel.recv() => { match trigger {