fix(metrics): preserve stop event and restore after controller exit

src/main.rs

@@ -264,6 +264,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
     let (reselect_tx, reselect_rx) = flume::bounded(1);
 
     let dropped = Arc::new(AtomicU32::new(0));
+    let metrics_shutdown = Arc::new(AtomicBool::new(false));
 
     let (metrics_tx, metrics_thread_handle) = if settings.observability.enabled {
         let (tx, rx) = flume::bounded(1000);
@@ -271,6 +272,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
             settings: settings.clone(),
             metrics_rx: rx,
             metrics_dropped: Arc::clone(&dropped),
+            shutdown_requested: Arc::clone(&metrics_shutdown),
         };
         let err_tx = error_tx.clone();
         let handle = thread::Builder::new()
@@ -420,7 +422,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() {
@@ -444,10 +446,17 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
         tags: &[],
     });
 
-    // Drop all MetricsSender instances and the raw tx so the metrics channel
+    let _ = ratecontroller_thread_handle.join();
-    // disconnects once all worker threads also drop their copies.
+    restore_shaper(
+        settings,
+        &mut main_traffic_control,
+        &down_shaper,
+        &up_shaper,
+    );
+
     drop(main_event_metrics);
     drop(metrics_tx);
+    metrics_shutdown.store(true, Ordering::Relaxed);
 
     if let Some(handle) = metrics_thread_handle {
         let deadline = Instant::now() + Duration::from_secs(1);
@@ -461,13 +470,6 @@ fn run(settings: &Settings) -> anyhow::Result<()> {
         }
     }
 
-    restore_shaper(
-        settings,
-        &mut main_traffic_control,
-        &down_shaper,
-        &up_shaper,
-    );
-
     result
 }
 

src/metrics.rs

@@ -2,7 +2,6 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
-use crate::SHUTDOWN;
 use crate::settings::Settings;
 use crate::settings::{MeasurementType, ObservabilityProtocol};
 use crate::time::Time;
@@ -12,7 +11,7 @@ use rustix::time::ClockId;
 use std::fmt::Write;
 use std::net::{IpAddr, TcpStream, UdpSocket};
 use std::sync::Arc;
-use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
 use std::time::{Duration, Instant};
 
 const MAX_RECONNECT_BACKOFF: u64 = 60;
@@ -228,6 +227,7 @@ pub struct Metrics {
     pub settings: Settings,
     pub metrics_dropped: Arc<AtomicU32>,
     pub metrics_rx: Receiver<(Metric, u64)>,
+    pub shutdown_requested: Arc<AtomicBool>,
 }
 
 impl Metrics {
@@ -261,7 +261,7 @@ impl Metrics {
         let mut last_flush = Instant::now();
 
         loop {
-            if SHUTDOWN.load(Ordering::Relaxed) {
+            if self.shutdown_requested.load(Ordering::Relaxed) {
                 let deadline = shutdown_deadline
                     .get_or_insert_with(|| Instant::now() + Duration::from_millis(400));
                 if Instant::now() >= *deadline {

src/ratecontroller.rs

@@ -336,7 +336,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()