use crate::{bimap::BiHashMap, timeout::MaybeTimeout};
use rustc_hash::FxHashMap;
use std::{net::Ipv4Addr, time::Duration};
#[derive(Debug)]
pub struct NetworkAddressTable {
addr_map: BiHashMap<u32, u32>,
timeouts: FxHashMap<(u32, u32), MaybeTimeout>,
}
impl NetworkAddressTable {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[profiling::function]
pub fn prune(&mut self) {
log::trace!("Pruning old network address mappings");
let now = std::time::Instant::now();
self.timeouts.retain(|(left, right), timeout| {
match timeout {
MaybeTimeout::Never => true,
MaybeTimeout::After { duration, start } => {
let should_retain = now.duration_since(*start) < *duration;
if !should_retain {
log::trace!(
"Mapping {:?} -> {:?} has timed out and will be removed",
left,
right
);
self.addr_map.remove(left, right);
}
should_retain
}
}
});
}
#[profiling::function]
pub fn insert_indefinite(&mut self, left: Ipv4Addr, right: Ipv4Addr) {
self.prune();
let (left, right) = (left.into(), right.into());
self.addr_map.insert(left, right);
self.timeouts.insert((left, right), MaybeTimeout::Never);
}
#[profiling::function]
pub fn insert(&mut self, left: Ipv4Addr, right: Ipv4Addr, duration: Duration) {
self.prune();
let (left, right) = (left.into(), right.into());
self.addr_map.insert(left, right);
self.timeouts.insert(
(left, right),
MaybeTimeout::After {
duration,
start: std::time::Instant::now(),
},
);
}
#[must_use]
#[profiling::function]
pub fn get_right(&self, left: &Ipv4Addr) -> Option<Ipv4Addr> {
self.addr_map
.get_right(&(*left).into())
.map(|addr| (*addr).into())
}
#[must_use]
#[profiling::function]
pub fn get_left(&self, right: &Ipv4Addr) -> Option<Ipv4Addr> {
self.addr_map
.get_left(&(*right).into())
.map(|addr| (*addr).into())
}
}
impl Default for NetworkAddressTable {
fn default() -> Self {
Self {
addr_map: BiHashMap::new(),
timeouts: FxHashMap::default(),
}
}
}