refactor(workers): supervise async tasks and retire legacy control workers

Activate Runtime and SqmaLaw under an async main with JoinSet supervision and Tokio signal handling. Replace the thread and error-channel scaffolding and atomic shutdown polling with task cancellation, and remove the superseded Baseliner and Ratecontroller implementations.

Cargo.lock

@@ -695,6 +695,7 @@ name = "sqm-autorate-rust"
 version = "0.4.1"
 dependencies = [
  "anyhow",
+ "async-trait",
  "config",
  "csv",
  "fastrand",

Cargo.toml

@@ -20,6 +20,7 @@ strip = "symbols"
 
 [dependencies]
 anyhow = "1.0"
+async-trait = "0.1.92"
 config = { version = "0.15.27", default-features = false }
 csv = "1.4.0"
 fastrand = "2"

src/baseliner.rs

@@ -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<MeasurementStream, ReflectorState>,
-}
-
-pub struct Baseliner {
-    pub settings: Settings,
-    pub reselect_trigger: mpsc::Sender<bool>,
-    pub start_time: Instant,
-    pub stats_rx: UnboundedReceiver<MeasurementEvent>,
-    pub control_tx: tokio::sync::watch::Sender<Option<ControlSnapshot>>,
-    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<MeasurementStream, ReflectorState>,
-    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<MeasurementStream, ReflectorState>) -> ControlSnapshot {
-    ControlSnapshot {
-        reflectors: reflectors.clone(),
-    }
-}
-
-impl Baseliner {
-    pub fn run(&mut self) -> anyhow::Result<()> {
-        let mut reflectors = HashMap::<MeasurementStream, ReflectorState>::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);
-    }
-}

src/controller.rs

@@ -74,7 +74,7 @@ impl Controller {
         }
     }
 
-    pub async fn run(mut self, initial_reflectors: Vec<SocketAddr>) -> Result<(), ControllerError> {
+    pub async fn run(mut self, initial_reflectors: Vec<SocketAddr>) -> anyhow::Result<()> {
         self.update_reflectors(initial_reflectors).await?;
         while let Some(reflectors) = self.reflector_rx.recv().await {
             self.update_reflectors(reflectors).await?;

src/law/context.rs

@@ -0,0 +1,45 @@
+use std::sync::{Arc, RwLock};
+
+use crate::{
+    measurement::model::{MeasurementEvent, MeasurementStream},
+    metrics::MetricsSender,
+    runtime::{Direction, Runtime, Shaper, TrafficRates},
+};
+use tokio::sync::{mpsc, watch};
+
+pub struct LawContext {
+    pub measurements: mpsc::UnboundedReceiver<MeasurementEvent>,
+    pub metrics: MetricsSender,
+    pub reselect_trigger: mpsc::Sender<bool>,
+    runtime: Arc<Runtime>,
+    pub targets: Arc<RwLock<Vec<MeasurementStream>>>,
+    pub traffic: watch::Receiver<Option<TrafficRates>>,
+}
+
+impl LawContext {
+    pub fn new(
+        measurements: mpsc::UnboundedReceiver<MeasurementEvent>,
+        metrics: MetricsSender,
+        reselect_trigger: mpsc::Sender<bool>,
+        runtime: Arc<Runtime>,
+        targets: Arc<RwLock<Vec<MeasurementStream>>>,
+        traffic: watch::Receiver<Option<TrafficRates>>,
+    ) -> Self {
+        Self {
+            measurements,
+            metrics,
+            reselect_trigger,
+            runtime,
+            targets,
+            traffic,
+        }
+    }
+
+    pub fn set_rate(&self, direction: Direction, rate_kbps: f64) -> anyhow::Result<()> {
+        self.runtime.set_rate(direction, rate_kbps)
+    }
+
+    pub fn shaper(&self, direction: Direction) -> Option<&Shaper> {
+        self.runtime.shaper(direction)
+    }
+}

src/law/mod.rs

@@ -0,0 +1,10 @@
+mod context;
+mod sqma;
+
+pub(crate) use context::LawContext;
+pub(crate) use sqma::{ControlSnapshot, SqmaLaw};
+
+#[async_trait::async_trait]
+pub trait Law: Send {
+    async fn run(self: Box<Self>, ctx: LawContext) -> anyhow::Result<()>;
+}

src/law/sqma.rs

@@ -0,0 +1,450 @@
+use std::{
+    collections::{HashMap, VecDeque},
+    time::{Duration, Instant},
+};
+
+use anyhow::Context;
+use log::{debug, warn};
+use tokio::sync::watch;
+
+use crate::{
+    law::{Law, LawContext},
+    measurement::model::{MeasurementEvent, MeasurementObservation, MeasurementStream},
+    metrics::Metric,
+    runtime::{Direction, Shaper},
+    settings::Settings,
+};
+
+const ANOMALY_THRESHOLD_MS: f64 = 5_000.0;
+const BASELINE_HALF_LIFE_SECS: f64 = 135.0;
+const RECENT_HALF_LIFE_SECS: f64 = 0.4;
+const STALE_AFTER: Duration = Duration::from_secs(30);
+
+struct RateState {
+    current_kbps: f64,
+    safe_rates: VecDeque<f64>,
+}
+
+impl RateState {
+    fn new(base_speed: f64, size: u32, rng: &mut fastrand::Rng) -> Self {
+        let mut rates = VecDeque::new();
+
+        for _ in 0..size {
+            rates.push_back((rng.f64() * 0.2 + 0.75) * base_speed);
+        }
+
+        Self {
+            current_kbps: base_speed,
+            safe_rates: rates,
+        }
+    }
+
+    fn calculate_rate(
+        &mut self,
+        shaper: &Shaper,
+        traffic_kbps: f64,
+        deltas: &[f64],
+        delay_threshold_ms: f64,
+        high_load_level: f64,
+        rng: &mut fastrand::Rng,
+    ) -> f64 {
+        let mut next_kbps = self.current_kbps;
+
+        if !deltas.is_empty() {
+            let delta = deltas[deltas.len() / 2];
+
+            if delta > 0.0 {
+                let load = traffic_kbps / self.current_kbps;
+
+                if delta < delay_threshold_ms && load > high_load_level {
+                    self.safe_rates.pop_front();
+                    self.safe_rates
+                        .push_back((self.current_kbps * load).floor());
+
+                    let max_rate = self
+                        .safe_rates
+                        .iter()
+                        .copied()
+                        .max_by(f64::total_cmp)
+                        .unwrap();
+                    next_kbps = self.current_kbps
+                        * (1.0 + 0.1 * (1.0_f64 - self.current_kbps / max_rate).max(0.0))
+                        + (shaper.base_kbps * 0.03);
+                }
+
+                if delta > delay_threshold_ms {
+                    let rnd_rate = self.safe_rates[rng.usize(..self.safe_rates.len())];
+                    next_kbps = rnd_rate.min(0.9 * self.current_kbps * load);
+                }
+            }
+        }
+
+        next_kbps.max(shaper.min_kbps).floor()
+    }
+}
+
+pub(crate) struct SqmaLaw {
+    settings: Settings,
+
+    reflectors: HashMap<MeasurementStream, ReflectorState>,
+
+    download: Option<RateState>,
+    upload: Option<RateState>,
+
+    fast_factor: f64,
+    slow_factor: f64,
+
+    rng: fastrand::Rng,
+
+    // TODO: Eventually the selector should do its own calculations,
+    // and not depend on the law for measurements
+    control_tx: watch::Sender<Option<ControlSnapshot>>,
+}
+
+#[async_trait::async_trait]
+impl Law for SqmaLaw {
+    async fn run(mut self: Box<Self>, mut ctx: LawContext) -> anyhow::Result<()> {
+        self.initialize(&ctx)?;
+
+        let interval = Duration::from_secs_f64(self.settings.advanced_settings.min_change_interval);
+
+        let mut ticker = tokio::time::interval_at(
+            tokio::time::Instant::now() + Duration::from_secs(10),
+            interval,
+        );
+
+        ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
+
+        loop {
+            tokio::select! {
+                event = ctx.measurements.recv() => {
+                    let event = event.context("Measurement channel closed")?;
+                    self.observe(event, &ctx)?;
+                }
+
+                _ = ticker.tick() => {
+                    self.evaluate(&mut ctx)?;
+
+                    self.reflectors.retain(|_, state| {
+                        Instant::now().saturating_duration_since(state.last_receive_at) <= STALE_AFTER
+                    });
+                }
+            }
+        }
+    }
+}
+
+impl SqmaLaw {
+    pub fn new(settings: Settings, control_tx: watch::Sender<Option<ControlSnapshot>>) -> Self {
+        let tick = settings.advanced_settings.tick_interval;
+
+        Self {
+            settings,
+            reflectors: HashMap::new(),
+            download: None,
+            upload: None,
+            fast_factor: (0.5_f64.ln() / (RECENT_HALF_LIFE_SECS / tick)).exp(),
+            slow_factor: (0.5_f64.ln() / (BASELINE_HALF_LIFE_SECS / tick)).exp(),
+            rng: fastrand::Rng::new(),
+            control_tx,
+        }
+    }
+
+    fn initialize(&mut self, ctx: &super::LawContext) -> anyhow::Result<()> {
+        let size = self.settings.advanced_settings.speed_hist_size;
+
+        if let Some(shaper) = ctx.shaper(Direction::Download) {
+            let initial = shaper.base_kbps * 0.6;
+
+            ctx.set_rate(Direction::Download, initial)?;
+
+            let mut state = RateState::new(shaper.base_kbps, size, &mut self.rng);
+            state.current_kbps = initial;
+
+            self.download = Some(state);
+        }
+
+        if let Some(shaper) = ctx.shaper(Direction::Upload) {
+            let initial = shaper.base_kbps * 0.6;
+
+            ctx.set_rate(Direction::Upload, initial)?;
+
+            let mut state = RateState::new(shaper.base_kbps, size, &mut self.rng);
+            state.current_kbps = initial;
+
+            self.upload = Some(state);
+        }
+
+        Ok(())
+    }
+
+    fn deltas(&self, targets: &[MeasurementStream], now: Instant) -> (Vec<f64>, Vec<f64>) {
+        let mut down = Vec::with_capacity(targets.len());
+        let mut up = Vec::with_capacity(targets.len());
+
+        let max_age = Duration::from_secs_f64(self.settings.advanced_settings.tick_interval * 2.0);
+
+        for target in targets {
+            let Some(state) = self.reflectors.get(target) else {
+                continue;
+            };
+
+            if now.saturating_duration_since(state.last_receive_at) >= max_age {
+                continue;
+            }
+
+            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: {}",
+                target.peer, target.source, down_delta, up_delta
+            );
+        }
+
+        down.sort_by(f64::total_cmp);
+        up.sort_by(f64::total_cmp);
+
+        (down, up)
+    }
+
+    fn evaluate(&mut self, ctx: &mut super::LawContext) -> anyhow::Result<()> {
+        let now = Instant::now();
+
+        let (down_deltas, up_deltas, target_count) = {
+            let targets = ctx.targets.read().expect("active streams lock poisoned");
+            let (down, up) = self.deltas(targets.as_slice(), now);
+            (down, up, targets.len())
+        };
+
+        let num_reflectors = self.settings.reflectors.num_reflectors as usize;
+        let minimum = num_reflectors.min(target_count);
+        if target_count == 0 || down_deltas.len() < minimum || up_deltas.len() < minimum {
+            warn!(
+                "Not enough delta values (D: {}, U: {}, need {}), triggering reselection",
+                down_deltas.len(),
+                up_deltas.len(),
+                minimum
+            );
+            let _ = ctx.reselect_trigger.try_send(true);
+        }
+
+        if down_deltas.is_empty() || up_deltas.is_empty() {
+            warn!("No reflector data available, dropping to minimum rates");
+            ctx.metrics.send(Metric::Event {
+                name: "reflector_unavailable",
+                reason: "",
+                reflector: None,
+                tags: &[],
+            });
+
+            for (direction, state) in [
+                (Direction::Download, self.download.as_mut()),
+                (Direction::Upload, self.upload.as_mut()),
+            ] {
+                let (Some(state), Some(shaper)) = (state, ctx.shaper(direction)) else {
+                    continue;
+                };
+
+                ctx.set_rate(direction, shaper.min_kbps)?;
+                state.current_kbps = shaper.min_kbps;
+            }
+
+            return Ok(());
+        }
+
+        let Some(traffic) = *ctx.traffic.borrow() else {
+            // No traffic sample available. Retain current rates.
+            ctx.metrics.send(Metric::Event {
+                name: "traffic_sample_unavailable",
+                reason: "",
+                reflector: None,
+                tags: &[],
+            });
+            return Ok(());
+        };
+
+        for (direction, state, deltas, traffic_kbps, delay_ms) in [
+            (
+                Direction::Download,
+                self.download.as_mut(),
+                down_deltas.as_slice(),
+                traffic.rx_kbps,
+                self.settings.advanced_settings.download_delay_ms,
+            ),
+            (
+                Direction::Upload,
+                self.upload.as_mut(),
+                up_deltas.as_slice(),
+                traffic.tx_kbps,
+                self.settings.advanced_settings.upload_delay_ms,
+            ),
+        ] {
+            let (Some(state), Some(shaper)) = (state, ctx.shaper(direction)) else {
+                continue;
+            };
+
+            let next = state.calculate_rate(
+                shaper,
+                traffic_kbps,
+                deltas,
+                delay_ms,
+                self.settings.advanced_settings.high_load_level,
+                &mut self.rng,
+            );
+
+            if next != state.current_kbps {
+                ctx.set_rate(direction, next)?;
+                state.current_kbps = next;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn observe(&mut self, event: MeasurementEvent, ctx: &super::LawContext) -> anyhow::Result<()> {
+        let observation = match event {
+            MeasurementEvent::Observation(reply) => reply,
+            MeasurementEvent::Loss(_) => return Ok(()),
+        };
+
+        let sample = EwmaStats::from_reply(&observation);
+        let (state, anomaly) = {
+            let state = self
+                .reflectors
+                .entry(observation.stream)
+                .or_insert_with(|| ReflectorState::new(sample, observation.observed_at));
+
+            let anomaly = state.update(
+                sample,
+                observation.observed_at,
+                self.slow_factor,
+                self.fast_factor,
+            );
+
+            (*state, anomaly)
+        };
+
+        if anomaly {
+            self.reflectors.remove(&observation.stream);
+            ctx.metrics.send(Metric::Event {
+                name: "reselection",
+                reason: "anomaly",
+                reflector: Some(observation.stream.peer),
+                tags: &[],
+            });
+
+            let _ = ctx.reselect_trigger.try_send(true);
+        } else {
+            ctx.metrics.send(Metric::Baseline {
+                reflector: observation.stream.peer,
+                baseline_up_ewma: state.baseline.up,
+                baseline_down_ewma: state.baseline.down,
+                recent_up_ewma: state.recent.up,
+                recent_down_ewma: state.recent.down,
+            });
+        }
+
+        self.control_tx.send_replace(Some(ControlSnapshot {
+            reflectors: self.reflectors.clone(),
+        }));
+
+        ctx.metrics.send(Metric::Ping {
+            reflector: observation.stream.peer,
+            measurement_type: observation.stream.source,
+            rtt: observation.rtt_ms(),
+            up_time: observation
+                .one_way
+                .map(|o| o.uplink.as_nanos() as f64 / 1000000.0)
+                .unwrap_or(observation.rtt_ms() / 2.0),
+            down_time: observation
+                .one_way
+                .map(|o| o.downlink.as_nanos() as f64 / 1000000.0)
+                .unwrap_or(observation.rtt_ms() / 2.0),
+        });
+
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct ControlSnapshot {
+    pub reflectors: HashMap<MeasurementStream, ReflectorState>,
+}
+
+#[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);
+    }
+
+    fn update(
+        &mut self,
+        sample: EwmaStats,
+        received_at: Instant,
+        slow_factor: f64,
+        fast_factor: f64,
+    ) -> bool {
+        if received_at.saturating_duration_since(self.last_receive_at) > STALE_AFTER {
+            self.reset(sample, received_at);
+        }
+
+        if sample.up > self.baseline.up + ANOMALY_THRESHOLD_MS
+            || sample.down > self.baseline.down + ANOMALY_THRESHOLD_MS
+        {
+            return true;
+        }
+
+        self.baseline.update(sample, slow_factor);
+        self.recent.update(sample, fast_factor);
+
+        self.baseline.down = self.baseline.down.min(self.recent.down);
+        self.baseline.up = self.baseline.up.min(self.recent.up);
+        self.last_receive_at = received_at;
+
+        false
+    }
+}

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<SocketAddr, ReflectorDefinition>,
     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<ReflectorSetup> {
     })
 }
 
-fn spawn_task<F, E>(
-    rt: &tokio::runtime::Runtime,
-    error_tx: &UnboundedSender<anyhow::Error>,
-    task: F,
-) -> tokio::task::JoinHandle<()>
-where
-    F: Future<Output = Result<(), E>> + Send + 'static,
-    E: Into<anyhow::Error> + Send + 'static,
-{
-    let error_tx = error_tx.clone();
+async fn run(settings: Settings) -> anyhow::Result<()> {
+    let mut workers = tokio::task::JoinSet::<anyhow::Result<()>>::new();
 
-    rt.spawn(async move {
-        if let Err(error) = task.await {
-            let _ = error_tx.send(error.into());
-        }
-    })
-}
-
-fn spawn_worker<F>(
-    name: &str,
-    error_tx: &UnboundedSender<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))
-}
-
-async fn wait_for_exit(error_rx: &mut UnboundedReceiver<anyhow::Error>) -> 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::<anyhow::Error>();
     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
 }

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())
                     }
                 }
 

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(());
-                    }
-                }
             }
         }
     }

src/ratecontroller.rs

@@ -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<f64> {
-    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<f64>,
-    load: f64,
-    next_rate: f64,
-    nrate: usize,
-    previous_bytes: i128,
-    prev_t: Instant,
-    safe_rates: Vec<f64>,
-    utilisation: f64,
-}
-
-impl State {
-    fn new(previous_bytes: i128, safe_rates: Vec<f64>) -> 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<f64>, Vec<f64>) {
-    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<Option<ControlSnapshot>>,
-    reflectors_lock: ArcRwLock<Vec<MeasurementStream>>,
-    reselect_trigger: mpsc::Sender<bool>,
-    state_dl: State,
-    state_ul: State,
-    rng: fastrand::Rng,
-    platform: Arc<Platform>,
-    metrics: MetricsSender,
-}
-
-impl Ratecontroller {
-    pub fn new(
-        settings: Settings,
-        control_snapshot_rx: tokio::sync::watch::Receiver<Option<ControlSnapshot>>,
-        reflectors_lock: ArcRwLock<Vec<MeasurementStream>>,
-        reselect_trigger: mpsc::Sender<bool>,
-        metrics: MetricsSender,
-        platform: Arc<Platform>,
-    ) -> anyhow::Result<Self> {
-        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<File> = None;
-        let mut speed_hist_fd_inner: File;
-        let mut stats_fd: Option<File> = 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::<Realtime>::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::<Realtime>::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(())
-    }
-}

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 {

src/runtime.rs

@@ -0,0 +1,196 @@
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
+use log::{error, warn};
+use tokio::{sync::watch, time};
+
+use crate::{
+    platform::{InterfaceStats, Platform},
+    settings::Settings,
+};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum Direction {
+    Download,
+    Upload,
+}
+
+pub(crate) struct Shaper {
+    interface: String,
+
+    pub(crate) base_kbps: f64,
+    pub(crate) min_kbps: f64,
+    pub(crate) max_kbps: Option<f64>,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct TrafficRates {
+    pub sampled_at: Instant,
+    pub generation: u16,
+    pub rx_kbps: f64,
+    pub tx_kbps: f64,
+}
+
+pub struct Runtime {
+    platform: Arc<Platform>,
+
+    traffic_interface: String,
+    sample_interval: Duration,
+    traffic_tx: watch::Sender<Option<TrafficRates>>,
+
+    download: Option<Shaper>,
+    upload: Option<Shaper>,
+    dry_run: bool,
+}
+
+impl Runtime {
+    pub(super) fn new(
+        platform: Arc<Platform>,
+        settings: Settings,
+        traffic_tx: watch::Sender<Option<TrafficRates>>,
+    ) -> Self {
+        Runtime {
+            platform,
+            traffic_interface: settings.network.upload_interface.clone(),
+            sample_interval: Duration::from_secs_f64(settings.advanced_settings.tick_interval),
+            traffic_tx,
+            download: Some(Shaper {
+                interface: settings.network.download_interface.clone(),
+                base_kbps: settings.network.download_base_kbits,
+                min_kbps: settings.network.download_min_kbits(),
+                max_kbps: None,
+            }),
+            upload: Some(Shaper {
+                interface: settings.network.upload_interface.clone(),
+                base_kbps: settings.network.upload_base_kbits,
+                min_kbps: settings.network.upload_min_kbits(),
+                max_kbps: None,
+            }),
+            dry_run: settings.advanced_settings.dry_run,
+        }
+    }
+
+    pub(super) async fn run(self: Arc<Self>) -> anyhow::Result<()> {
+        let mut generation = 0u16;
+
+        let mut previous = self
+            .platform
+            .read_stats(&self.traffic_interface)
+            .ok()
+            .map(|stats| (Instant::now(), stats));
+
+        let mut ticker = time::interval_at(
+            time::Instant::now() + self.sample_interval,
+            self.sample_interval,
+        );
+
+        ticker.set_missed_tick_behavior(time::MissedTickBehavior::Delay);
+
+        loop {
+            ticker.tick().await;
+            self.sample_traffic(&mut previous, &mut generation);
+        }
+    }
+
+    pub fn restore_shaper(&self) -> anyhow::Result<()> {
+        for direction in [Direction::Download, Direction::Upload] {
+            let Some(shaper) = self.shaper(direction) else {
+                continue;
+            };
+
+            if let Err(e) = self.set_rate(direction, shaper.base_kbps) {
+                error!("Unable to restore shaper for {direction:?}: {e:?}");
+            }
+        }
+
+        Ok(())
+    }
+
+    fn sample_traffic(
+        &self,
+        previous: &mut Option<(Instant, InterfaceStats)>,
+        generation: &mut u16,
+    ) {
+        let stats = match self.platform.read_stats(&self.traffic_interface) {
+            Ok(stats) => stats,
+            Err(err) => {
+                log::warn!("Failed to read interface statistics: {err}");
+
+                if previous.take().is_some() {
+                    *generation = generation.wrapping_add(1);
+                }
+
+                self.traffic_tx.send_replace(None);
+                return;
+            }
+        };
+
+        let now = Instant::now();
+
+        let Some((last_at, last_stats)) = previous.replace((now, stats)) else {
+            self.traffic_tx.send_replace(None);
+            return;
+        };
+
+        if stats.rx_bytes < last_stats.rx_bytes || stats.tx_bytes < last_stats.tx_bytes {
+            *generation = generation.wrapping_add(1);
+            self.traffic_tx.send_replace(None);
+            return;
+        }
+
+        let elapsed = now.duration_since(last_at).as_secs_f64();
+
+        const KBITS_PER_BYTE: f64 = 8.0 / 1_000.0;
+
+        let rates = TrafficRates {
+            sampled_at: now,
+            generation: *generation,
+            rx_kbps: (stats.rx_bytes - last_stats.rx_bytes) as f64 * KBITS_PER_BYTE / elapsed,
+            tx_kbps: (stats.tx_bytes - last_stats.tx_bytes) as f64 * KBITS_PER_BYTE / elapsed,
+        };
+
+        self.traffic_tx.send_replace(Some(rates));
+    }
+
+    pub(super) fn shaper(&self, direction: Direction) -> Option<&Shaper> {
+        match direction {
+            Direction::Download => self.download.as_ref(),
+            Direction::Upload => self.upload.as_ref(),
+        }
+    }
+
+    pub(super) fn set_rate(&self, direction: Direction, rate_kbps: f64) -> anyhow::Result<()> {
+        let shaper = self
+            .shaper(direction)
+            .ok_or_else(|| anyhow::anyhow!("No shaper configured for {direction:?}"))?;
+
+        anyhow::ensure!(
+            rate_kbps.is_finite() && rate_kbps >= 0.0,
+            "Invalid requested shaper rate: {rate_kbps}"
+        );
+
+        if let Some(max_kbps) = shaper.max_kbps
+            && rate_kbps > max_kbps
+        {
+            warn!(
+                "Law requested {direction:?} rate above configured maximum: \
+                 {rate_kbps} > {max_kbps} kbps"
+            );
+        }
+
+        if rate_kbps < shaper.min_kbps {
+            warn!(
+                "Law requested {direction:?} rate below configured minimum: \
+                 {rate_kbps} < {} kbps",
+                shaper.min_kbps
+            );
+        }
+
+        self.platform
+            .set_rate(&shaper.interface, rate_kbps.round() as u64, self.dry_run)?;
+
+        Ok(())
+    }
+}