diff --git a/src/main.rs b/src/main.rs index 6ee7a81..7e54c81 100644 --- a/src/main.rs +++ b/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 - // disconnects once all worker threads also drop their copies. + let _ = ratecontroller_thread_handle.join(); + 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 } diff --git a/src/metrics.rs b/src/metrics.rs index a1fd40a..eca9fa2 100644 --- a/src/metrics.rs +++ b/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, pub metrics_rx: Receiver<(Metric, u64)>, + pub shutdown_requested: Arc, } 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 { diff --git a/src/ratecontroller.rs b/src/ratecontroller.rs index fb9798a..13c5f15 100644 --- a/src/ratecontroller.rs +++ b/src/ratecontroller.rs @@ -336,7 +336,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()