Lochnair-Labs/sqm-autorate-rust
2
0 Fork 0
Code Issues 5 Pull Requests 2 Actions 4 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Nils Andreas Svee committed 2026-10-08 14:15:14 +07:00
1 parent afedfbeb65
commit 89c7d75947
4 files changed
+257 -194
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/main.rs
+3 -1
View File
@@ -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
+114 -187
View File
@@ -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 {
socket: IcmpSocket4,
identifier: u16,
inflight: Arc<InFlightProbes>,
bindings: HashMap<MeasurementStreamId, IcmpBinding>,
inflight: HashMap<InFlightProbeKey, InFlightProbe>,
origin: Instant,
binding_rx: Receiver<Vec<IcmpBinding>>,
observation_tx: Sender<MeasurementEvent>,
}
pub struct IcmpEngine {
sender: IcmpSender,
receiver: IcmpReceiver,
async fn wait_for_deadline(deadline: Option<Instant>) {
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<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()?;
let tx_socket = rx_socket.try_clone()?;
socket.enable_receive_metadata()?;
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: 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<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> {
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::<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) {
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::<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;