From 708d88fc57baa5f85df5ec988f792f3da52acd6e Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Thu, 24 Sep 2026 17:21:48 +0700 Subject: [PATCH] fix(shaper): restore base rates and recover recreated CAKE --- src/baseliner.rs | 11 ++- src/main.rs | 138 ++++++++++++++++++++++++++++++---- src/platform/linux/netlink.rs | 50 +++++++++++- src/ratecontroller.rs | 65 ++++++++++++++-- 4 files changed, 240 insertions(+), 24 deletions(-) diff --git a/src/baseliner.rs b/src/baseliner.rs index 1f05116..a078489 100644 --- a/src/baseliner.rs +++ b/src/baseliner.rs @@ -9,7 +9,7 @@ use crate::SHUTDOWN; use crate::metrics::{Metric, MetricsSender}; use crate::pinger::PingReply; use crate::settings::Settings; -use flume::{Receiver, Sender}; +use flume::{Receiver, RecvTimeoutError, Sender}; use log::{debug, info}; use std::collections::HashMap; use std::net::IpAddr; @@ -165,7 +165,14 @@ impl Baseliner { return Ok(()); } - let reply = self.stats_rx.recv()?; + let reply = match self.stats_rx.recv_timeout(Duration::from_millis(500)) { + Ok(reply) => reply, + Err(RecvTimeoutError::Timeout) => continue, + Err(RecvTimeoutError::Disconnected) if SHUTDOWN.load(Ordering::Relaxed) => { + return Ok(()); + } + Err(error) => return Err(error.into()), + }; let reflector = reply.reflector; let (state, anomaly) = process_reply( &mut reflectors, diff --git a/src/main.rs b/src/main.rs index ad6c000..e581b0f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -121,16 +121,24 @@ fn initialize_shaper_with_settle_time( settings.network.upload_min_kbits(), ); - traffic_control.set_rate( + if let Err(error) = traffic_control.set_rate( &down, settings.network.download_min_kbits() as u64, settings.advanced_settings.dry_run, - )?; - traffic_control.set_rate( + ) { + SHUTDOWN.store(true, Ordering::Relaxed); + restore_shaper(settings, traffic_control, &down, &up); + return Err(error.into()); + } + if let Err(error) = traffic_control.set_rate( &up, settings.network.upload_min_kbits() as u64, settings.advanced_settings.dry_run, - )?; + ) { + SHUTDOWN.store(true, Ordering::Relaxed); + restore_shaper(settings, traffic_control, &down, &up); + return Err(error.into()); + } info!( "Sleeping for {} seconds to give the shaper a chance to control existing bloat", @@ -188,6 +196,20 @@ fn restore_shaper( } } +struct ShaperGuard<'a, T: TrafficControlBackend> { + settings: &'a Settings, + traffic_control: &'a mut T, + down: T::Handle, + up: T::Handle, +} + +impl Drop for ShaperGuard<'_, T> { + fn drop(&mut self) { + SHUTDOWN.store(true, Ordering::Relaxed); + restore_shaper(self.settings, self.traffic_control, &self.down, &self.up); + } +} + fn setup_reflectors(settings: &Settings) -> anyhow::Result { let configured = settings.load_reflectors()?; let configured_count = configured.len(); @@ -338,6 +360,12 @@ fn run(settings: &Settings) -> anyhow::Result<()> { let mut main_traffic_control = traffic_control_backend(); let (down_shaper, up_shaper) = initialize_shaper(settings, &mut main_traffic_control)?; + let shaper_guard = ShaperGuard { + settings, + traffic_control: &mut main_traffic_control, + down: down_shaper, + up: up_shaper, + }; let err_tx = error_tx.clone(); let listener_peers = Arc::clone(&reflector_peers); @@ -409,7 +437,13 @@ fn run(settings: &Settings) -> anyhow::Result<()> { } // Give the baseliner time to collect initial samples before adjusting rates. - sleep(Duration::from_secs(10)); + let baseline_start = Instant::now(); + while baseline_start.elapsed() < Duration::from_secs(10) { + if SHUTDOWN.load(Ordering::Relaxed) { + return Ok(()); + } + sleep(Duration::from_millis(100)); + } let mut ratecontroller = Ratecontroller::new( settings.clone(), @@ -420,7 +454,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { )?; let err_tx = error_tx.clone(); - thread::Builder::new() + let ratecontroller_thread_handle = thread::Builder::new() .name("ratecontroller".to_string()) .spawn(move || { if let Err(error) = ratecontroller.run() { @@ -449,17 +483,13 @@ fn run(settings: &Settings) -> anyhow::Result<()> { drop(main_event_metrics); drop(metrics_tx); + let _ = ratecontroller_thread_handle.join(); + drop(shaper_guard); + if let Some(handle) = metrics_thread_handle { let _ = handle.join(); } - restore_shaper( - settings, - &mut main_traffic_control, - &down_shaper, - &up_shaper, - ); - result } @@ -476,3 +506,85 @@ fn main() -> anyhow::Result<()> { run(&settings) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::settings::{ + AdvancedSettings, NetworkSettings, ObservabilitySettings, OutputSettings, + }; + + #[derive(Default)] + struct FakeControl { + calls: Vec<(String, u64)>, + fail_upload_min: bool, + } + + impl TrafficControlBackend for FakeControl { + type Error = std::io::Error; + type Handle = String; + + fn find_shaper(&mut self, interface: &str) -> Result { + Ok(interface.into()) + } + + fn set_rate( + &mut self, + shaper: &Self::Handle, + rate: u64, + _: bool, + ) -> Result<(), Self::Error> { + self.calls.push((shaper.clone(), rate)); + if self.fail_upload_min && shaper == "eth0" && rate == 1_000 { + return Err(std::io::Error::other("simulated startup failure")); + } + Ok(()) + } + } + + #[test] + fn startup_failures_restore_both_base_rates() { + let settings = Settings { + network: NetworkSettings { + download_interface: "ifb0".into(), + upload_interface: "eth0".into(), + download_base_kbits: 10_000.0, + download_min_percent: 20.0, + upload_base_kbits: 5_000.0, + upload_min_percent: 20.0, + }, + output: OutputSettings::default(), + observability: ObservabilitySettings::default(), + advanced_settings: AdvancedSettings::default(), + }; + let mut control = FakeControl { + fail_upload_min: true, + ..FakeControl::default() + }; + assert!( + initialize_shaper_with_settle_time(&settings, &mut control, Duration::ZERO).is_err() + ); + SHUTDOWN.store(false, Ordering::Relaxed); + assert_eq!( + control.calls[2..], + [("ifb0".into(), 10_000), ("eth0".into(), 5_000)] + ); + + control.calls.clear(); + control.fail_upload_min = false; + let (down, up) = + initialize_shaper_with_settle_time(&settings, &mut control, Duration::ZERO).unwrap(); + let guard = ShaperGuard { + settings: &settings, + traffic_control: &mut control, + down, + up, + }; + drop(guard); + SHUTDOWN.store(false, Ordering::Relaxed); + assert_eq!( + control.calls[2..], + [("ifb0".into(), 10_000), ("eth0".into(), 5_000)] + ); + } +} diff --git a/src/platform/linux/netlink.rs b/src/platform/linux/netlink.rs index 5742979..31af174 100644 --- a/src/platform/linux/netlink.rs +++ b/src/platform/linux/netlink.rs @@ -5,12 +5,16 @@ use std::io; use std::str::Utf8Error; -use log::info; +use log::{info, warn}; use netlink_bindings::{rt_link, tc}; use netlink_socket2::NetlinkSocket; +use std::thread::sleep; +use std::time::{Duration, Instant}; use thiserror::Error; +use crate::SHUTDOWN; use crate::platform::{InterfaceStats, InterfaceStatsProvider, TrafficControlBackend}; +use std::sync::atomic::Ordering; #[derive(Debug, Error)] pub(crate) enum NetlinkError { @@ -36,12 +40,29 @@ pub(crate) enum NetlinkError { Utf8Error(#[from] Utf8Error), } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct Qdisc { + ifname: String, ifindex: i32, parent: u32, } +impl NetlinkError { + fn is_qdisc_loss(&self) -> bool { + match self { + Self::InterfaceNotFound(_) | Self::NoQdiscFound(_) => true, + Self::Netlink(error) => { + matches!(error.raw_os_error(), Some(libc::ENODEV | libc::ENOENT)) + } + Self::Reply(error) => matches!( + error.as_io_error().raw_os_error(), + Some(libc::ENODEV | libc::ENOENT) + ), + _ => false, + } + } +} + #[derive(Debug, Default)] pub(crate) struct Netlink {} @@ -96,6 +117,7 @@ impl Netlink { == "cake" { return Ok(Qdisc { + ifname: ifname.to_owned(), ifindex, parent: header.parent, }); @@ -111,7 +133,7 @@ impl Netlink { } fn set_qdisc_rate( - qdisc: Qdisc, + qdisc: &Qdisc, bandwidth_kbit: u64, dry_run: bool, ) -> Result<(), NetlinkError> { @@ -167,6 +189,26 @@ impl TrafficControlBackend for Netlink { bandwidth_kbit: u64, dry_run: bool, ) -> Result<(), Self::Error> { - Self::set_qdisc_rate(*shaper, bandwidth_kbit, dry_run) + let mut last_warning = None; + loop { + let result = Self::qdisc_from_ifname(&shaper.ifname) + .and_then(|current| Self::set_qdisc_rate(¤t, bandwidth_kbit, dry_run)); + match result { + Ok(()) => return Ok(()), + Err(error) if error.is_qdisc_loss() && !SHUTDOWN.load(Ordering::Relaxed) => { + if last_warning + .is_none_or(|at: Instant| at.elapsed() >= Duration::from_secs(30)) + { + warn!( + "CAKE on {} is unavailable: {error}; waiting for SQM", + shaper.ifname + ); + last_warning = Some(Instant::now()); + } + sleep(Duration::from_millis(500)); + } + Err(error) => return Err(error), + } + } } } diff --git a/src/ratecontroller.rs b/src/ratecontroller.rs index fb9798a..a7ef573 100644 --- a/src/ratecontroller.rs +++ b/src/ratecontroller.rs @@ -297,6 +297,8 @@ impl Ratecontroller { let mut lastchg_t = Instant::now(); let mut lastdump_t = Instant::now(); + let mut last_reassert_t = Instant::now(); + let mut stats_missing = false; self.request_initial_rates()?; @@ -336,7 +338,22 @@ impl Ratecontroller { info!("Rate controller shutting down"); return Ok(()); } - sleep(sleep_time); + let wake_at = Instant::now() + sleep_time; + while Instant::now() < wake_at { + if SHUTDOWN.load(Ordering::Relaxed) { + info!("Rate controller shutting down"); + return Ok(()); + } + sleep( + wake_at + .saturating_duration_since(Instant::now()) + .min(Duration::from_millis(100)), + ); + } + if SHUTDOWN.load(Ordering::Relaxed) { + info!("Rate controller shutting down"); + return Ok(()); + } let now_t = Instant::now(); if now_t.duration_since(lastchg_t).as_secs_f64() @@ -345,10 +362,48 @@ impl Ratecontroller { // 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( - &mut self.stats_provider, - &self.settings.network.upload_interface, - )?; + (self.state_dl.current_bytes, self.state_ul.current_bytes) = + match get_interface_stats( + &mut self.stats_provider, + &self.settings.network.upload_interface, + ) { + Ok(stats) => { + if stats_missing { + info!("Interface statistics available again"); + stats_missing = false; + } + stats + } + Err(error) => { + if !stats_missing { + warn!( + "Interface statistics unavailable: {error}; waiting for interface recovery" + ); + stats_missing = true; + } + continue; + } + }; + if SHUTDOWN.load(Ordering::Relaxed) { + return Ok(()); + } + + // SQM may replace CAKE while the desired rate remains unchanged. + if !self.settings.advanced_settings.dry_run + && now_t.duration_since(last_reassert_t) >= Duration::from_secs(5) + { + self.traffic_control.set_rate( + &self.state_dl.shaper, + self.state_dl.current_rate.round() as u64, + false, + )?; + self.traffic_control.set_rate( + &self.state_ul.shaper, + self.state_ul.current_rate.round() as u64, + false, + )?; + last_reassert_t = now_t; + } self.update_deltas()?; if self.state_dl.deltas.is_empty() || self.state_ul.deltas.is_empty() {