fix(shaper): restore base rates and recover recreated CAKE

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,

src/main.rs

@@ -121,16 +121,24 @@ fn initialize_shaper_with_settle_time<T: TrafficControlBackend>(
         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<T: TrafficControlBackend>(
     }
 }
 
+struct ShaperGuard<'a, T: TrafficControlBackend> {
+    settings: &'a Settings,
+    traffic_control: &'a mut T,
+    down: T::Handle,
+    up: T::Handle,
+}
+
+impl<T: TrafficControlBackend> 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<ReflectorSetup> {
     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<Self::Handle, Self::Error> {
+            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)]
+        );
+    }
+}

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(&current, 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),
+            }
+        }
     }
 }

src/ratecontroller.rs

@@ -297,6 +297,8 @@ impl<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
 
         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<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
                 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<S: InterfaceStatsProvider, T: TrafficControlBackend> Ratecontroller<S, T> {
                 // 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.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() {