From 6bcfa6a6e2e72e6433806e9d05a6179b5601a3f8 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Wed, 28 Oct 2026 18:01:20 +0700 Subject: [PATCH] refactor: replace Flume channels with Tokio and std mpsc --- Cargo.lock | 129 --------------- Cargo.toml | 1 - src/baseliner.rs | 59 ++----- src/controller.rs | 23 +-- src/main.rs | 49 +++--- src/measurement/icmp.rs | 23 ++- src/measurement/irtt.rs | 24 ++- src/metrics.rs | 10 +- src/ratecontroller.rs | 154 ++---------------- src/reflector_selector.rs | 331 +++++--------------------------------- 10 files changed, 133 insertions(+), 670 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 04ff6cc..c08293f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -69,12 +69,6 @@ dependencies = [ "hybrid-array", ] -[[package]] -name = "bumpalo" -version = "3.20.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" - [[package]] name = "bytes" version = "1.12.1" @@ -234,9 +228,6 @@ name = "fastrand" version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" -dependencies = [ - "getrandom", -] [[package]] name = "find-msvc-tools" @@ -244,43 +235,6 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484" -[[package]] -name = "flume" -version = "0.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" -dependencies = [ - "fastrand", - "futures-core", - "futures-sink", - "spin", -] - -[[package]] -name = "futures-core" -version = "0.3.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" - -[[package]] -name = "futures-sink" -version = "0.3.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" - -[[package]] -name = "getrandom" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" -dependencies = [ - "cfg-if", - "js-sys", - "libc", - "r-efi", - "wasm-bindgen", -] - [[package]] name = "glob" version = "0.3.4" @@ -374,16 +328,6 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" -[[package]] -name = "js-sys" -version = "0.3.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7883d941dae510fb2d978fc3fe018c71c9e2892fd38854de3e8b92c2e5ad9cc5" -dependencies = [ - "cfg-if", - "wasm-bindgen", -] - [[package]] name = "libc" version = "0.2.191" @@ -513,12 +457,6 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "once_cell" -version = "1.21.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" - [[package]] name = "parking_lot" version = "0.12.5" @@ -582,12 +520,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "r-efi" -version = "6.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" - [[package]] name = "redox_syscall" version = "0.5.18" @@ -656,12 +588,6 @@ dependencies = [ "windows-sys", ] -[[package]] -name = "rustversion" -version = "1.0.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" - [[package]] name = "ryu" version = "1.0.23" @@ -764,15 +690,6 @@ dependencies = [ "windows-sys", ] -[[package]] -name = "spin" -version = "0.9.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" -dependencies = [ - "lock_api", -] - [[package]] name = "sqm-autorate-rust" version = "0.4.1" @@ -781,7 +698,6 @@ dependencies = [ "config", "csv", "fastrand", - "flume", "icmp-socket2", "irtt-client", "libc", @@ -975,51 +891,6 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" -[[package]] -name = "wasm-bindgen" -version = "0.2.129" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bb54f33acc68fd454578d9820b0bde1a1a3d17aa17bb7b6595806d02886d409" -dependencies = [ - "cfg-if", - "once_cell", - "rustversion", - "wasm-bindgen-macro", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-macro" -version = "0.2.129" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e29d0c35b16e224a7eeb5cd2d25e3e1968fbd65604117b44d3b789d00ee8535" -dependencies = [ - "quote", - "wasm-bindgen-macro-support", -] - -[[package]] -name = "wasm-bindgen-macro-support" -version = "0.2.129" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f501a8bc3719dba86ef8ae4728879c08001bea749eb1333ac5b91e040e2a6b7" -dependencies = [ - "bumpalo", - "proc-macro2", - "quote", - "syn 3.0.9", - "wasm-bindgen-shared", -] - -[[package]] -name = "wasm-bindgen-shared" -version = "0.2.129" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23f0c9c52aa7cd7d77769a4cfe2a9adb1b331f489a41d912ce14513d5ab995c6" -dependencies = [ - "unicode-ident", -] - [[package]] name = "winapi-util" version = "0.1.11" diff --git a/Cargo.toml b/Cargo.toml index 02b0024..a448e87 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,6 @@ anyhow = "1.0" config = { version = "0.15.27", default-features = false } csv = "1.4.0" fastrand = "2" -flume = { version = "0.12.0", default-features = false, features = ["async", "eventual-fairness", "select"] } icmp-socket2 = { version = "0.6.4", features = ["tokio"] } irtt-client = { version = "0.7.0", optional = true, features = ["ancillary", "tokio"] } libc = "0.2" diff --git a/src/baseliner.rs b/src/baseliner.rs index 8363050..cb49820 100644 --- a/src/baseliner.rs +++ b/src/baseliner.rs @@ -9,12 +9,13 @@ use crate::SHUTDOWN; use crate::measurement::model::{MeasurementEvent, MeasurementObservation, MeasurementStream}; use crate::metrics::{Metric, MetricsSender}; use crate::settings::Settings; -use flume::{Receiver, Sender}; +use anyhow::Context; use log::{debug, info}; use std::collections::HashMap; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use tokio::sync::mpsc::{self, UnboundedReceiver}; const STALE_AFTER: Duration = Duration::from_secs(30); const ANOMALY_THRESHOLD_MS: f64 = 5_000.0; @@ -70,17 +71,15 @@ impl ReflectorState { #[derive(Clone, Debug)] pub struct ControlSnapshot { - pub generated_at: Instant, pub reflectors: HashMap, } pub struct Baseliner { pub settings: Settings, - pub reselect_trigger: Sender, + pub reselect_trigger: mpsc::Sender, pub start_time: Instant, - pub stats_rx: Receiver, - pub control_tx: Sender, - pub selection_tx: Option>, + pub stats_rx: UnboundedReceiver, + pub control_tx: tokio::sync::watch::Sender>, pub baseline_metrics: MetricsSender, pub event_metrics: MetricsSender, pub ping_metrics: MetricsSender, @@ -142,18 +141,14 @@ fn process_reply( (*state, anomaly) } -fn control_snapshot( - reflectors: &HashMap, - generated_at: Instant, -) -> ControlSnapshot { +fn control_snapshot(reflectors: &HashMap) -> ControlSnapshot { ControlSnapshot { - generated_at, reflectors: reflectors.clone(), } } impl Baseliner { - pub fn run(&self) -> anyhow::Result<()> { + pub fn run(&mut self) -> anyhow::Result<()> { let mut reflectors = HashMap::::new(); let slow_factor = ewma_factor( @@ -171,7 +166,10 @@ impl Baseliner { return Ok(()); } - let event = self.stats_rx.recv()?; + let event = self + .stats_rx + .blocking_recv() + .context("measurement channel closed")?; let reply = match event { MeasurementEvent::Observation(reply) => reply, @@ -232,11 +230,8 @@ impl Baseliner { reflector, state.recent.up, state.recent.down ); - let snapshot = control_snapshot(&reflectors, Instant::now()); - if let Some(selection_tx) = &self.selection_tx { - let _ = selection_tx.send(snapshot.clone()); - } - self.control_tx.send(snapshot)?; + let snapshot = control_snapshot(&reflectors); + self.control_tx.send(Some(snapshot))?; } } } @@ -454,32 +449,4 @@ mod tests { assert_eq!(updated.last_receive_at, received_at); assert_close(reflectors[&stream].recent.up, 200.0); } - - #[test] - fn snapshot_contains_updated_combined_state() { - let start = Instant::now(); - let received_at = start + Duration::from_secs(1); - let generated_at = received_at + Duration::from_millis(1); - let reflector = "192.0.2.1".parse().unwrap(); - let mut reflectors = HashMap::new(); - - let (updated, anomaly) = process_reply( - &mut reflectors, - &reply(reflector, 20.0, 30.0, received_at), - start, - 0.9, - 0.5, - ); - let snapshot = control_snapshot(&reflectors, generated_at); - let stream = reply(reflector, 0.0, 0.0, received_at).stream; - - assert!(!anomaly); - assert_eq!(snapshot.generated_at, generated_at); - assert_close( - snapshot.reflectors[&stream].baseline.down, - updated.baseline.down, - ); - assert_close(snapshot.reflectors[&stream].recent.up, updated.recent.up); - assert_eq!(snapshot.reflectors[&stream].last_receive_at, received_at); - } } diff --git a/src/controller.rs b/src/controller.rs index 4572985..7bc07bf 100644 --- a/src/controller.rs +++ b/src/controller.rs @@ -5,8 +5,11 @@ use std::{ time::Duration, }; -use flume::{Receiver, SendError, Sender}; use thiserror::Error; +use tokio::sync::{ + mpsc::UnboundedReceiver, + watch::{self, error::SendError}, +}; #[cfg(feature = "irtt")] use crate::measurement::irtt::IrttBinding; @@ -28,10 +31,10 @@ pub enum ControllerError { } pub struct Controller { - reflector_rx: Receiver>, - icmp_binding_tx: Sender>, + reflector_rx: UnboundedReceiver>, + icmp_binding_tx: watch::Sender>, #[cfg(feature = "irtt")] - irtt_binding_tx: Sender>, + irtt_binding_tx: watch::Sender>, reflectors: HashMap, active_streams: Arc>>, endpoint_streams: Arc>>, @@ -47,9 +50,9 @@ impl Controller { reason = "Constructor wires the controller channels and shared stream state" )] pub fn new( - reflector_rx: Receiver>, - icmp_binding_tx: Sender>, - #[cfg(feature = "irtt")] irtt_binding_tx: Sender>, + reflector_rx: UnboundedReceiver>, + icmp_binding_tx: watch::Sender>, + #[cfg(feature = "irtt")] irtt_binding_tx: watch::Sender>, reflectors: HashMap, active_streams: Arc>>, endpoint_streams: Arc>>, @@ -73,7 +76,7 @@ impl Controller { pub async fn run(mut self, initial_reflectors: Vec) -> Result<(), ControllerError> { self.update_reflectors(initial_reflectors).await?; - while let Ok(reflectors) = self.reflector_rx.recv_async().await { + while let Some(reflectors) = self.reflector_rx.recv().await { self.update_reflectors(reflectors).await?; } Ok(()) @@ -140,9 +143,9 @@ impl Controller { .endpoint_streams .write() .expect("endpoint stream lock poisoned") = endpoint_streams; - self.icmp_binding_tx.send_async(icmp).await?; + self.icmp_binding_tx.send(icmp)?; #[cfg(feature = "irtt")] - self.irtt_binding_tx.send_async(irtt).await?; + self.irtt_binding_tx.send(irtt)?; Ok(()) } } diff --git a/src/main.rs b/src/main.rs index 1bda3c4..38f2abf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -32,7 +32,6 @@ use crate::reflector_selector::ReflectorSelector; use crate::settings::{ReflectorDefinition, ReflectorProtocol, Settings}; use crate::util::RwLockExt; use ::log::{info, warn}; -use flume::RecvTimeoutError; use std::collections::HashMap; use std::net::{IpAddr, SocketAddr, ToSocketAddrs}; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; @@ -40,6 +39,8 @@ use std::sync::{Arc, RwLock}; use std::thread; use std::thread::{JoinHandle, sleep}; use std::time::{Duration, Instant}; +use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; +use tokio::sync::watch; struct ReflectorSetup { peers: Arc>>, @@ -266,7 +267,7 @@ fn setup_reflectors(settings: &Settings) -> anyhow::Result { fn spawn_task( rt: &tokio::runtime::Runtime, - error_tx: &flume::Sender, + error_tx: &UnboundedSender, task: F, ) -> tokio::task::JoinHandle<()> where @@ -284,7 +285,7 @@ where fn spawn_worker( name: &str, - error_tx: &flume::Sender, + error_tx: &UnboundedSender, worker: F, ) -> anyhow::Result> where @@ -302,16 +303,16 @@ where .map_err(|e| anyhow::anyhow!(e)) } -fn wait_for_exit(error_rx: &flume::Receiver) -> anyhow::Result<()> { +async fn wait_for_exit(error_rx: &mut UnboundedReceiver) -> anyhow::Result<()> { loop { - match error_rx.recv_timeout(Duration::from_secs(1)) { - Ok(error) => { + match tokio::time::timeout(Duration::from_secs(1), error_rx.recv()).await { + Ok(Some(error)) => { return Err(anyhow::anyhow!("worker exited with error: {error}")); } - Err(RecvTimeoutError::Disconnected) => { + Ok(None) => { return Ok(()); } - Err(RecvTimeoutError::Timeout) => { + Err(_) => { if SHUTDOWN.load(Ordering::Relaxed) { info!("Received shutdown signal"); return Ok(()); @@ -345,20 +346,14 @@ fn run(settings: &Settings) -> anyhow::Result<()> { reselection_enabled, } = setup_reflectors(settings)?; - let (control_snapshot_tx, control_snapshot_rx) = flume::unbounded(); - let (error_tx, error_rx) = flume::unbounded::(); - let (icmp_binding_tx, icmp_binding_rx) = flume::unbounded(); + let (control_snapshot_tx, control_snapshot_rx) = watch::channel(None); + let (error_tx, mut error_rx) = tokio::sync::mpsc::unbounded_channel::(); + let (icmp_binding_tx, icmp_binding_rx) = watch::channel(Vec::new()); #[cfg(feature = "irtt")] - let (irtt_binding_tx, irtt_binding_rx) = flume::unbounded(); - let (measurement_tx, measurement_rx) = flume::unbounded(); - let (reflector_tx, reflector_rx) = flume::unbounded(); - let (reselect_tx, reselect_rx) = flume::bounded(1); - let (selection_snapshot_tx, selection_snapshot_rx) = if reselection_enabled { - let (tx, rx) = flume::unbounded(); - (Some(tx), Some(rx)) - } else { - (None, None) - }; + let (irtt_binding_tx, irtt_binding_rx) = watch::channel(Vec::new()); + let (measurement_tx, measurement_rx) = tokio::sync::mpsc::unbounded_channel(); + let (reflector_tx, reflector_rx) = tokio::sync::mpsc::unbounded_channel(); + let (reselect_tx, reselect_rx) = tokio::sync::mpsc::channel(1); let source = match settings.advanced_settings.measurement_type { ReflectorProtocol::IcmpEcho => MeasurementSource::IcmpEcho, @@ -403,7 +398,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { let dropped = Arc::new(AtomicU32::new(0)); let (metrics_tx, metrics_thread_handle) = if settings.observability.enabled { - let (tx, rx) = flume::bounded(1000); + let (tx, rx) = std::sync::mpsc::sync_channel(1000); let metrics = Metrics { settings: settings.clone(), metrics_rx: rx, @@ -439,13 +434,12 @@ fn run(settings: &Settings) -> anyhow::Result<()> { }, }); - let baseliner = Baseliner { + let mut baseliner = Baseliner { settings: settings.clone(), reselect_trigger: reselect_tx.clone(), start_time, stats_rx: measurement_rx, control_tx: control_snapshot_tx, - selection_tx: selection_snapshot_tx, baseline_metrics, event_metrics: event_metrics.clone(), ping_metrics, @@ -466,8 +460,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { if reselection_enabled { let reflector_selector = ReflectorSelector { settings: settings.clone(), - snapshot_rx: selection_snapshot_rx - .expect("reselection snapshot receiver must exist when reselection is enabled"), + snapshot_rx: control_snapshot_rx.clone(), reflector_peers_lock: Arc::clone(&reflector_peers), reflector_pool, endpoint_streams, @@ -475,7 +468,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { metrics: event_metrics, reflector_tx, }; - spawn_worker("reselection", &error_tx, move || reflector_selector.run())?; + spawn_task(&rt, &error_tx, reflector_selector.run()); } // Give the baseliner time to collect initial samples before adjusting rates. @@ -494,7 +487,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { // Drop the original sender so the channel disconnects if all workers exit cleanly. drop(error_tx); - let result = wait_for_exit(&error_rx); + let result = rt.block_on(wait_for_exit(&mut error_rx)); // Make the error path request shutdown too, rather than leaving the other workers running. SHUTDOWN.store(true, Ordering::Relaxed); diff --git a/src/measurement/icmp.rs b/src/measurement/icmp.rs index cbdeb59..fa980b8 100644 --- a/src/measurement/icmp.rs +++ b/src/measurement/icmp.rs @@ -5,7 +5,6 @@ use std::{ time::{Duration, Instant, SystemTime}, }; -use flume::{Receiver, Sender}; use icmp_socket2::{ AsyncIcmpSocket, IcmpSocket4, Icmpv4Message, Icmpv4Packet, packet::{IcmpPacketBuildError, WithEchoRequest, WithTimestampRequest}, @@ -13,6 +12,7 @@ use icmp_socket2::{ }; use log::debug; use thiserror::Error; +use tokio::sync::{mpsc::UnboundedSender, watch}; use crate::{ measurement::model::{ @@ -24,6 +24,7 @@ use crate::{ time::{ClockSample, Realtime}, }; +#[derive(Clone)] pub struct IcmpBinding { pub stream: MeasurementStream, pub schedule: ProbeSchedule, @@ -103,8 +104,8 @@ pub struct IcmpEngine { origin: Instant, - binding_rx: Receiver>, - observation_tx: Sender, + binding_rx: watch::Receiver>, + observation_tx: UnboundedSender, } async fn wait_for_deadline(deadline: Option) { @@ -123,8 +124,8 @@ impl IcmpEngine { pub fn new( origin: Instant, #[allow(unused)] settings: &NetworkSettings, - observation_tx: Sender, - binding_rx: Receiver>, + observation_tx: UnboundedSender, + binding_rx: watch::Receiver>, ) -> Result { #[allow(unused_mut)] let mut socket = IcmpSocket4::new()?; @@ -420,12 +421,11 @@ impl IcmpEngine { for probe in expired { self.observation_tx - .send_async(MeasurementEvent::Loss(MeasurementLoss { + .send(MeasurementEvent::Loss(MeasurementLoss { stream: probe.stream, started_at: probe.sent_at, detected_at: now, })) - .await .map_err(|_| IcmpError::EventChannelClosed)?; } @@ -446,10 +446,10 @@ impl IcmpEngine { let deadline = self.next_deadline(); tokio::select! { - bindings = self.binding_rx.recv_async() => { - let bindings = bindings - .map_err(|_| IcmpError::BindingChannelClosed)?; + bindings = self.binding_rx.changed() => { + bindings.map_err(|_| IcmpError::BindingChannelClosed)?; + let bindings = self.binding_rx.borrow_and_update().clone(); self.set_bindings(bindings)?; } @@ -468,8 +468,7 @@ impl IcmpEngine { ) { Ok(observation) => { self.observation_tx - .send_async(MeasurementEvent::Observation(observation)) - .await + .send(MeasurementEvent::Observation(observation)) .map_err(|_| IcmpError::EventChannelClosed)?; } diff --git a/src/measurement/irtt.rs b/src/measurement/irtt.rs index ad6b818..51957e7 100644 --- a/src/measurement/irtt.rs +++ b/src/measurement/irtt.rs @@ -1,6 +1,5 @@ use std::{collections::HashMap, net::SocketAddr}; -use flume::{Receiver, Sender}; use irtt_client::{ ClientConfig, ClientEvent, NegotiationPolicy, OpenPolicy, managed::{ @@ -12,7 +11,7 @@ use irtt_client::{ }; use log::{debug, trace, warn}; use thiserror::Error; -use tokio::sync::broadcast::error::RecvError; +use tokio::sync::{broadcast::error::RecvError, mpsc::UnboundedSender, watch}; use crate::SHUTDOWN; use crate::measurement::model::{ @@ -51,6 +50,7 @@ pub enum IrttError { ManagedAbandoned, } +#[derive(Clone)] pub struct IrttBinding { stream: MeasurementStream, port: u16, @@ -69,16 +69,16 @@ pub struct IrttEngine { bindings: HashMap, targets: Vec, - binding_rx: Receiver>, - observation_tx: Sender, + binding_rx: watch::Receiver>, + observation_tx: UnboundedSender, } impl IrttEngine { pub fn new( interval: std::time::Duration, settings: &NetworkSettings, - observation_tx: Sender, - binding_rx: Receiver>, + observation_tx: UnboundedSender, + binding_rx: watch::Receiver>, ) -> Result<(ManagedClientTask, Self), IrttError> { let config = managed_client_config(interval, settings); let (task, handle) = ManagedClient::task(config, Vec::new())?; @@ -158,10 +158,10 @@ impl IrttEngine { retry.tick().await; loop { tokio::select! { - bindings = self.binding_rx.recv_async() => { - let bindings = bindings - .map_err(|_| IrttError::BindingChannelClosed)?; + bindings = self.binding_rx.changed() => { + bindings.map_err(|_| IrttError::BindingChannelClosed)?; + let bindings = self.binding_rx.borrow_and_update().clone(); self.set_bindings(bindings).await?; } @@ -273,8 +273,7 @@ impl IrttEngine { }; self.observation_tx - .send_async(MeasurementEvent::Observation(observation)) - .await + .send(MeasurementEvent::Observation(observation)) .map_err(|_| IrttError::EventChannelClosed)?; } @@ -294,8 +293,7 @@ impl IrttEngine { }; self.observation_tx - .send_async(MeasurementEvent::Loss(loss)) - .await + .send(MeasurementEvent::Loss(loss)) .map_err(|_| IrttError::EventChannelClosed)?; } diff --git a/src/metrics.rs b/src/metrics.rs index fac27c9..3f71143 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -6,12 +6,14 @@ use crate::measurement::model::MeasurementSource; use crate::settings::ObservabilityProtocol; use crate::settings::Settings; use crate::time::{ClockSample, Realtime}; -use flume::{Receiver, RecvTimeoutError, Sender}; use log::{error, info, warn}; use std::fmt::Write; use std::net::{IpAddr, TcpStream, UdpSocket}; use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::mpsc; +use std::sync::mpsc::RecvTimeoutError; +use std::sync::mpsc::SyncSender; use std::time::Duration; const MAX_RECONNECT_BACKOFF: u64 = 60; @@ -127,12 +129,12 @@ pub enum Metric { #[derive(Clone)] pub struct MetricsSender { - tx: Option>, + tx: Option>, dropped: Arc, } impl MetricsSender { - pub fn new(tx: Sender<(Metric, u64)>, dropped: Arc) -> Self { + pub fn new(tx: SyncSender<(Metric, u64)>, dropped: Arc) -> Self { Self { tx: Some(tx), dropped, @@ -159,7 +161,7 @@ impl MetricsSender { pub struct Metrics { pub settings: Settings, pub metrics_dropped: Arc, - pub metrics_rx: Receiver<(Metric, u64)>, + pub metrics_rx: mpsc::Receiver<(Metric, u64)>, } impl Metrics { diff --git a/src/ratecontroller.rs b/src/ratecontroller.rs index 867e299..e6363ca 100644 --- a/src/ratecontroller.rs +++ b/src/ratecontroller.rs @@ -16,13 +16,13 @@ use crate::platform::{ use crate::settings::Settings; use crate::time::{ClockSample, Realtime}; use crate::util::{ArcRwLock, RwLockExt}; -use flume::{Receiver, Sender}; use log::{debug, info, warn}; use std::fs::File; use std::io::Write; use std::sync::atomic::Ordering; use std::thread::sleep; use std::time::{Duration, Instant}; +use tokio::sync::mpsc; #[derive(Copy, Clone, Debug, PartialEq)] enum Direction { @@ -84,20 +84,6 @@ impl State { } } -fn drain_latest_snapshot( - snapshot_rx: &Receiver, - latest: &mut Option, -) { - while let Ok(snapshot) = snapshot_rx.try_recv() { - if latest - .as_ref() - .is_none_or(|current| snapshot.generated_at >= current.generated_at) - { - *latest = Some(snapshot); - } - } -} - fn deltas_from_snapshot( snapshot: Option<&ControlSnapshot>, reflectors: &[MeasurementStream], @@ -137,10 +123,9 @@ fn deltas_from_snapshot( pub struct Ratecontroller { settings: Settings, - control_snapshot: Option, - control_snapshot_rx: Receiver, + control_snapshot_rx: tokio::sync::watch::Receiver>, reflectors_lock: ArcRwLock>, - reselect_trigger: Sender, + reselect_trigger: mpsc::Sender, state_dl: State, state_ul: State, rng: fastrand::Rng, @@ -156,9 +141,9 @@ impl Ratecontroller { )] fn new_with_backends( settings: Settings, - control_snapshot_rx: Receiver, + control_snapshot_rx: tokio::sync::watch::Receiver>, reflectors_lock: ArcRwLock>, - reselect_trigger: Sender, + reselect_trigger: mpsc::Sender, metrics: MetricsSender, mut stats_provider: S, mut traffic_control: T, @@ -183,7 +168,6 @@ impl Ratecontroller { Ok(Self { settings, - control_snapshot: None, control_snapshot_rx, reflectors_lock, reselect_trigger, @@ -483,15 +467,17 @@ impl Ratecontroller { } fn update_deltas(&mut self) -> anyhow::Result<()> { - drain_latest_snapshot(&self.control_snapshot_rx, &mut self.control_snapshot); let now_t = Instant::now(); let reflectors = self.reflectors_lock.read_anyhow()?; - let (down_deltas, up_deltas) = deltas_from_snapshot( - self.control_snapshot.as_ref(), - &reflectors, - now_t, - self.settings.advanced_settings.tick_interval, - ); + let (down_deltas, up_deltas) = { + let snapshot = self.control_snapshot_rx.borrow(); + deltas_from_snapshot( + snapshot.as_ref(), + &reflectors, + now_t, + self.settings.advanced_settings.tick_interval, + ) + }; self.state_dl.deltas = down_deltas; self.state_ul.deltas = up_deltas; @@ -517,9 +503,9 @@ impl Ratecontroller { impl Ratecontroller { pub fn new( settings: Settings, - control_snapshot_rx: Receiver, + control_snapshot_rx: tokio::sync::watch::Receiver>, reflectors_lock: ArcRwLock>, - reselect_trigger: Sender, + reselect_trigger: mpsc::Sender, metrics: MetricsSender, ) -> anyhow::Result { Self::new_with_backends( @@ -534,111 +520,3 @@ impl Ratecontroller { ) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::baseliner::{EwmaStats, ReflectorState}; - use crate::measurement::model::{MeasurementSource, MeasurementStreamId}; - use std::collections::HashMap; - use std::net::IpAddr; - - fn stream(id: u64, peer: IpAddr) -> MeasurementStream { - MeasurementStream { - id: MeasurementStreamId(id), - peer, - source: MeasurementSource::IcmpTimestamp, - } - } - - #[test] - fn fixed_seed_reproduces_initial_safe_rates() { - let mut first = fastrand::Rng::with_seed(42); - let mut second = fastrand::Rng::with_seed(42); - - let first_rates = generate_initial_speeds(&mut first, 100_000.0, 100); - let second_rates = generate_initial_speeds(&mut second, 100_000.0, 100); - - assert_eq!(first_rates, second_rates); - } - - fn reflector_state( - baseline_down: f64, - baseline_up: f64, - recent_down: f64, - recent_up: f64, - last_receive_at: Instant, - ) -> ReflectorState { - ReflectorState { - baseline: EwmaStats { - down: baseline_down, - up: baseline_up, - }, - recent: EwmaStats { - down: recent_down, - up: recent_up, - }, - last_receive_at, - } - } - - fn snapshot(generated_at: Instant, marker: f64) -> ControlSnapshot { - let reflector = stream(0, "192.0.2.1".parse().unwrap()); - ControlSnapshot { - generated_at, - reflectors: HashMap::from([( - reflector, - reflector_state(marker, marker, marker, marker, generated_at), - )]), - } - } - - #[test] - fn drains_queued_snapshots_and_keeps_the_newest() { - let start = Instant::now(); - let newest_at = start + Duration::from_secs(2); - let (tx, rx) = flume::unbounded(); - tx.send(snapshot(start + Duration::from_secs(1), 1.0)) - .unwrap(); - tx.send(snapshot(newest_at, 2.0)).unwrap(); - let mut latest = None; - - drain_latest_snapshot(&rx, &mut latest); - - let latest = latest.unwrap(); - assert!(rx.is_empty()); - assert_eq!(latest.generated_at, newest_at); - assert_eq!( - latest.reflectors[&stream(0, "192.0.2.1".parse().unwrap())] - .recent - .down, - 2.0 - ); - } - - #[test] - fn calculates_sorted_deltas_from_active_fresh_snapshot_state() { - let now = Instant::now(); - let first = stream(0, "192.0.2.1".parse().unwrap()); - let second = stream(1, "192.0.2.2".parse().unwrap()); - let stale = stream(2, "192.0.2.3".parse().unwrap()); - let inactive = stream(3, "192.0.2.4".parse().unwrap()); - let snapshot = ControlSnapshot { - generated_at: now, - reflectors: HashMap::from([ - (first, reflector_state(10.0, 20.0, 14.0, 28.0, now)), - (second, reflector_state(10.0, 20.0, 11.0, 22.0, now)), - ( - stale, - reflector_state(10.0, 20.0, 110.0, 220.0, now - Duration::from_secs(3)), - ), - (inactive, reflector_state(10.0, 20.0, 60.0, 80.0, now)), - ]), - }; - - let (down, up) = deltas_from_snapshot(Some(&snapshot), &[first, second, stale], now, 1.0); - - assert_eq!(down, vec![1.0, 4.0]); - assert_eq!(up, vec![2.0, 8.0]); - } -} diff --git a/src/reflector_selector.rs b/src/reflector_selector.rs index e1993a6..b173358 100644 --- a/src/reflector_selector.rs +++ b/src/reflector_selector.rs @@ -11,93 +11,13 @@ use crate::measurement::model::MeasurementStream; use crate::metrics::{Metric, MetricsSender}; use crate::settings::Settings; use crate::util::RwLockExt; -use flume::{Receiver, RecvError, RecvTimeoutError, Selector, Sender}; use log::{debug, info}; use std::net::SocketAddr; use std::sync::atomic::Ordering; use std::sync::{Arc, RwLock}; -use std::thread::sleep; use std::time::{Duration, Instant}; - -enum SelectorEvent { - Trigger(Result), - Snapshot(Result), -} - -#[derive(Debug, Eq, PartialEq)] -enum WaitOutcome { - Triggered, - DeadlineReached, - ChannelsClosed, -} - -fn wait_for_trigger_or_deadline( - trigger_rx: &Receiver, - snapshot_rx: &Receiver, - latest_snapshot: &mut Option, - 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, - latest_snapshot: &mut Option, - 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; - } - } - } -} +use tokio::sync::mpsc::{self, UnboundedSender}; +use tokio::sync::watch; fn recent_rtt_for_peer( snapshot: Option<&ControlSnapshot>, @@ -113,29 +33,24 @@ fn recent_rtt_for_peer( pub struct ReflectorSelector { pub settings: Settings, - pub snapshot_rx: Receiver, + pub snapshot_rx: watch::Receiver>, pub reflector_peers_lock: Arc>>, - pub reflector_tx: Sender>, + pub reflector_tx: UnboundedSender>, pub reflector_pool: Vec, pub endpoint_streams: Arc>>, - pub trigger_channel: Receiver, + pub trigger_channel: mpsc::Receiver, pub metrics: MetricsSender, } impl ReflectorSelector { - pub fn run(&self) -> anyhow::Result<()> { + pub async fn run(mut 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.settings.advanced_settings.tick_interval * std::f64::consts::PI, ); // Initial wait of several seconds to allow some OWD data to build up - receive_snapshots_until( - &self.snapshot_rx, - &mut latest_snapshot, - Instant::now() + baseline_sleep_time, - ); + tokio::time::sleep(baseline_sleep_time).await; loop { if SHUTDOWN.load(Ordering::Relaxed) { @@ -143,18 +58,14 @@ impl ReflectorSelector { return 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(()); + let triggered = tokio::select! { + trigger = self.trigger_channel.recv() => { + match trigger { + Some(_) => true, + None => return Ok(()), + } } + _ = tokio::time::sleep(selector_sleep_time) => false, }; reselection_count += 1; info!("Starting reselection [#{}]", reselection_count); @@ -172,54 +83,51 @@ impl ReflectorSelector { } let mut next_peers: Vec = Vec::new(); - let mut reflectors_peers = self.reflector_peers_lock.write_anyhow()?; - let current_peers = reflectors_peers.clone(); + let current_peers = { + let mut reflectors_peers = self.reflector_peers_lock.write_anyhow()?; + let current_peers = reflectors_peers.clone(); - // Include all current peers - for reflector in reflectors_peers.iter() { - debug!("Current peer: {}", reflector); - next_peers.push(*reflector); - } + // Include all current peers + for reflector in reflectors_peers.iter() { + debug!("Current peer: {}", reflector); + next_peers.push(*reflector); + } - if self.reflector_pool.is_empty() { - continue; - } - for _ in 0..20 { - let next_candidate = - &self.reflector_pool[fastrand::usize(..self.reflector_pool.len())]; - if next_peers.contains(next_candidate) { + if self.reflector_pool.is_empty() { continue; } - debug!("Next candidate: {}", next_candidate); - next_peers.push(*next_candidate); - } + 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(); - self.reflector_tx.send(next_peers.clone())?; + // Clone next_peers because we need it again after the baseline sleep + // to iterate over candidates for RTT measurement. + *reflectors_peers = next_peers.clone(); + self.reflector_tx.send(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); + current_peers + }; debug!("Waiting for candidates to be baselined"); // Wait for several seconds to allow all reflectors to be re-baselined - receive_snapshots_until( - &self.snapshot_rx, - &mut latest_snapshot, - Instant::now() + baseline_sleep_time, - ); + tokio::time::sleep(baseline_sleep_time).await; // Re-acquire the lock when we wake up again - reflectors_peers = self.reflector_peers_lock.write_anyhow()?; + let mut reflectors_peers = self.reflector_peers_lock.write_anyhow()?; let mut candidates = Vec::new(); + let snapshot = self.snapshot_rx.borrow(); for peer in next_peers { let stream = self.endpoint_streams.read_anyhow()?.get(&peer).copied(); if let Some(rtt) = recent_rtt_for_peer( - latest_snapshot.as_ref(), + snapshot.as_ref(), stream.as_ref(), Instant::now(), Duration::from_secs_f64(self.settings.advanced_settings.tick_interval * 2.0), @@ -285,158 +193,3 @@ impl ReflectorSelector { } } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::baseliner::{EwmaStats, ReflectorState}; - use crate::measurement::model::MeasurementSource; - use std::collections::HashMap; - use std::net::IpAddr; - 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( - crate::measurement::model::MeasurementStream { - id: crate::measurement::model::MeasurementStreamId(0), - peer, - source: MeasurementSource::IcmpTimestamp, - }, - 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: IpAddr = "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(), - latest_snapshot.as_ref().unwrap().reflectors.keys().next(), - generated_at + Duration::from_secs(1), - Duration::from_secs(2) - ), - Some(42) - ); - } -}