diff --git a/src/ping.rs b/src/ping.rs index ba60250..46849c7 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -6,17 +6,18 @@ //! the host answers pings — a different thing, often throttled or blocked — //! and would need raw sockets or an external `ping` binary. //! -//! Flow: each round samples one random host from every segment, probes them -//! concurrently, and keeps the ones that pass the filters — repeating until -//! `count` valid IPs are gathered (or the `max_probe` budget is hit), then -//! returns the lowest-latency ones. +//! Flow: probe random hosts across all segments concurrently, keeping the ones +//! that pass the filters, until `count` valid IPs are gathered (or the +//! `max_probe` budget is hit). Each segment is sampled in proportion to how +//! often it has yielded valid IPs, so persistently failing segments fade out +//! and the budget flows to productive ones. Returns the lowest-latency IPs. use std::collections::HashSet; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::{Duration, Instant}; use anyhow::Result; -use futures::stream::{self, StreamExt}; +use futures::stream::{FuturesUnordered, StreamExt}; use ipnet::IpNet; use rand::Rng; use tokio::net::TcpStream; @@ -122,49 +123,87 @@ pub async fn run( let mut collected: Vec = Vec::new(); let mut probed = 0usize; - let (port, times, to, concurrency) = (cfg.port, cfg.times, cfg.timeout, cfg.concurrency.max(1)); - // Lazily generate fresh random hosts, round-robin across all segments. A - // continuous stream (vs discrete rounds) keeps `concurrency` probes always - // in flight, so a slow unreachable IP only blocks its own slot. + // Per-segment adaptive sampling. CF segments are wildly uneven — many are + // unroutable from a given vantage point. Each segment is drawn in + // proportion to its Laplace-smoothed valid-hit rate, so a segment that + // keeps failing is sampled less and the budget flows to productive ones. + let mut attempts = vec![0u32; nets.len()]; + let mut hits = vec![0u32; nets.len()]; let mut seen: HashSet = HashSet::new(); let mut rng = rand::thread_rng(); - let mut idx = 0usize; - let candidates = std::iter::from_fn(move || { - let max_tries = nets.len().saturating_mul(50).max(1); - for _ in 0..max_tries { - let net = &nets[idx % nets.len()]; - idx = idx.wrapping_add(1); - let ip = random_host(net, &mut rng); - if seen.insert(ip) { - return Some(ip); + + // Keep `concurrency` probes in flight, refilling one as each finishes so the + // weights learned so far steer every new draw. + let mut inflight = FuturesUnordered::new(); + for _ in 0..concurrency { + match draw_candidate(&nets, &attempts, &hits, &mut seen, &mut rng) { + Some((seg, ip)) => { + attempts[seg] += 1; + inflight.push(probe(seg, ip, port, times, to)); } + None => break, } - None // every segment effectively exhausted - }); + } - let mut probes = stream::iter( - candidates.map(|ip| async move { (ip, probe_ip(ip, port, times, to).await) }), - ) - .buffer_unordered(concurrency); - - while let Some((ip, outcome)) = probes.next().await { + while let Some((seg, ip, outcome)) = inflight.next().await { probed += 1; let status = cfg.classify(&outcome); let latency = (outcome.received > 0).then_some(outcome.avg_latency); if matches!(status, ProbeStatus::Ok) { + hits[seg] += 1; collected.push(PingResult { ip, latency: outcome.avg_latency }); } on_progress(Progress::Probing { ip, status, latency, valid: collected.len(), probed }); if collected.len() >= cfg.count || probed >= cfg.max_probe { break; } + // Refill the freed slot, biased by the updated weights. + if let Some((seg, ip)) = draw_candidate(&nets, &attempts, &hits, &mut seen, &mut rng) { + attempts[seg] += 1; + inflight.push(probe(seg, ip, port, times, to)); + } } Ok(finalize(collected, cfg.count)) } +/// Choose a segment index with probability proportional to its Laplace-smoothed +/// valid-hit rate `(hits+1)/(attempts+2)`. Untried segments start at 0.5 and +/// fade toward 0 the longer they go without yielding a valid IP. +fn weighted_segment(attempts: &[u32], hits: &[u32], rng: &mut impl Rng) -> usize { + let weight = |i: usize| (hits[i] as f64 + 1.0) / (attempts[i] as f64 + 2.0); + let total: f64 = (0..attempts.len()).map(weight).sum(); + let mut pick = rng.gen_range(0.0..total); + for i in 0..attempts.len() { + pick -= weight(i); + if pick < 0.0 { + return i; + } + } + attempts.len() - 1 +} + +/// Draw a fresh, never-probed host, choosing the segment by weight. Returns the +/// segment index and IP, or None if no new host turned up (segments exhausted). +fn draw_candidate( + nets: &[IpNet], + attempts: &[u32], + hits: &[u32], + seen: &mut HashSet, + rng: &mut impl Rng, +) -> Option<(usize, IpAddr)> { + for _ in 0..64 { + let seg = weighted_segment(attempts, hits, rng); + let ip = random_host(&nets[seg], rng); + if seen.insert(ip) { + return Some((seg, ip)); + } + } + None +} + /// Sort by latency ascending and keep the best `count`. fn finalize(mut results: Vec, count: usize) -> Vec { results.sort_by_key(|r| r.latency); @@ -172,6 +211,19 @@ fn finalize(mut results: Vec, count: usize) -> Vec { results } +/// Probe one IP and tag the result with its segment index, so the scheduler can +/// credit the right segment. One named future type keeps `FuturesUnordered` +/// happy (distinct async blocks would be distinct types). +async fn probe( + seg: usize, + ip: IpAddr, + port: u16, + times: usize, + to: Duration, +) -> (usize, IpAddr, ProbeOutcome) { + (seg, ip, probe_ip(ip, port, times, to).await) +} + /// Probe one IP `times` times, returning how many succeeded and the average /// latency over the successful probes. async fn probe_ip(ip: IpAddr, port: u16, times: usize, to: Duration) -> ProbeOutcome {