sqm-autorate-rust/src/main.rs

extern crate core;

mod baseliner;
mod config;
mod log;
mod metrics;
mod netlink;
mod pinger;
mod pinger_icmp;
mod pinger_icmp_ts;
mod ratecontroller;
mod reflector_selector;
mod time;
mod util;

use crate::baseliner::{Baseliner, ReflectorStats};
use crate::metrics::Metrics;
use ::log::{debug, info};
use std::collections::HashMap;
use std::net::IpAddr;
use std::str::FromStr;
use std::sync::mpsc::{channel, sync_channel};
use std::sync::{Arc, Mutex, RwLock};
use std::thread::sleep;
use std::time::Duration;
use std::time::Instant;
use std::{process, thread};

use crate::config::{Config, MeasurementType};
use crate::netlink::Netlink;
use crate::pinger::{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;

const VERSION: &str = env!("CARGO_PKG_VERSION");

fn main() -> anyhow::Result<()> {
    println!("Starting sqm-autorate version {}", VERSION);

    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;

    // 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 default_reflectors = [
        IpAddr::from_str("9.9.9.9")?,
        IpAddr::from_str("8.238.120.14")?,
        IpAddr::from_str("74.82.42.42")?,
        IpAddr::from_str("194.242.2.2")?,
        IpAddr::from_str("208.67.222.222")?,
        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());
        }
        false => {
            let mut peers = reflector_peers_lock.write().unwrap();
            peers.extend_from_slice(&default_reflectors);
        }
    }

    let (baseliner_stats_sender, baseliner_stats_receiver) = channel();
    let (reselect_sender, reselect_receiver) = channel();

    let (metrics_sender, metrics_receiver) = sync_channel(1000);

    let ping_tx = config
        .observability_export_ping_metrics
        .then(|| metrics_sender.clone());

    let baseline_tx = config
        .observability_export_baseline_metrics
        .then(|| metrics_sender.clone());

    let event_tx = config
        .observability_export_events
        .then(|| metrics_sender.clone());

    let (mut pinger_receiver, mut pinger_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 baseliner = Baseliner {
        config: config.clone(),
        owd_baseline: owd_baseline.clone(),
        owd_recent: owd_recent.clone(),
        reselect_trigger: reselect_sender.clone(),
        start_time: start_t,
        stats_receiver: baseliner_stats_receiver,
        metrics_sender: metrics_sender.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 (error_tx, error_rx) = channel::<anyhow::Error>();

    let err_tx = error_tx.clone();
    let metrics_sender_clone = metrics_sender.clone();
    let reflector_peers_lock_clone = reflector_peers_lock.clone();
    thread::Builder::new()
        .name("receiver".to_string())
        .spawn(move || {
            if let Err(e) = pinger_receiver.listen(
                id,
                config.measurement_type,
                reflector_peers_lock_clone,
                baseliner_stats_sender,
                metrics_sender_clone,
            ) {
                let _ = err_tx.send(e);
            }
        })?;

    let err_tx = error_tx.clone();
    thread::Builder::new()
        .name("baseliner".to_string())
        .spawn(move || {
            if let Err(e) = baseliner.run() {
                let _ = err_tx.send(e);
            }
        })?;

    let err_tx = error_tx.clone();
    let reflector_peers_lock_clone = reflector_peers_lock.clone();
    thread::Builder::new()
        .name("sender".to_string())
        .spawn(move || {
            if let Err(e) = pinger_sender.send(
                id,
                config.measurement_type,
                reflector_peers_lock_clone,
                config.tick_interval,
            ) {
                let _ = err_tx.send(e);
            }
        })?;

    if reflector_pool_size > config.num_reflectors as usize {
        let reflector_selector = ReflectorSelector {
            config: config.clone(),
            owd_recent: owd_recent.clone(),
            reflector_peers_lock: reflector_peers_lock.clone(),
            reflector_pool,
            trigger_channel: reselect_receiver,
            metrics_sender: metrics_sender.clone(),
        };
        let err_tx = error_tx.clone();
        thread::Builder::new()
            .name("reselection".to_string())
            .spawn(move || {
                if let Err(e) = reflector_selector.run() {
                    let _ = err_tx.send(e);
                }
            })?;
    }

    // Sleep 10 seconds before we start adjusting speeds
    sleep(Duration::new(10, 0));

    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 mut ratecontroller = Ratecontroller::new(
        config.clone(),
        owd_baseline,
        owd_recent,
        reflector_peers_lock,
        reselect_sender,
        dl_direction,
        ul_direction,
        metrics_sender.clone(),
    )?;

    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);
            }
        })?;

    let metrics = Metrics {
        config,
        metrics_receiver,
    };

    let err_tx = error_tx.clone();
    thread::Builder::new()
        .name("metrics".to_string())
        .spawn(move || {
            if let Err(e) = metrics.run() {
                let _ = err_tx.send(e);
            }
        })?;

    // Drop original sender so recv() unblocks if all threads exit cleanly
    drop(error_tx);

    // Wait for first error
    match error_rx.recv() {
        Ok(e) => Err(anyhow::anyhow!("thread exited with error: {e}")),
        Err(_) => Ok(()), // all senders dropped = all threads exited without error
    }
}