refactor(baseliner): Replace shared Mutex state with thread-safe control snapshots

Eliminate lock contention between the baseliner, ratecontroller, and selector by moving to a message-passing architecture using immutable snapshots of reflector state. This improves data consistency and simplifies cross-thread communication.

src/reflector_selector.rs

@@ -5,22 +5,108 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
+use crate::Config;
 use crate::SHUTDOWN;
+use crate::baseliner::ControlSnapshot;
 use crate::metrics::{Metric, MetricsSender};
-use crate::util::{ArcMutex, MutexExt, RwLockExt};
-use crate::{Config, ReflectorStats};
-use flume::Receiver;
+use crate::util::RwLockExt;
+use flume::{Receiver, RecvError, RecvTimeoutError, Selector};
 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;
+use std::time::{Duration, Instant};
+
+enum SelectorEvent {
+    Trigger(Result<bool, RecvError>),
+    Snapshot(Result<ControlSnapshot, RecvError>),
+}
+
+#[derive(Debug, Eq, PartialEq)]
+enum WaitOutcome {
+    Triggered,
+    DeadlineReached,
+    ChannelsClosed,
+}
+
+fn wait_for_trigger_or_deadline(
+    trigger_rx: &Receiver<bool>,
+    snapshot_rx: &Receiver<ControlSnapshot>,
+    latest_snapshot: &mut Option<ControlSnapshot>,
+    deadline: Instant,
+) -> WaitOutcome {
+    let mut trigger_connected = true;
+    let mut snapshot_connected = true;
+
+    loop {
+        if Instant::now() >= deadline {
+            return WaitOutcome::DeadlineReached;
+        }
+
+        let event = match (trigger_connected, snapshot_connected) {
+            (true, true) => Selector::new()
+                .recv(trigger_rx, SelectorEvent::Trigger)
+                .recv(snapshot_rx, SelectorEvent::Snapshot)
+                .wait_deadline(deadline)
+                .ok(),
+            (true, false) => match trigger_rx.recv_deadline(deadline) {
+                Ok(trigger) => Some(SelectorEvent::Trigger(Ok(trigger))),
+                Err(RecvTimeoutError::Timeout) => None,
+                Err(RecvTimeoutError::Disconnected) => {
+                    Some(SelectorEvent::Trigger(Err(RecvError::Disconnected)))
+                }
+            },
+            (false, true) => match snapshot_rx.recv_deadline(deadline) {
+                Ok(snapshot) => Some(SelectorEvent::Snapshot(Ok(snapshot))),
+                Err(RecvTimeoutError::Timeout) => None,
+                Err(RecvTimeoutError::Disconnected) => {
+                    Some(SelectorEvent::Snapshot(Err(RecvError::Disconnected)))
+                }
+            },
+            (false, false) => return WaitOutcome::ChannelsClosed,
+        };
+
+        match event {
+            Some(SelectorEvent::Trigger(Ok(_))) => return WaitOutcome::Triggered,
+            Some(SelectorEvent::Trigger(Err(_))) => trigger_connected = false,
+            Some(SelectorEvent::Snapshot(Ok(snapshot))) => *latest_snapshot = Some(snapshot),
+            Some(SelectorEvent::Snapshot(Err(_))) => snapshot_connected = false,
+            None => return WaitOutcome::DeadlineReached,
+        }
+    }
+}
+
+fn receive_snapshots_until(
+    snapshot_rx: &Receiver<ControlSnapshot>,
+    latest_snapshot: &mut Option<ControlSnapshot>,
+    deadline: Instant,
+) {
+    loop {
+        if Instant::now() >= deadline {
+            return;
+        }
+
+        match snapshot_rx.recv_deadline(deadline) {
+            Ok(snapshot) => *latest_snapshot = Some(snapshot),
+            Err(RecvTimeoutError::Timeout) => return,
+            Err(RecvTimeoutError::Disconnected) => {
+                sleep(deadline.saturating_duration_since(Instant::now()));
+                return;
+            }
+        }
+    }
+}
+
+fn recent_rtt_for_peer(snapshot: Option<&ControlSnapshot>, peer: &IpAddr) -> Option<u64> {
+    snapshot
+        .and_then(|snapshot| snapshot.reflectors.get(peer))
+        .map(|state| (state.recent.down + state.recent.up) as u64)
+}
 
 pub struct ReflectorSelector {
     pub config: Config,
-    pub owd_recent: ArcMutex<HashMap<IpAddr, ReflectorStats>>,
+    pub snapshot_rx: Receiver<ControlSnapshot>,
     pub reflector_peers_lock: Arc<RwLock<Vec<IpAddr>>>,
     pub reflector_pool: Vec<IpAddr>,
     pub trigger_channel: Receiver<bool>,
@@ -31,10 +117,15 @@ impl ReflectorSelector {
     pub fn run(&self) -> anyhow::Result<()> {
         let mut selector_sleep_time = Duration::new(30, 0);
         let mut reselection_count = 0;
+        let mut latest_snapshot = None;
         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);
+        receive_snapshots_until(
+            &self.snapshot_rx,
+            &mut latest_snapshot,
+            Instant::now() + baseline_sleep_time,
+        );
 
         loop {
             if SHUTDOWN.load(Ordering::Relaxed) {
@@ -46,10 +137,19 @@ impl ReflectorSelector {
              * 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();
+            let triggered = match wait_for_trigger_or_deadline(
+                &self.trigger_channel,
+                &self.snapshot_rx,
+                &mut latest_snapshot,
+                Instant::now() + selector_sleep_time,
+            ) {
+                WaitOutcome::Triggered => true,
+                WaitOutcome::DeadlineReached => false,
+                WaitOutcome::ChannelsClosed => {
+                    info!("Reflector selector channels closed, shutting down");
+                    return Ok(());
+                }
+            };
             reselection_count += 1;
             info!("Starting reselection [#{}]", reselection_count);
             self.metrics.send(Metric::Event {
@@ -93,17 +193,19 @@ impl ReflectorSelector {
 
             debug!("Waiting for candidates to be baselined");
             // Wait for several seconds to allow all reflectors to be re-baselined
-            sleep(baseline_sleep_time);
+            receive_snapshots_until(
+                &self.snapshot_rx,
+                &mut latest_snapshot,
+                Instant::now() + 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;
+                if let Some(rtt) = recent_rtt_for_peer(latest_snapshot.as_ref(), &peer) {
                     candidates.push((peer, rtt));
                     info!("Candidate reflector: {} RTT: {}", peer, rtt);
                 } else {
@@ -153,3 +255,147 @@ impl ReflectorSelector {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::baseliner::{EwmaStats, ReflectorState};
+    use std::collections::HashMap;
+    use std::thread;
+
+    fn snapshot(generated_at: Instant) -> ControlSnapshot {
+        ControlSnapshot {
+            generated_at,
+            reflectors: HashMap::new(),
+        }
+    }
+
+    fn snapshot_with_rtt(generated_at: Instant, peer: IpAddr, rtt: f64) -> ControlSnapshot {
+        let mut reflectors = HashMap::new();
+        reflectors.insert(
+            peer,
+            ReflectorState {
+                baseline: EwmaStats {
+                    down: rtt / 2.0,
+                    up: rtt / 2.0,
+                },
+                recent: EwmaStats {
+                    down: rtt / 2.0,
+                    up: rtt / 2.0,
+                },
+                last_receive_at: generated_at,
+            },
+        );
+        ControlSnapshot {
+            generated_at,
+            reflectors,
+        }
+    }
+
+    #[test]
+    fn wait_retains_the_newest_snapshot() {
+        let (_trigger_tx, trigger_rx) = flume::unbounded();
+        let (snapshot_tx, snapshot_rx) = flume::unbounded();
+        let generated_at = Instant::now();
+        snapshot_tx.send(snapshot(generated_at)).unwrap();
+        snapshot_tx
+            .send(snapshot(generated_at + Duration::from_millis(1)))
+            .unwrap();
+        snapshot_tx
+            .send(snapshot(generated_at + Duration::from_millis(2)))
+            .unwrap();
+        let mut latest_snapshot = None;
+
+        let result = wait_for_trigger_or_deadline(
+            &trigger_rx,
+            &snapshot_rx,
+            &mut latest_snapshot,
+            Instant::now() + Duration::from_millis(20),
+        );
+
+        assert_eq!(result, WaitOutcome::DeadlineReached);
+        assert_eq!(
+            latest_snapshot.unwrap().generated_at,
+            generated_at + Duration::from_millis(2)
+        );
+    }
+
+    #[test]
+    fn snapshots_do_not_extend_the_deadline() {
+        let (_trigger_tx, trigger_rx) = flume::unbounded();
+        let (snapshot_tx, snapshot_rx) = flume::unbounded();
+        let producer = thread::spawn(move || {
+            let stop_at = Instant::now() + Duration::from_millis(500);
+            while Instant::now() < stop_at {
+                if snapshot_tx.send(snapshot(Instant::now())).is_err() {
+                    break;
+                }
+                sleep(Duration::from_millis(1));
+            }
+        });
+        let mut latest_snapshot = None;
+        let started_at = Instant::now();
+
+        let result = wait_for_trigger_or_deadline(
+            &trigger_rx,
+            &snapshot_rx,
+            &mut latest_snapshot,
+            started_at + Duration::from_millis(30),
+        );
+        let elapsed = started_at.elapsed();
+        drop(snapshot_rx);
+        producer.join().unwrap();
+
+        assert_eq!(result, WaitOutcome::DeadlineReached);
+        assert!(latest_snapshot.is_some());
+        assert!(elapsed < Duration::from_millis(250), "elapsed: {elapsed:?}");
+    }
+
+    #[test]
+    fn trigger_interrupts_the_wait() {
+        let (trigger_tx, trigger_rx) = flume::unbounded();
+        let (_snapshot_tx, snapshot_rx) = flume::unbounded();
+        trigger_tx.send(true).unwrap();
+        let mut latest_snapshot = None;
+        let started_at = Instant::now();
+
+        let result = wait_for_trigger_or_deadline(
+            &trigger_rx,
+            &snapshot_rx,
+            &mut latest_snapshot,
+            started_at + Duration::from_secs(1),
+        );
+
+        assert_eq!(result, WaitOutcome::Triggered);
+        assert!(started_at.elapsed() < Duration::from_millis(250));
+    }
+
+    #[test]
+    fn baseline_wait_retains_snapshot_for_evaluation() {
+        let (snapshot_tx, snapshot_rx) = flume::unbounded();
+        let peer = "192.0.2.1".parse().unwrap();
+        let generated_at = Instant::now();
+        snapshot_tx
+            .send(snapshot_with_rtt(generated_at, peer, 20.0))
+            .unwrap();
+        snapshot_tx
+            .send(snapshot_with_rtt(
+                generated_at + Duration::from_millis(1),
+                peer,
+                42.0,
+            ))
+            .unwrap();
+        let mut latest_snapshot = None;
+
+        receive_snapshots_until(
+            &snapshot_rx,
+            &mut latest_snapshot,
+            Instant::now() + Duration::from_millis(20),
+        );
+
+        assert_eq!(
+            recent_rtt_for_peer(latest_snapshot.as_ref(), &peer),
+            Some(42)
+        );
+    }
+}