refactor(measurement): migrate ICMP engine to Tokio async runtime

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"

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]

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,

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<HashMap<InFlightProbeKey, InFlightProbe>>;
+type InFlightProbes = HashMap<InFlightProbeKey, InFlightProbe>;
+
+pub struct IcmpEngine {
+    socket: AsyncIcmpV4Socket,
 
-struct IcmpSender {
-    socket: IcmpSocket4,
     identifier: u16,
     next_sequence: u16,
-    bindings: HashMap<MeasurementStreamId, IcmpBinding>,
-    inflight: Arc<InFlightProbes>,
-    origin: Instant,
-    binding_rx: Receiver<Vec<IcmpBinding>>,
-}
 
-struct IcmpReceiver {
+    bindings: HashMap<MeasurementStreamId, IcmpBinding>,
-    socket: IcmpSocket4,
+    inflight: HashMap<InFlightProbeKey, InFlightProbe>,
-    identifier: u16,
+
-    inflight: Arc<InFlightProbes>,
+    origin: Instant,
+
+    binding_rx: Receiver<Vec<IcmpBinding>>,
     observation_tx: Sender<MeasurementEvent>,
 }
 
-pub struct IcmpEngine {
+async fn wait_for_deadline(deadline: Option<Instant>) {
-    sender: IcmpSender,
+    match deadline {
-    receiver: IcmpReceiver,
+        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<Vec<IcmpBinding>>,
     ) -> Result<Self, IcmpError> {
         #[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()?;
+        socket.enable_receive_metadata()?;
-        let tx_socket = rx_socket.try_clone()?;
         let identifier = (std::process::id() & 0xffff) as u16;
-        let inflight: Arc<InFlightProbes> = Arc::new(Mutex::new(HashMap::new()));
+        let inflight: InFlightProbes = HashMap::new();
 
         Ok(Self {
-            sender: IcmpSender {
+            socket: socket.into_tokio()?,
-                socket: tx_socket,
+            identifier,
-                identifier,
+            next_sequence: 0,
-                next_sequence: 0,
+            bindings: HashMap::new(),
-                bindings: HashMap::new(),
+            inflight,
-                inflight: Arc::clone(&inflight),
+            origin,
-                origin,
+            binding_rx,
-                binding_rx,
+            observation_tx,
-            },
-            receiver: IcmpReceiver {
-                socket: rx_socket,
-                identifier,
-                inflight,
-                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<IcmpBinding>) -> 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> {
+    async fn send_due_probes(&mut self, now: Instant) -> 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> {
         let due = self
             .bindings
             .iter()
@@ -291,7 +222,7 @@ impl IcmpSender {
             .collect::<Vec<_>>();
 
         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) {
+        self.socket.send_to(peer, packet).await?;
-            return Err(err.into());
+        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 probe = self
-                    let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned");
+                    .inflight
-
+                    .remove(&key)
-                    inflight
+                    .ok_or(IcmpError::NoMatchingProbe(key))?;
-                        .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 probe = self
-                    let mut inflight = self.inflight.lock().expect("ICMP inflight mutex poisoned");
+                    .inflight
-
+                    .remove(&key)
-                    inflight
+                    .ok_or(IcmpError::NoMatchingProbe(key))?;
-                        .remove(&key)
-                        .ok_or(IcmpError::NoMatchingProbe(key))?
-                };
 
                 let now_ms = if let Some(timestamp) = kernel_rx_timestamp {
                     ClockSample::<Realtime>::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;