diff --git a/Cargo.lock b/Cargo.lock index c08293f..8ba93b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -695,6 +695,7 @@ name = "sqm-autorate-rust" version = "0.4.1" dependencies = [ "anyhow", + "async-trait", "config", "csv", "fastrand", diff --git a/Cargo.toml b/Cargo.toml index a448e87..5161f4f 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/law/context.rs b/src/law/context.rs new file mode 100644 index 0000000..ce2e64d --- /dev/null +++ b/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, + pub metrics: MetricsSender, + pub reselect_trigger: mpsc::Sender, + runtime: Arc, + pub targets: Arc>>, + pub traffic: watch::Receiver>, +} + +impl LawContext { + pub fn new( + measurements: mpsc::UnboundedReceiver, + metrics: MetricsSender, + reselect_trigger: mpsc::Sender, + runtime: Arc, + targets: Arc>>, + traffic: watch::Receiver>, + ) -> 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) + } +} diff --git a/src/law/mod.rs b/src/law/mod.rs new file mode 100644 index 0000000..c690838 --- /dev/null +++ b/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, ctx: LawContext) -> anyhow::Result<()>; +} diff --git a/src/law/sqma.rs b/src/law/sqma.rs new file mode 100644 index 0000000..9c71330 --- /dev/null +++ b/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, +} + +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, + + download: Option, + upload: Option, + + 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>, +} + +#[async_trait::async_trait] +impl Law for SqmaLaw { + async fn run(mut self: Box, 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>) -> 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, Vec) { + 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, +} + +#[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 + } +} diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 0000000..842230d --- /dev/null +++ b/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, +} + +#[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, + + traffic_interface: String, + sample_interval: Duration, + traffic_tx: watch::Sender>, + + download: Option, + upload: Option, + dry_run: bool, +} + +impl Runtime { + pub(super) fn new( + platform: Arc, + settings: Settings, + traffic_tx: watch::Sender>, + ) -> 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) -> 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(()) + } +}