Lochnair-Labs/sqm-autorate-rust
2
0 Fork 0
Code Issues 5 Pull Requests 2 Actions 8 Packages Projects Releases Wiki Activity

refactor(main): modularize main loop and improve resource cleanup logic

Browse Source
This commit is contained in:
Nils Andreas Svee committed 2026-07-14 22:53:01 +07:00
1 parent 8b0c37bc83
commit b56c8b72cc
1 file changed
+215 -150
Download Patch File Download Diff File Expand all files Collapse all files
src/main.rs
+215 -150
View File
@@ -22,7 +22,7 @@ mod util;
use crate::baseliner::{Baseliner, ReflectorStats};
use crate::metrics::{Metric, Metrics, MetricsSender};
use ::log::{debug, info};
use ::log::{debug, info, warn};
use flume::RecvTimeoutError;
use std::collections::HashMap;
use std::net::IpAddr;
@@ -30,22 +30,30 @@ use std::str::FromStr;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::sleep;
use std::time::Duration;
use std::time::Instant;
use std::time::{Duration, Instant};
use std::{process, thread};
use crate::config::{Config, MeasurementType};
use crate::netlink::Netlink;
use crate::netlink::{Netlink, Qdisc};
use crate::pinger::{InFlightProbeCache, PingListener, PingSender};
use crate::pinger_icmp::{PingerICMPEchoListener, PingerICMPEchoSender};
use crate::pinger_icmp_ts::{PingerICMPTimestampListener, PingerICMPTimestampSender};
use crate::ratecontroller::{Ratecontroller, StatsDirection};
use crate::reflector_selector::ReflectorSelector;
struct ReflectorSetup {
peers: Arc<RwLock<Vec<IpAddr>>>,
pool: Vec<IpAddr>,
reselection_enabled: bool,
active_count: usize,
}
type PingListenerBox = Box<dyn PingListener + Send>;
type PingSenderBox = Box<dyn PingSender + Send>;
pub static SHUTDOWN: AtomicBool = AtomicBool::new(false);
extern "C" fn signal_handler(_: libc::c_int) {
println!("Signal received, shutting down sqm-autorate-rust");
SHUTDOWN.store(true, Ordering::Relaxed);
}
@@ -77,9 +85,47 @@ fn compute_inflight_probe_capacity(
.max(INFLIGHT_CAPACITY_MIN)
}
fn main() -> anyhow::Result<()> {
println!("Starting sqm-autorate-rust version {}", VERSION);
fn create_pinger(measurement_type: MeasurementType) -> (PingListenerBox, PingSenderBox) {
match measurement_type {
MeasurementType::Icmp => (
Box::new(PingerICMPEchoListener {}),
Box::new(PingerICMPEchoSender {}),
),
MeasurementType::IcmpTimestamps => (
Box::new(PingerICMPTimestampListener {}),
Box::new(PingerICMPTimestampSender {}),
),
MeasurementType::Ntp | MeasurementType::TcpTimestamps => {
todo!()
}
}
}
fn initialize_shaper(config: &Config) -> anyhow::Result<(Qdisc, Qdisc)> {
let down = Netlink::qdisc_from_ifname(&config.download_interface)?;
let up = Netlink::qdisc_from_ifname(&config.upload_interface)?;
info!(
"Setting shaper rates to minimum (D/L): {} / {}",
config.download_min_kbits, config.upload_min_kbits,
);
Netlink::set_qdisc_rate(down, config.download_min_kbits as u64, config.dry_run)?;
Netlink::set_qdisc_rate(up, config.upload_min_kbits as u64, config.dry_run)?;
let settle_time = Duration::from_secs(2);
info!(
"Sleeping for {} seconds to give the shaper a chance to control existing bloat",
settle_time.as_secs_f64(),
);
sleep(settle_time);
Ok((down, up))
}
fn install_signal_handlers() {
unsafe {
libc::signal(
libc::SIGINT,
@@ -90,25 +136,54 @@ fn main() -> anyhow::Result<()> {
signal_handler as *const () as libc::sighandler_t,
);
}
let config = Config::new()?;
log::init(config.log_level)?;
if config.dry_run {
info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***");
}
let mut reflectors = config.load_reflectors()?;
let start_t = Instant::now();
// The identifier field in ICMP is only 2 bytes
// so take the last 2 bytes of the PID as the ID
let id = (process::id() & 0xFFFF) as u16;
fn interface_stats_directions(config: &Config) -> (StatsDirection, StatsDirection) {
let download = if config.download_interface.starts_with("ifb")
|| config.download_interface.starts_with("veth")
{
StatsDirection::TX
} else {
StatsDirection::RX
};
// Create data structures shared by different threads
let owd_baseline = Arc::new(Mutex::new(HashMap::<IpAddr, ReflectorStats>::new()));
let owd_recent = Arc::new(Mutex::new(HashMap::<IpAddr, ReflectorStats>::new()));
let reflector_peers_lock = Arc::new(RwLock::new(Vec::<IpAddr>::new()));
let mut reflector_pool = Vec::<IpAddr>::new();
let reflector_pool_size = reflectors.len();
let upload = if config.upload_interface.starts_with("ifb")
|| config.upload_interface.starts_with("veth")
{
StatsDirection::RX
} else {
StatsDirection::TX
};
(download, upload)
}
fn probe_identifier() -> u16 {
// The ICMP identifier is only two bytes.
(process::id() & 0xffff) as u16
}
fn restore_shaper(config: &Config, down: Qdisc, up: Qdisc) {
info!(
"Restoring base shaper rates (D/L): {} / {}",
config.download_base_kbits, config.upload_base_kbits,
);
if let Err(error) =
Netlink::set_qdisc_rate(down, config.download_base_kbits as u64, config.dry_run)
{
warn!("Failed to restore download shaper rate: {error}");
}
if let Err(error) = Netlink::set_qdisc_rate(up, config.upload_base_kbits as u64, config.dry_run)
{
warn!("Failed to restore upload shaper rate: {error}");
}
}
fn setup_reflectors(config: &Config) -> anyhow::Result<ReflectorSetup> {
let configured = config.load_reflectors()?;
let configured_count = configured.len();
let default_reflectors = [
IpAddr::from_str("9.9.9.9")?,
@@ -119,17 +194,57 @@ fn main() -> anyhow::Result<()> {
IpAddr::from_str("94.140.14.14")?,
];
match reflector_pool_size > config.num_reflectors as usize {
true => {
let mut peers = reflector_peers_lock.write().unwrap();
peers.extend_from_slice(&default_reflectors);
reflector_pool.append(reflectors.as_mut());
let reselection_enabled = configured_count > config.num_reflectors as usize;
let pool = if reselection_enabled {
configured
} else {
Vec::new()
};
Ok(ReflectorSetup {
peers: Arc::new(RwLock::new(default_reflectors.to_vec())),
pool,
reselection_enabled,
active_count: default_reflectors.len().max(config.num_reflectors as usize),
})
}
false => {
let mut peers = reflector_peers_lock.write().unwrap();
peers.extend_from_slice(&default_reflectors);
fn wait_for_exit(error_rx: &flume::Receiver<anyhow::Error>) -> anyhow::Result<()> {
loop {
match error_rx.recv_timeout(Duration::from_secs(1)) {
Ok(error) => {
return Err(anyhow::anyhow!("thread exited with error: {error}"));
}
Err(RecvTimeoutError::Disconnected) => {
return Ok(());
}
Err(RecvTimeoutError::Timeout) => {
if SHUTDOWN.load(Ordering::Relaxed) {
info!("Received shutdown signal");
return Ok(());
}
}
}
}
}
fn run(config: Config) -> anyhow::Result<()> {
if config.dry_run {
info!("*** MONITORING MODE ACTIVE — qdisc rates will NOT be changed ***");
}
let start_time = Instant::now();
let probe_id = probe_identifier();
let ReflectorSetup {
peers: reflector_peers,
pool: reflector_pool,
reselection_enabled,
active_count: active_reflector_count,
} = setup_reflectors(&config)?;
let owd_baseline = Arc::new(Mutex::new(HashMap::<IpAddr, ReflectorStats>::new()));
let owd_recent = Arc::new(Mutex::new(HashMap::<IpAddr, ReflectorStats>::new()));
let (baseliner_stats_tx, baseliner_stats_rx) = flume::unbounded();
let (error_tx, error_rx) = flume::unbounded::<anyhow::Error>();
@@ -148,8 +263,8 @@ fn main() -> anyhow::Result<()> {
let handle = thread::Builder::new()
.name("metrics".to_string())
.spawn(move || {
if let Err(e) = metrics.run() {
let _ = err_tx.send(e);
if let Err(error) = metrics.run() {
let _ = err_tx.send(error);
}
})?;
(Some(tx), Some(handle))
@@ -181,79 +296,47 @@ fn main() -> anyhow::Result<()> {
},
});
let (mut ping_listener, mut ping_sender) = match config.measurement_type {
MeasurementType::Icmp => (
Box::new(PingerICMPEchoListener {}) as Box<dyn PingListener + Send>,
Box::new(PingerICMPEchoSender {}) as Box<dyn PingSender + Send>,
),
MeasurementType::IcmpTimestamps => (
Box::new(PingerICMPTimestampListener {}) as Box<dyn PingListener + Send>,
Box::new(PingerICMPTimestampSender {}) as Box<dyn PingSender + Send>,
),
MeasurementType::Ntp | MeasurementType::TcpTimestamps => {
todo!()
}
};
let reselection_enabled = reflector_pool_size > config.num_reflectors as usize;
let active_reflector_count = default_reflectors.len().max(config.num_reflectors as usize);
let (mut ping_listener, mut ping_sender) = create_pinger(config.measurement_type);
let inflight_probe_capacity =
compute_inflight_probe_capacity(&config, active_reflector_count, reselection_enabled);
info!(
"In-flight probe cache capacity: {} (active_reflectors={}, reselection_enabled={}, tick_interval_s={})",
inflight_probe_capacity, active_reflector_count, reselection_enabled, config.tick_interval
);
let inflight: InFlightProbeCache =
Arc::new(Mutex::new(HashMap::with_capacity(inflight_probe_capacity)));
let baseliner = Baseliner {
config: config.clone(),
owd_baseline: owd_baseline.clone(),
owd_recent: owd_recent.clone(),
owd_baseline: Arc::clone(&owd_baseline),
owd_recent: Arc::clone(&owd_recent),
reselect_trigger: reselect_tx.clone(),
start_time: start_t,
start_time,
stats_rx: baseliner_stats_rx,
baseline_metrics,
event_metrics: event_metrics.clone(),
};
let down_qdisc = Netlink::qdisc_from_ifname(config.download_interface.as_str())?;
let up_qdisc = Netlink::qdisc_from_ifname(config.upload_interface.as_str())?;
/* Set initial TC values to minimum
* so there should be no initial bufferbloat to
* fool the baseliner
*/
info!(
"Setting shaper rates to minimum (D/L): {} / {}",
config.download_min_kbits, config.upload_min_kbits
);
Netlink::set_qdisc_rate(down_qdisc, config.download_min_kbits as u64, config.dry_run)?;
Netlink::set_qdisc_rate(up_qdisc, config.upload_min_kbits as u64, config.dry_run)?;
// Sleep for a few seconds to give the shaper a chance
// to control the queue if load is heavy
let settle_sleep_time = Duration::new(2, 0);
info!(
"Sleeping for {} to give the shaper a chance to get in control if there's bloat",
settle_sleep_time.as_secs_f64()
);
sleep(settle_sleep_time);
let (down_qdisc, up_qdisc) = initialize_shaper(&config)?;
let err_tx = error_tx.clone();
let reflector_peers_lock_clone = reflector_peers_lock.clone();
let inflight_listener = inflight.clone();
let listener_peers = Arc::clone(&reflector_peers);
let listener_inflight = Arc::clone(&inflight);
let measurement_type = config.measurement_type;
thread::Builder::new()
.name("receiver".to_string())
.spawn(move || {
if let Err(e) = ping_listener.listen(
id,
config.measurement_type,
reflector_peers_lock_clone,
inflight_listener,
if let Err(error) = ping_listener.listen(
probe_id,
measurement_type,
listener_peers,
listener_inflight,
baseliner_stats_tx,
ping_metrics,
) {
let _ = err_tx.send(e);
let _ = err_tx.send(error);
}
})?;
@@ -261,25 +344,27 @@ fn main() -> anyhow::Result<()> {
thread::Builder::new()
.name("baseliner".to_string())
.spawn(move || {
if let Err(e) = baseliner.run() {
let _ = err_tx.send(e);
if let Err(error) = baseliner.run() {
let _ = err_tx.send(error);
}
})?;
let err_tx = error_tx.clone();
let reflector_peers_lock_clone = reflector_peers_lock.clone();
let inflight_sender = inflight.clone();
let sender_peers = Arc::clone(&reflector_peers);
let sender_inflight = Arc::clone(&inflight);
let measurement_type = config.measurement_type;
let tick_interval = config.tick_interval;
thread::Builder::new()
.name("sender".to_string())
.spawn(move || {
if let Err(e) = ping_sender.send(
id,
config.measurement_type,
reflector_peers_lock_clone,
inflight_sender,
config.tick_interval,
if let Err(error) = ping_sender.send(
probe_id,
measurement_type,
sender_peers,
sender_inflight,
tick_interval,
) {
let _ = err_tx.send(e);
let _ = err_tx.send(error);
}
})?;
@@ -288,8 +373,8 @@ fn main() -> anyhow::Result<()> {
if reselection_enabled {
let reflector_selector = ReflectorSelector {
config: config.clone(),
owd_recent: owd_recent.clone(),
reflector_peers_lock: reflector_peers_lock.clone(),
owd_recent: Arc::clone(&owd_recent),
reflector_peers_lock: Arc::clone(&reflector_peers),
reflector_pool,
trigger_channel: reselect_rx,
metrics: event_metrics,
@@ -298,75 +383,53 @@ fn main() -> anyhow::Result<()> {
thread::Builder::new()
.name("reselection".to_string())
.spawn(move || {
if let Err(e) = reflector_selector.run() {
let _ = err_tx.send(e);
if let Err(error) = reflector_selector.run() {
let _ = err_tx.send(error);
}
})?;
}
// Sleep 10 seconds before we start adjusting speeds
sleep(Duration::new(10, 0));
// Give the baseliner time to collect initial samples before adjusting rates.
sleep(Duration::from_secs(10));
let dl_direction = if config.download_interface.starts_with("ifb")
|| config.download_interface.starts_with("veth")
{
StatsDirection::TX
} else {
StatsDirection::RX
};
let ul_direction = if config.upload_interface.starts_with("ifb")
|| config.upload_interface.starts_with("veth")
{
StatsDirection::RX
} else {
StatsDirection::TX
};
let (dl_direction, ul_direction) = interface_stats_directions(&config);
debug!(
"Download direction: {}:{:?}",
config.download_interface, dl_direction
);
debug!(
"Upload direction: {}:{:?}",
config.upload_interface, ul_direction
);
let mut ratecontroller = Ratecontroller::new(
config.clone(),
owd_baseline,
owd_recent,
reflector_peers_lock,
reflector_peers,
reselect_tx,
dl_direction,
ul_direction,
rate_metrics,
)?;
debug!(
"Download direction: {}:{:?}",
config.download_interface, dl_direction
);
debug!(
"Upload direction: {}:{:?}",
config.upload_interface, ul_direction
);
let err_tx = error_tx.clone();
thread::Builder::new()
.name("ratecontroller".to_string())
.spawn(move || {
if let Err(e) = ratecontroller.run() {
let _ = err_tx.send(e);
if let Err(error) = ratecontroller.run() {
let _ = err_tx.send(error);
}
})?;
// Drop original sender so the channel disconnects if all threads exit cleanly
// Drop the original sender so the channel disconnects if all workers exit cleanly.
drop(error_tx);
let result = loop {
match error_rx.recv_timeout(Duration::from_secs(1)) {
Ok(e) => break Err(anyhow::anyhow!("thread exited with error: {e}")),
Err(RecvTimeoutError::Disconnected) => break Ok(()),
Err(RecvTimeoutError::Timeout) => {
if SHUTDOWN.load(Ordering::Relaxed) {
info!("Received shutdown signal");
break Ok(());
}
}
}
};
let result = wait_for_exit(&error_rx);
// Make the error path request shutdown too, rather than leaving the other workers running.
SHUTDOWN.store(true, Ordering::Relaxed);
let stopping_reason = if result.is_err() { "error" } else { "signal" };
main_event_metrics.send(Metric::Event {
@@ -377,24 +440,26 @@ fn main() -> anyhow::Result<()> {
});
// Drop all MetricsSender instances and the raw tx so the metrics channel
// disconnects once all threads also drop their copies, allowing the
// metrics thread to drain and exit cleanly.
// disconnects once all worker threads also drop their copies.
drop(main_event_metrics);
drop(metrics_tx);
if let Some(handle) = metrics_thread_handle {
let _ = handle.join();
}
info!(
"Restoring base shaper rates (D/L): {} / {}",
config.download_base_kbits, config.upload_base_kbits
);
let _ = Netlink::set_qdisc_rate(
down_qdisc,
config.download_base_kbits as u64,
config.dry_run,
);
let _ = Netlink::set_qdisc_rate(up_qdisc, config.upload_base_kbits as u64, config.dry_run);
restore_shaper(&config, down_qdisc, up_qdisc);
result
}
fn main() -> anyhow::Result<()> {
println!("Starting sqm-autorate-rust version {}", VERSION);
install_signal_handlers();
let config = Config::new()?;
log::init(config.log_level)?;
run(config)
}