sqm-autorate-rust/src/reflector_selector.rs

// SPDX-FileCopyrightText: 2022-Present Charles Corrigan mailto:chas-iot@runegate.org (github @chas-iot)
// SPDX-FileCopyrightText: 2022-Present Daniel Lakeland mailto:dlakelan@street-artists.org (github @dlakelan)
// SPDX-FileCopyrightText: 2022-Present Mark Baker mailto:mark@vpost.net (github @Fail-Safe)
// SPDX-FileCopyrightText: 2022-Present Nils Andreas Svee mailto:contact@lochnair.net (github @Lochnair)
//
// SPDX-License-Identifier: MPL-2.0

use crate::SHUTDOWN;
use crate::metrics::{Metric, MetricsSender};
use crate::util::{ArcMutex, MutexExt, RwLockExt};
use crate::{Config, ReflectorStats};
use flume::Receiver;
use log::{debug, info};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::atomic::Ordering;
use std::sync::{Arc, RwLock};
use std::thread::sleep;
use std::time::Duration;

pub struct ReflectorSelector {
    pub config: Config,
    pub owd_recent: ArcMutex<HashMap<IpAddr, ReflectorStats>>,
    pub reflector_peers_lock: Arc<RwLock<Vec<IpAddr>>>,
    pub reflector_pool: Vec<IpAddr>,
    pub trigger_channel: Receiver<bool>,
    pub metrics: MetricsSender,
}

impl ReflectorSelector {
    pub fn run(&self) -> anyhow::Result<()> {
        let mut selector_sleep_time = Duration::new(30, 0);
        let mut reselection_count = 0;
        let baseline_sleep_time =
            Duration::from_secs_f64(self.config.tick_interval * std::f64::consts::PI);
        // Initial wait of several seconds to allow some OWD data to build up
        sleep(baseline_sleep_time);

        loop {
            if SHUTDOWN.load(Ordering::Relaxed) {
                info!("Reflector selector shutting down");
                return Ok(());
            }
            /*
             * Selection is triggered either by some other thread triggering it through the channel,
             * or it passes the timeout. In any case we don't care about the result of this function,
             * so we ignore the result of it.
             */
            let triggered = self
                .trigger_channel
                .recv_timeout(selector_sleep_time)
                .is_ok();
            reselection_count += 1;
            info!("Starting reselection [#{}]", reselection_count);
            self.metrics.send(Metric::Event {
                name: "reflector_selection",
                reason: if triggered { "triggered" } else { "timeout" },
                reflector: None,
                tags: &[],
            });

            // After 40 reselections, slow down to every 15 minutes
            if reselection_count > 40 {
                selector_sleep_time = Duration::new(self.config.peer_reselection_time * 60, 0);
            }

            let mut next_peers: Vec<IpAddr> = Vec::new();
            let mut reflectors_peers = self.reflector_peers_lock.write_anyhow()?;

            // Include all current peers
            for reflector in reflectors_peers.iter() {
                debug!("Current peer: {}", reflector);
                next_peers.push(*reflector);
            }

            for _ in 0..20 {
                let next_candidate =
                    &self.reflector_pool[fastrand::usize(..self.reflector_pool.len())];
                if next_peers.contains(next_candidate) {
                    continue;
                }
                debug!("Next candidate: {}", next_candidate);
                next_peers.push(*next_candidate);
            }

            // Clone next_peers because we need it again after the baseline sleep
            // to iterate over candidates for RTT measurement.
            *reflectors_peers = next_peers.clone();

            // Drop the MutexGuard explicitly, as Rust won't unlock the mutex by default
            // until the guard goes out of scope
            drop(reflectors_peers);

            debug!("Waiting for candidates to be baselined");
            // Wait for several seconds to allow all reflectors to be re-baselined
            sleep(baseline_sleep_time);

            // Re-acquire the lock when we wake up again
            reflectors_peers = self.reflector_peers_lock.write_anyhow()?;

            let mut candidates = Vec::new();
            let owd_recent = self.owd_recent.lock_anyhow()?;

            for peer in next_peers {
                if owd_recent.contains_key(&peer) {
                    let rtt = (owd_recent[&peer].down_ewma + owd_recent[&peer].up_ewma) as u64;
                    candidates.push((peer, rtt));
                    info!("Candidate reflector: {} RTT: {}", peer, rtt);
                } else {
                    info!(
                        "No data found from candidate reflector: {} - skipping",
                        peer
                    );
                }
            }

            // Sort the candidates table now by ascending RTT
            candidates.sort_by(|a, b| a.1.cmp(&b.1));

            // Now we will just limit the candidates down to 2 * num_reflectors
            let mut num_reflectors = self.config.num_reflectors;
            let candidate_pool_num = (2 * num_reflectors) as usize;
            candidates.truncate(candidate_pool_num);

            for (candidate, rtt) in candidates.iter() {
                info!("Fastest candidate {}: {}", candidate, rtt);
            }

            // Shuffle the deck so we avoid overwhelming good reflectors (Fisher-Yates)
            for i in (1_usize..candidates.len()).rev() {
                let j = fastrand::usize(0..=i);
                candidates.swap(i, j);
            }

            if (candidates.len() as u8) < num_reflectors {
                num_reflectors = candidates.len() as u8;
            }

            let mut new_peers = Vec::new();
            for i in 0..num_reflectors {
                let peer = candidates[i as usize].0;
                new_peers.push(peer);
                info!("New selected peer: {}", peer);
                self.metrics.send(Metric::Event {
                    name: "reflector_selected",
                    reason: "",
                    reflector: Some(peer),
                    tags: &[],
                });
            }

            *reflectors_peers = new_peers;
        }
    }
}