From 89c7d75947b818ff7a2b0ff9f3d5bee070c53029 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Mon, 5 Oct 2026 01:15:19 +0700 Subject: [PATCH] refactor(measurement): migrate ICMP engine to Tokio async runtime --- Cargo.lock | 133 ++++++++++++++++++ Cargo.toml | 13 +- src/main.rs | 4 +- src/measurement/icmp.rs | 301 +++++++++++++++------------------------- 4 files changed, 257 insertions(+), 194 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bde2836..c996483 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,17 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "async-trait" +version = "0.1.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -55,6 +66,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + [[package]] name = "cc" version = "1.3.0" @@ -183,9 +200,23 @@ 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.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + [[package]] name = "getrandom" version = "0.4.3" @@ -217,9 +248,11 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "816dc56a4f7f76e42b0616818600b05528cd346cb3843276e82df982d5c3ee26" dependencies = [ + "async-trait", "nix", "socket2", "thiserror", + "tokio", ] [[package]] @@ -325,6 +358,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "mio" +version = "1.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1788edb87fdc09c7e26304471e2f5be8cdefb1b6930d6e3985fc02ff53bf86ee" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + [[package]] name = "netlink-bindings" version = "0.3.5" @@ -371,12 +415,41 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "pathdiff" version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "prettyplease" version = "0.2.37" @@ -411,6 +484,15 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "regex" version = "1.13.1" @@ -552,6 +634,22 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "smallvec" +version = "1.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e" + [[package]] name = "socket2" version = "0.6.5" @@ -589,6 +687,7 @@ dependencies = [ "rustix", "serde", "thiserror", + "tokio", "uci-derive", ] @@ -655,6 +754,34 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "tokio" +version = "1.53.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e95f91fcc7a621e8b030f6aa23c71fe9838ae2fb4d8118b75602a328f5144044" +dependencies = [ + "bytes", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.6", +] + [[package]] name = "toml" version = "1.1.3+spec-1.1.0" @@ -725,6 +852,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "wasi" +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.126" diff --git a/Cargo.toml b/Cargo.toml index 120c8d5..cd61932 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,15 +20,16 @@ strip = "symbols" [dependencies] anyhow = "1.0" -icmp-socket2 = "0.6.4" -fastrand = "2" -rustix = { version = "1.1", features = ["fs", "system", "thread", "time"] } -thiserror = "2.0" -libc = "0.2" -flume = { version = "0.12.0", default-features = false, features = ["eventual-fairness", "select"] } config = { version = "0.15.25", 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"] } +libc = "0.2" +rustix = { version = "1.1", features = ["fs", "system", "thread", "time"] } serde = { version = "1", features = ["derive"] } +thiserror = "2.0" +tokio = { version = "1", features = ["full"] } uci-derive = { path = "crates/uci-derive" } [dependencies.log] diff --git a/src/main.rs b/src/main.rs index 2c15b45..40b0281 100644 --- a/src/main.rs +++ b/src/main.rs @@ -264,7 +264,9 @@ fn run(settings: &Settings) -> anyhow::Result<()> { offset: Duration::from_secs(0), }, )])?; - engine.run()?; + + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.spawn(engine.run()); let ReflectorSetup { peers: reflector_peers, diff --git a/src/measurement/icmp.rs b/src/measurement/icmp.rs index 80d2099..22a6d49 100644 --- a/src/measurement/icmp.rs +++ b/src/measurement/icmp.rs @@ -2,20 +2,18 @@ use std::{ collections::{HashMap, HashSet}, io, net::IpAddr, - sync::{Arc, Mutex, atomic::Ordering}, - thread, time::{Duration, Instant, SystemTime}, }; -use flume::{Receiver, RecvTimeoutError, Sender}; +use flume::{Receiver, Sender}; use icmp_socket2::{ - IcmpSocket, IcmpSocket4, Icmpv4Message, Icmpv4Packet, + AsyncIcmpSocket, IcmpSocket4, Icmpv4Message, Icmpv4Packet, packet::{IcmpPacketBuildError, WithEchoRequest, WithTimestampRequest}, + tokio::AsyncIcmpV4Socket, }; use thiserror::Error; use crate::{ - SHUTDOWN, measurement::model::{ MeasurementEvent, MeasurementLoss, MeasurementObservation, MeasurementSource, MeasurementStream, MeasurementStreamId, OneWayClock, OneWayLatency, ProbeSchedule, @@ -78,6 +76,9 @@ pub enum IcmpError { #[error("event channel closed")] EventChannelClosed, + + #[error("binding channel closed")] + BindingChannelClosed, } #[derive(Debug, Clone, Copy)] @@ -87,28 +88,33 @@ struct InFlightProbe { } type InFlightProbeKey = (IpAddr, MeasurementSource, u16); -type InFlightProbes = Mutex>; +type InFlightProbes = HashMap; + +pub struct IcmpEngine { + socket: AsyncIcmpV4Socket, -struct IcmpSender { - socket: IcmpSocket4, identifier: u16, next_sequence: u16, - bindings: HashMap, - inflight: Arc, - origin: Instant, - binding_rx: Receiver>, -} -struct IcmpReceiver { - socket: IcmpSocket4, - identifier: u16, - inflight: Arc, + bindings: HashMap, + inflight: HashMap, + + origin: Instant, + + binding_rx: Receiver>, observation_tx: Sender, } -pub struct IcmpEngine { - sender: IcmpSender, - receiver: IcmpReceiver, +async fn wait_for_deadline(deadline: Option) { + match deadline { + Some(deadline) => { + tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; + } + + None => { + std::future::pending::<()>().await; + } + } } impl IcmpEngine { @@ -119,80 +125,39 @@ impl IcmpEngine { binding_rx: Receiver>, ) -> Result { #[allow(unused_mut)] - let mut rx_socket = IcmpSocket4::new()?; + let mut socket = IcmpSocket4::new()?; #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] if let Some(device) = &settings.measurement_bind_device { - IcmpSocket::bind_device(&mut rx_socket, device)?; + icmp_socket2::IcmpSocket::bind_device(&mut socket, device)?; } #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))] if let Some(mark) = settings.measurement_mark { - IcmpSocket::set_mark(&mut rx_socket, mark)?; + icmp_socket2::IcmpSocket::set_mark(&mut socket, mark)?; } #[cfg(target_os = "freebsd")] if let Some(fib) = settings.measurement_fib { - IcmpSocket::set_fib(&mut rx_socket, fib)?; + icmp_socket2::IcmpSocket::set_fib(&mut socket, fib)?; } - rx_socket.enable_receive_metadata()?; - let tx_socket = rx_socket.try_clone()?; + socket.enable_receive_metadata()?; let identifier = (std::process::id() & 0xffff) as u16; - let inflight: Arc = Arc::new(Mutex::new(HashMap::new())); + let inflight: InFlightProbes = HashMap::new(); Ok(Self { - sender: IcmpSender { - socket: tx_socket, - identifier, - next_sequence: 0, - bindings: HashMap::new(), - inflight: Arc::clone(&inflight), - origin, - binding_rx, - }, - receiver: IcmpReceiver { - socket: rx_socket, - identifier, - inflight, - observation_tx, - }, + socket: socket.into_tokio()?, + identifier, + next_sequence: 0, + bindings: HashMap::new(), + inflight, + origin, + binding_rx, + observation_tx, }) } - pub fn run(self) -> Result<(), IcmpError> { - let Self { sender, receiver } = self; - - let sender_handle = thread::spawn(move || { - let result = sender.run(); - SHUTDOWN.store(true, Ordering::Relaxed); - result - }); - let sender_thread = sender_handle.thread().clone(); - - let receiver_handle = thread::spawn(move || { - let result = receiver.run(); - SHUTDOWN.store(true, Ordering::Relaxed); - sender_thread.unpark(); - result - }); - - let sender_result = sender_handle.join().expect("ICMP sender thread panicked"); - - SHUTDOWN.store(true, Ordering::Relaxed); - - let receiver_result = receiver_handle - .join() - .expect("ICMP receiver thread panicked"); - - sender_result?; - receiver_result?; - - Ok(()) - } -} - -impl IcmpSender { fn set_bindings(&mut self, desired: Vec) -> Result<(), IcmpError> { // Validate the complete desired state before mutating anything. for binding in &desired { @@ -249,41 +214,7 @@ impl IcmpSender { Ok(()) } - fn run(mut self) -> Result<(), IcmpError> { - while !SHUTDOWN.load(Ordering::Relaxed) { - let deadline = self.next_deadline(); - - match deadline { - Some(deadline) => match self.binding_rx.recv_deadline(deadline) { - Ok(bindings) => { - self.set_bindings(bindings)?; - } - - Err(RecvTimeoutError::Timeout) => { - self.send_due_probes(Instant::now())?; - } - - Err(RecvTimeoutError::Disconnected) => { - break; - } - }, - - None => { - // Nothing scheduled. Don't exit. - // Just sleep until configuration arrives. - let bindings = self - .binding_rx - .recv() - .map_err(|_| IcmpError::EventChannelClosed)?; - self.set_bindings(bindings)?; - } - } - } - - Ok(()) - } - - fn send_due_probes(&mut self, now: Instant) -> Result<(), IcmpError> { + async fn send_due_probes(&mut self, now: Instant) -> Result<(), IcmpError> { let due = self .bindings .iter() @@ -291,7 +222,7 @@ impl IcmpSender { .collect::>(); for stream_id in due { - self.send_probe(stream_id)?; + self.send_probe(stream_id).await?; let binding = self .bindings @@ -306,7 +237,7 @@ impl IcmpSender { Ok(()) } - fn send_probe(&mut self, stream_id: MeasurementStreamId) -> Result<(), IcmpError> { + async fn send_probe(&mut self, stream_id: MeasurementStreamId) -> Result<(), IcmpError> { let stream = self .bindings .get(&stream_id) @@ -333,12 +264,9 @@ impl IcmpSender { let sent_at = Instant::now(); let key = (stream.peer, stream.source, sequence); - let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned"); - inflight.insert(key, InFlightProbe { stream, sent_at }); - if let Err(err) = self.socket.send_to(peer, packet) { - return Err(err.into()); - } + self.socket.send_to(peer, packet).await?; + self.inflight.insert(key, InFlightProbe { stream, sent_at }); Ok(()) } @@ -349,59 +277,6 @@ impl IcmpSender { .map(|binding| binding.next_probe_at) .min() } -} - -impl IcmpReceiver { - fn run(mut self) -> Result<(), IcmpError> { - const RECEIVE_TIMEOUT: Duration = Duration::from_millis(50); - const HOUSEKEEPING_INTERVAL: Duration = Duration::from_millis(500); - - self.socket.set_timeout(Some(RECEIVE_TIMEOUT)); - - let mut next_housekeeping = Instant::now() + HOUSEKEEPING_INTERVAL; - - while !SHUTDOWN.load(Ordering::Relaxed) { - match self.socket.rcv_from_with_meta() { - Ok(result) => { - match self.handle_reply( - result.packet, - result.peer.as_socket().unwrap().ip(), - result.received_at, - result.kernel_rx_timestamp, - ) { - Ok(observation) => { - self.observation_tx - .send(MeasurementEvent::Observation(observation)) - .map_err(|_| IcmpError::EventChannelClosed)?; - } - Err( - IcmpError::WrongIdentifier { .. } - | IcmpError::NoMatchingProbe { .. } - | IcmpError::InvalidPacket(_), - ) => { - // Not one of ours, malformed, or a late reply. - } - Err(err) => return Err(err), - } - } - Err(err) - if matches!( - err.kind(), - io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut - ) => {} - Err(err) => return Err(err.into()), - } - - let now = Instant::now(); - - if now >= next_housekeeping { - self.expire_inflight(now)?; - next_housekeeping = now + HOUSEKEEPING_INTERVAL; - } - } - - Ok(()) - } fn handle_reply( &mut self, @@ -424,13 +299,10 @@ impl IcmpReceiver { } let key = (peer, MeasurementSource::IcmpEcho, sequence); - let probe = { - let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned"); - - inflight - .remove(&key) - .ok_or(IcmpError::NoMatchingProbe(key))? - }; + let probe = self + .inflight + .remove(&key) + .ok_or(IcmpError::NoMatchingProbe(key))?; Ok(MeasurementObservation { stream: probe.stream, @@ -462,13 +334,10 @@ impl IcmpReceiver { } let key = (peer, MeasurementSource::IcmpTimestamp, sequence); - let probe = { - let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned"); - - inflight - .remove(&key) - .ok_or(IcmpError::NoMatchingProbe(key))? - }; + let probe = self + .inflight + .remove(&key) + .ok_or(IcmpError::NoMatchingProbe(key))?; let now_ms = if let Some(timestamp) = kernel_rx_timestamp { ClockSample::::from(timestamp).as_time_since_midnight() @@ -496,14 +365,13 @@ impl IcmpReceiver { } } - fn expire_inflight(&self, now: Instant) -> Result<(), IcmpError> { + fn expire_inflight(&mut self, now: Instant) -> Result<(), IcmpError> { const PROBE_TIMEOUT: Duration = Duration::from_secs(2); let expired = { - let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned"); let mut expired = Vec::new(); - inflight.retain(|_, probe| { + self.inflight.retain(|_, probe| { if now.saturating_duration_since(probe.sent_at) >= PROBE_TIMEOUT { expired.push(*probe); false @@ -527,6 +395,65 @@ impl IcmpReceiver { Ok(()) } + + pub async fn run(mut self) -> Result<(), IcmpError> { + const HOUSEKEEPING_INTERVAL: Duration = Duration::from_millis(500); + + let mut housekeeping = tokio::time::interval_at( + tokio::time::Instant::now() + HOUSEKEEPING_INTERVAL, + HOUSEKEEPING_INTERVAL, + ); + + housekeeping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + let deadline = self.next_deadline(); + + tokio::select! { + bindings = self.binding_rx.recv_async() => { + let bindings = bindings + .map_err(|_| IcmpError::BindingChannelClosed)?; + + self.set_bindings(bindings)?; + } + + _ = wait_for_deadline(deadline) => { + self.send_due_probes(Instant::now()).await?; + } + + result = self.socket.rcv_from_with_meta() => { + let result = result?; + + match self.handle_reply( + result.packet, + result.peer.as_socket().unwrap().ip(), + result.received_at, + result.kernel_rx_timestamp, + ) { + Ok(observation) => { + self.observation_tx + .send(MeasurementEvent::Observation(observation)) + .map_err(|_| IcmpError::EventChannelClosed)?; + } + + Err( + IcmpError::WrongIdentifier { .. } + | IcmpError::NoMatchingProbe { .. } + | IcmpError::InvalidPacket(_), + ) => { + // Not ours, malformed, or late. + } + + Err(err) => return Err(err), + } + } + + _ = housekeeping.tick() => { + self.expire_inflight(Instant::now())?; + } + } + } + } } const MS_PER_DAY: i64 = 86_400_000;