103 lines
2.8 KiB
Rust
103 lines
2.8 KiB
Rust
use std::{collections::HashMap, net::IpAddr, time::Duration};
|
|
|
|
use flume::{Receiver, SendError, Sender};
|
|
use thiserror::Error;
|
|
|
|
use crate::measurement::{
|
|
icmp::IcmpBinding,
|
|
model::{MeasurementSource, MeasurementStream, MeasurementStreamId, ProbeSchedule},
|
|
};
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum ControllerError {
|
|
#[error("ICMP channel error")]
|
|
IcmpChannelError(#[from] SendError<Vec<IcmpBinding>>),
|
|
}
|
|
|
|
pub struct Controller {
|
|
reflector_rx: Receiver<Vec<IpAddr>>,
|
|
icmp_binding_tx: Sender<Vec<IcmpBinding>>,
|
|
|
|
source: MeasurementSource,
|
|
period: Duration,
|
|
|
|
next_stream_id: u64,
|
|
stream_ids: HashMap<(IpAddr, MeasurementSource), MeasurementStreamId>,
|
|
}
|
|
|
|
impl Controller {
|
|
pub fn new(
|
|
reflector_rx: Receiver<Vec<IpAddr>>,
|
|
icmp_binding_tx: Sender<Vec<IcmpBinding>>,
|
|
) -> Self {
|
|
Self {
|
|
reflector_rx,
|
|
icmp_binding_tx,
|
|
source: MeasurementSource::IcmpTimestamp,
|
|
period: Duration::from_millis(500),
|
|
next_stream_id: 0,
|
|
stream_ids: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub async fn run(mut self, initial_reflectors: Vec<IpAddr>) -> Result<(), ControllerError> {
|
|
self.update_reflectors(initial_reflectors).await?;
|
|
|
|
while let Ok(reflectors) = self.reflector_rx.recv_async().await {
|
|
self.update_reflectors(reflectors).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn stream_for(&mut self, peer: IpAddr) -> MeasurementStream {
|
|
let key = (peer, self.source);
|
|
|
|
let id = *self.stream_ids.entry(key).or_insert_with(|| {
|
|
let id = MeasurementStreamId(self.next_stream_id);
|
|
self.next_stream_id += 1;
|
|
id
|
|
});
|
|
|
|
MeasurementStream {
|
|
id,
|
|
peer,
|
|
source: self.source,
|
|
}
|
|
}
|
|
|
|
async fn update_reflectors(&mut self, reflectors: Vec<IpAddr>) -> Result<(), ControllerError> {
|
|
let count = reflectors.len();
|
|
|
|
let bindings = reflectors
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(index, peer)| {
|
|
let stream = self.stream_for(peer);
|
|
|
|
IcmpBinding::new(
|
|
stream,
|
|
ProbeSchedule {
|
|
period: self.period,
|
|
offset: evenly_spaced_offset(self.period, index, count),
|
|
},
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
self.icmp_binding_tx
|
|
.send_async(bindings)
|
|
.await
|
|
.map_err(|e| ControllerError::IcmpChannelError(e))
|
|
}
|
|
}
|
|
|
|
fn evenly_spaced_offset(period: Duration, index: usize, count: usize) -> Duration {
|
|
let nanos = period.as_nanos() * index as u128 / count as u128;
|
|
|
|
Duration::new(
|
|
(nanos / 1_000_000_000) as u64,
|
|
(nanos % 1_000_000_000) as u32,
|
|
)
|
|
}
|