use std::time::Duration;
use tokio::process::Command;
pub const RESOLVE: Duration = Duration::from_secs(5);
pub const QUICK: Duration = Duration::from_secs(5);
pub const SLOW: Duration = Duration::from_secs(10);
pub const TRACE: Duration = Duration::from_secs(60);
pub fn ping_budget(count: u32) -> Duration {
Duration::from_secs(2 * count as u64 + 4)
}
pub async fn retry_probe<T, F, Fut>(attempts: u32, delay: Duration, mut op: F) -> Option<T>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Option<T>>,
{
for attempt in 0..attempts {
if let Some(v) = op().await {
return Some(v);
}
if attempt + 1 < attempts {
tokio::time::sleep(delay).await;
}
}
None
}
pub async fn harvest_or<T>(handle: tokio::task::JoinHandle<T>, fallback: T) -> T {
handle.abort();
match handle.await {
Ok(v) => v,
Err(_) => fallback,
}
}
pub async fn run_with_timeout(mut cmd: Command, dur: Duration) -> Option<std::process::Output> {
match tokio::time::timeout(dur, cmd.output()).await {
Ok(Ok(output)) => Some(output),
Ok(Err(_)) => None,
Err(_) => None,
}
}
pub async fn lookup_host_timeout(addr: String, dur: Duration) -> Option<Vec<std::net::SocketAddr>> {
match tokio::time::timeout(dur, tokio::net::lookup_host(addr)).await {
Ok(Ok(addrs)) => Some(addrs.collect()),
Ok(Err(_)) => None,
Err(_) => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn fast_command_returns_some() {
#[cfg(windows)]
let mut cmd = {
let mut c = Command::new("cmd");
c.args(["/C", "exit", "0"]);
c
};
#[cfg(unix)]
let mut cmd = Command::new("true");
let _ = &mut cmd;
let out = run_with_timeout(cmd, QUICK).await;
assert!(out.is_some(), "fast command should finish within budget");
}
#[tokio::test]
async fn slow_command_times_out_to_none() {
#[cfg(windows)]
let cmd = {
let mut c = Command::new("cmd");
c.args(["/C", "ping", "-n", "3", "127.0.0.1"]);
c
};
#[cfg(unix)]
let cmd = {
let mut c = Command::new("sleep");
c.arg("2");
c
};
let out = run_with_timeout(cmd, Duration::from_millis(1)).await;
assert!(out.is_none(), "command exceeding budget should yield None");
}
#[tokio::test]
async fn missing_binary_returns_none() {
let cmd = Command::new("nd300-definitely-not-a-real-binary-xyz");
let out = run_with_timeout(cmd, QUICK).await;
assert!(out.is_none(), "missing binary should yield None");
}
#[tokio::test]
async fn localhost_resolves_to_some() {
let addrs = lookup_host_timeout("localhost:80".to_string(), RESOLVE).await;
assert!(
addrs.is_some_and(|v| !v.is_empty()),
"localhost:80 should resolve to at least one address"
);
}
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
#[tokio::test]
async fn retry_probe_first_success_calls_once() {
let calls = Arc::new(AtomicU32::new(0));
let c = calls.clone();
let out = retry_probe(3, Duration::from_millis(1), move || {
let c = c.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
Some(42)
}
})
.await;
assert_eq!(out, Some(42));
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn retry_probe_retries_then_succeeds() {
let calls = Arc::new(AtomicU32::new(0));
let c = calls.clone();
let out = retry_probe(3, Duration::from_millis(1), move || {
let c = c.clone();
async move {
let n = c.fetch_add(1, Ordering::SeqCst);
if n == 0 {
None
} else {
Some("ok")
}
}
})
.await;
assert_eq!(out, Some("ok"));
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn retry_probe_exhausts_to_none() {
let calls = Arc::new(AtomicU32::new(0));
let c = calls.clone();
let out: Option<u8> = retry_probe(2, Duration::from_millis(1), move || {
let c = c.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
None
}
})
.await;
assert_eq!(out, None);
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[test]
fn ping_budget_scales_with_count() {
assert_eq!(ping_budget(1), Duration::from_secs(6));
assert_eq!(ping_budget(6), Duration::from_secs(16));
assert_eq!(ping_budget(30), Duration::from_secs(64));
}
#[tokio::test]
async fn harvest_or_returns_finished_value() {
let handle = tokio::spawn(async { 42 });
tokio::task::yield_now().await;
let v = harvest_or(handle, 0).await;
assert_eq!(v, 42);
}
#[tokio::test]
async fn harvest_or_falls_back_for_pending() {
let handle = tokio::spawn(std::future::pending::<u32>());
let v = harvest_or(handle, 7).await;
assert_eq!(v, 7);
}
}