From 6a69f4574e4d8ec028c325d6e3c8afa555363c75 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 13 Oct 2026 16:01:02 +0700 Subject: [PATCH] feat(irtt): make IRTT measurement optional and refine engine behavior Conditionally compile IRTT functionality behind the `irtt` feature flag to allow building smaller binaries without IRTT dependencies. Enhance IRTT engine shutdown, configuration, and clarify ICMP protocol names. --- src/controller.rs | 27 +++++++++++----- src/main.rs | 22 +++++++++++-- src/measurement/irtt.rs | 69 ++++++++++++++++++++++++++++++++-------- src/measurement/model.rs | 2 +- src/settings/mod.rs | 34 ++++++++++---------- 5 files changed, 112 insertions(+), 42 deletions(-) diff --git a/src/controller.rs b/src/controller.rs index 188f4ef..258c414 100644 --- a/src/controller.rs +++ b/src/controller.rs @@ -8,10 +8,11 @@ use std::{ use flume::{Receiver, SendError, Sender}; use thiserror::Error; +#[cfg(feature = "irtt")] +use crate::measurement::irtt::IrttBinding; use crate::{ measurement::{ icmp::IcmpBinding, - irtt::IrttBinding, model::{MeasurementSource, MeasurementStream, MeasurementStreamId, ProbeSchedule}, }, settings::ReflectorDefinition, @@ -21,6 +22,7 @@ use crate::{ pub enum ControllerError { #[error("ICMP channel error")] IcmpChannelError(#[from] SendError>), + #[cfg(feature = "irtt")] #[error("IRTT channel error")] IrttChannelError(#[from] SendError>), } @@ -28,6 +30,7 @@ pub enum ControllerError { pub struct Controller { reflector_rx: Receiver>, icmp_binding_tx: Sender>, + #[cfg(feature = "irtt")] irtt_binding_tx: Sender>, reflectors: HashMap, active_streams: Arc>>, @@ -42,7 +45,7 @@ impl Controller { pub fn new( reflector_rx: Receiver>, icmp_binding_tx: Sender>, - irtt_binding_tx: Sender>, + #[cfg(feature = "irtt")] irtt_binding_tx: Sender>, reflectors: HashMap, active_streams: Arc>>, endpoint_streams: Arc>>, @@ -51,6 +54,7 @@ impl Controller { Self { reflector_rx, icmp_binding_tx, + #[cfg(feature = "irtt")] irtt_binding_tx, reflectors, active_streams, @@ -90,18 +94,21 @@ impl Controller { ) -> Result<(), ControllerError> { let count = endpoints.len(); let mut icmp = Vec::new(); + #[cfg(feature = "irtt")] let mut irtt = Vec::new(); let mut streams = Vec::with_capacity(count); let mut endpoint_streams = HashMap::new(); for (index, endpoint) in endpoints.into_iter().enumerate() { - let Some((port, hmac)) = self.reflectors.get(&endpoint).map(|definition| { - ( - definition.irtt_port.unwrap_or(2112), - definition.irtt_hmac.clone(), - ) - }) else { + let Some(_definition) = self.reflectors.get(&endpoint) else { continue; }; + #[cfg(feature = "irtt")] + let (port, hmac) = { + ( + _definition.irtt_port.unwrap_or(2112), + _definition.irtt_hmac.clone(), + ) + }; let stream = self.stream_for(endpoint); let schedule = ProbeSchedule { period: self.period, @@ -111,9 +118,12 @@ impl Controller { MeasurementSource::IcmpEcho | MeasurementSource::IcmpTimestamp => { icmp.push(IcmpBinding::new(stream, schedule)) } + #[cfg(feature = "irtt")] MeasurementSource::Irtt => { irtt.push(IrttBinding::new(stream, schedule, port, hmac)) } + #[cfg(not(feature = "irtt"))] + MeasurementSource::Irtt => unreachable!("IRTT feature is checked before startup"), _ => unreachable!("unsupported reflector measurement source"), } streams.push(stream); @@ -128,6 +138,7 @@ impl Controller { .write() .expect("endpoint stream lock poisoned") = endpoint_streams; self.icmp_binding_tx.send_async(icmp).await?; + #[cfg(feature = "irtt")] self.irtt_binding_tx.send_async(irtt).await?; Ok(()) } diff --git a/src/main.rs b/src/main.rs index 5977206..cef3c66 100644 --- a/src/main.rs +++ b/src/main.rs @@ -22,6 +22,7 @@ mod util; use crate::baseliner::Baseliner; use crate::controller::Controller; use crate::measurement::icmp::IcmpEngine; +#[cfg(feature = "irtt")] use crate::measurement::irtt::IrttEngine; use crate::measurement::model::{MeasurementSource, MeasurementStream}; use crate::metrics::{Metric, Metrics, MetricsSender}; @@ -279,6 +280,12 @@ fn wait_for_exit(error_rx: &flume::Receiver) -> anyhow::Result<() } fn run(settings: &Settings) -> anyhow::Result<()> { + #[cfg(not(feature = "irtt"))] + if settings.advanced_settings.measurement_type == ReflectorProtocol::Irtt { + anyhow::bail!( + "IRTT measurement selected, but this binary was built without the `irtt` feature" + ); + } let rt = tokio::runtime::Runtime::new()?; if settings.advanced_settings.dry_run { @@ -299,6 +306,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { 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(); + #[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(); @@ -318,6 +326,7 @@ fn run(settings: &Settings) -> anyhow::Result<()> { let controller = Controller::new( reflector_rx, icmp_binding_tx, + #[cfg(feature = "irtt")] irtt_binding_tx, reflectors, Arc::clone(&active_streams), @@ -333,13 +342,20 @@ fn run(settings: &Settings) -> anyhow::Result<()> { icmp_binding_rx, )? }; - let (_irtt_task, irtt_engine) = { + #[cfg(feature = "irtt")] + let (irtt_task, irtt_engine) = { let _guard = rt.enter(); - IrttEngine::new(start_time, measurement_tx, irtt_binding_rx)? + IrttEngine::new( + start_time, + Duration::from_secs_f64(settings.advanced_settings.tick_interval), + measurement_tx, + irtt_binding_rx, + )? }; spawn_task(&rt, &error_tx, icmp_engine.run()); - spawn_task(&rt, &error_tx, irtt_engine.run()); + #[cfg(feature = "irtt")] + spawn_task(&rt, &error_tx, irtt_engine.run(irtt_task)); let dropped = Arc::new(AtomicU32::new(0)); diff --git a/src/measurement/irtt.rs b/src/measurement/irtt.rs index c100ca6..665a33d 100644 --- a/src/measurement/irtt.rs +++ b/src/measurement/irtt.rs @@ -2,7 +2,7 @@ use std::{collections::HashMap, net::SocketAddr, time::Instant}; use flume::{Receiver, Sender}; use irtt_client::{ - ClientConfig, ClientEvent, + ClientConfig, ClientEvent, NegotiationPolicy, OpenPolicy, managed::{ ManagedClient, ManagedClientConfig, ManagedClientHandle, ManagedClientTask, ManagedCommandAcknowledgement, ManagedCommandApplyError, ManagedCommandError, @@ -15,6 +15,7 @@ use log::{debug, trace, warn}; use thiserror::Error; use tokio::sync::broadcast::error::RecvError; +use crate::SHUTDOWN; use crate::measurement::model::{ MeasurementEvent, MeasurementLoss, MeasurementObservation, MeasurementStream, MeasurementStreamId, OneWayClock, OneWayLatency, ProbeSchedule, SignedDuration, @@ -52,6 +53,9 @@ pub enum IrttError { pub struct IrttBinding { stream: MeasurementStream, + /// Retained for future hybrid-mode scheduling. IRTT currently uses the + /// global interval configured on its managed client. + #[allow(dead_code)] schedule: ProbeSchedule, port: u16, hmac: Option, @@ -87,20 +91,11 @@ pub struct IrttEngine { impl IrttEngine { pub fn new( origin: Instant, + interval: std::time::Duration, observation_tx: Sender, binding_rx: Receiver>, ) -> Result<(ManagedClientTask, Self), IrttError> { - let config = ManagedClientConfig { - client: ClientConfig { - request: irtt_client::SessionRequest { - duration: None, - ..Default::default() - }, - ..Default::default() - }, - completion: ManagedCompletionPolicy::ExplicitStop, - ..Default::default() - }; + let config = managed_client_config(interval); let (task, handle) = ManagedClient::task(config, Vec::new())?; Ok(( @@ -142,11 +137,19 @@ impl IrttEngine { Ok(()) } - pub async fn run(mut self) -> Result<(), IrttError> { + pub async fn run(mut self, task: ManagedClientTask) -> Result<(), IrttError> { + let mut task = Box::pin(task); let mut events = self.client.subscribe()?; loop { tokio::select! { + outcome = task.as_mut() => { + return match outcome.end_reason { + ManagedEndReason::StopRequested => Ok(()), + end_reason => Err(IrttError::ManagedFailure(end_reason)), + }; + } + bindings = self.binding_rx.recv_async() => { let bindings = bindings .map_err(|_| IrttError::BindingChannelClosed)?; @@ -338,11 +341,51 @@ impl IrttEngine { } } } + + _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => { + if SHUTDOWN.load(std::sync::atomic::Ordering::Relaxed) { + let stop = self.client.stop(); + tokio::pin!(stop); + tokio::select! { + _ = &mut stop => {}, + outcome = task.as_mut() => { + return match outcome.end_reason { + ManagedEndReason::StopRequested => Ok(()), + end_reason => Err(IrttError::ManagedFailure(end_reason)), + }; + } + } + let outcome = task.await; + return match outcome.end_reason { + ManagedEndReason::StopRequested => Ok(()), + end_reason => Err(IrttError::ManagedFailure(end_reason)), + }; + } + } } } } } +fn managed_client_config(interval: std::time::Duration) -> ManagedClientConfig { + ManagedClientConfig { + client: ClientConfig { + request: irtt_client::SessionRequest { + duration: None, + interval, + ..Default::default() + }, + open: OpenPolicy { + negotiation: NegotiationPolicy::Strict, + ..Default::default() + }, + ..Default::default() + }, + completion: ManagedCompletionPolicy::ExplicitStop, + ..Default::default() + } +} + fn map_one_way(sample: irtt_client::OneWayDelaySample) -> Option { let uplink = i64::try_from(sample.client_to_server?.as_nanos()).ok()?; diff --git a/src/measurement/model.rs b/src/measurement/model.rs index 12b7bf5..aab7fb5 100644 --- a/src/measurement/model.rs +++ b/src/measurement/model.rs @@ -20,7 +20,7 @@ pub enum MeasurementSource { impl Display for MeasurementSource { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - MeasurementSource::IcmpEcho => write!(f, "icmp"), + MeasurementSource::IcmpEcho => write!(f, "icmp-echo"), MeasurementSource::IcmpTimestamp => write!(f, "icmp-timestamps"), MeasurementSource::Irtt => write!(f, "irtt"), MeasurementSource::TcpSeqAck => write!(f, "tcp-seqack"), diff --git a/src/settings/mod.rs b/src/settings/mod.rs index 514a9c4..a59aabd 100644 --- a/src/settings/mod.rs +++ b/src/settings/mod.rs @@ -89,7 +89,7 @@ pub enum ReflectorProtocol { impl Display for ReflectorProtocol { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - ReflectorProtocol::IcmpEcho => write!(f, "icmp"), + ReflectorProtocol::IcmpEcho => write!(f, "icmp-echo"), ReflectorProtocol::IcmpTimestamps => write!(f, "icmp-timestamps"), ReflectorProtocol::Irtt => write!(f, "irtt"), } @@ -101,8 +101,8 @@ impl FromStr for ReflectorProtocol { fn from_str(s: &str) -> Result { match s.to_lowercase().as_str() { - "icmp" => Ok(ReflectorProtocol::IcmpEcho), - "icmp-timestamps" | "icmp_timestamps" => Ok(ReflectorProtocol::IcmpTimestamps), + "icmp-echo" => Ok(ReflectorProtocol::IcmpEcho), + "icmp-timestamps" => Ok(ReflectorProtocol::IcmpTimestamps), "irtt" => Ok(ReflectorProtocol::Irtt), _ => Err(ConfigError::Message(format!( "invalid measurement type: {s}" @@ -137,15 +137,15 @@ pub(crate) struct ReflectorDefinition { pub(crate) country: Option, pub(crate) protocols: Vec, pub(crate) irtt_port: Option, + #[allow(dead_code)] pub(crate) irtt_hmac: Option, } #[derive(Debug, Deserialize)] struct ReflectorRow { - #[serde(alias = "reflector_ip")] host: String, country: Option, - protocols: Option, + protocols: String, irtt_port: Option, irtt_hmac: Option, } @@ -246,18 +246,18 @@ impl ReflectorSettings { .map_err(|error| ConfigError::Message(error.to_string()))?; for row in reader.deserialize::() { let row = row.map_err(|error| ConfigError::Message(error.to_string()))?; - let protocols = match row.protocols { - Some(protocols) => protocols - .split(['|', ';']) - .map(str::trim) - .filter(|protocol| !protocol.is_empty()) - .map(ReflectorProtocol::from_str) - .collect::, _>>()?, - None => vec![ - ReflectorProtocol::IcmpEcho, - ReflectorProtocol::IcmpTimestamps, - ], - }; + let protocols = row + .protocols + .split(['|', ';']) + .map(str::trim) + .filter(|protocol| !protocol.is_empty()) + .map(ReflectorProtocol::from_str) + .collect::, _>>()?; + if protocols.is_empty() { + return Err(ConfigError::Message(format!( + "reflector row in {path} has no protocols" + ))); + } self.entries.push(ReflectorDefinition { host: row.host, country: row.country,