1 files changed, 179 insertions, 0 deletions
diff --git a/src/swarm/rtt.rs b/src/swarm/rtt.rs
new file mode 100644
index 0000000..ed6cc73
--- /dev/null
+++ b/src/swarm/rtt.rs
@@ -0,0 +1,179 @@
+//! Round-trip times to peers, measured from traffic we send anyway: PING → PONG and
+//! HELLO → HELLO_REPLY. No probes of our own.
+//!
+//! A sample only counts when exactly one request was outstanding: with two, we can't tell
+//! which one the answer belongs to (an answer to the older one would look too fast).
+//! Samples are smoothed (exponential moving average, weight 1/4) to ride out jitter.
+//!
+//! What we use is the *effective* round-trip time: the smoothed value, raised to the age
+//! of the oldest request still unanswered (the round trip takes at least that long), so
+//! a peer that stops answering drifts away by itself. Without a recent sample (nodes
+//! move; an old value says little) it is unknown. Silence alone doesn't count: we don't
+//! ask non-neighbours anything, so their quiet is ours.
+
+use std::collections::HashMap;
+use std::time::Duration;
+
+use tokio::time::Instant;
+
+#[derive(Debug, Default)]
+pub(crate) struct RttTable {
+    pings: HashMap<String, Outstanding>,
+    hellos: HashMap<String, Outstanding>,
+    smoothed: HashMap<String, Sample>,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct Sample {
+    rtt: Duration,
+    at: Instant,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct Outstanding {
+    sent: Instant,
+    count: u32,
+}
+
+impl RttTable {
+    pub(crate) fn ping_sent(&mut self, peer: &str, now: Instant) {
+        Self::sent(&mut self.pings, peer, now);
+    }
+
+    pub(crate) fn pong_received(&mut self, peer: &str, now: Instant) {
+        if let Some(sample) = Self::answered(&mut self.pings, peer, now) {
+            self.record(peer, sample, now);
+        }
+    }
+
+    pub(crate) fn hello_sent(&mut self, peer: &str, now: Instant) {
+        Self::sent(&mut self.hellos, peer, now);
+    }
+
+    pub(crate) fn hello_reply_received(&mut self, peer: &str, now: Instant) {
+        if let Some(sample) = Self::answered(&mut self.hellos, peer, now) {
+            self.record(peer, sample, now);
+        }
+    }
+
+    /// The effective round-trip time: the smoothed value if sampled within `fresh_for`,
+    /// raised to the age of the oldest unanswered request; None without a recent sample.
+    pub(crate) fn effective(
+        &self,
+        peer: &str,
+        now: Instant,
+        fresh_for: Duration,
+    ) -> Option<Duration> {
+        let sample = self
+            .smoothed
+            .get(peer)
+            .filter(|sample| now.saturating_duration_since(sample.at) < fresh_for)?;
+        let waiting = [self.pings.get(peer), self.hellos.get(peer)]
+            .into_iter()
+            .flatten()
+            .map(|outstanding| now.saturating_duration_since(outstanding.sent))
+            .max()
+            .unwrap_or_default();
+        Some(sample.rtt.max(waiting))
+    }
+
+    /// The effective round-trip time of every peer that has one.
+    pub(crate) fn all_effective(
+        &self,
+        now: Instant,
+        fresh_for: Duration,
+    ) -> impl Iterator<Item = (&str, Duration)> {
+        self.smoothed.keys().filter_map(move |peer| {
+            self.effective(peer, now, fresh_for)
+                .map(|rtt| (peer.as_str(), rtt))
+        })
+    }
+
+    fn sent(map: &mut HashMap<String, Outstanding>, peer: &str, now: Instant) {
+        map.entry(peer.to_owned())
+            .and_modify(|outstanding| outstanding.count += 1)
+            .or_insert(Outstanding {
+                sent: now,
+                count: 1,
+            });
+    }
+
+    fn answered(
+        map: &mut HashMap<String, Outstanding>,
+        peer: &str,
+        now: Instant,
+    ) -> Option<Duration> {
+        let outstanding = map.remove(peer)?;
+        (outstanding.count == 1).then(|| now.saturating_duration_since(outstanding.sent))
+    }
+
+    fn record(&mut self, peer: &str, sample: Duration, now: Instant) {
+        let rtt = match self.smoothed.get(peer) {
+            Some(old) => (old.rtt * 3 + sample) / 4,
+            None => sample,
+        };
+        self.smoothed
+            .insert(peer.to_owned(), Sample { rtt, at: now });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const FRESH: Duration = Duration::from_secs(300);
+
+    #[test]
+    fn samples_only_unambiguous_answers_and_smooths() {
+        let mut table = RttTable::default();
+        let start = Instant::now();
+        let at = |ms| start + Duration::from_millis(ms);
+        table.ping_sent("a", start);
+        table.pong_received("a", at(100));
+        assert_eq!(
+            table.effective("a", at(100), FRESH),
+            Some(Duration::from_millis(100))
+        );
+
+        table.ping_sent("a", at(100));
+        table.pong_received("a", at(300));
+        assert_eq!(
+            table.effective("a", at(300), FRESH),
+            Some(Duration::from_millis(125))
+        );
+
+        // Two PINGs out: the PONG can't be attributed.
+        table.ping_sent("a", at(300));
+        table.ping_sent("a", at(400));
+        table.pong_received("a", at(450));
+        assert_eq!(
+            table.effective("a", at(450), FRESH),
+            Some(Duration::from_millis(125))
+        );
+
+        // An unsolicited PONG changes nothing.
+        table.pong_received("b", start);
+        assert_eq!(table.effective("b", start, FRESH), None);
+    }
+
+    #[test]
+    fn unanswered_requests_raise_it_and_old_samples_expire() {
+        let mut table = RttTable::default();
+        let start = Instant::now();
+        table.ping_sent("a", start);
+        table.pong_received("a", start + Duration::from_millis(50));
+
+        // A PING unanswered for 40 s: at least 40 s.
+        let sent = start + Duration::from_secs(60);
+        table.ping_sent("a", sent);
+        let later = sent + Duration::from_secs(40);
+        assert_eq!(
+            table.effective("a", later, FRESH),
+            Some(Duration::from_secs(40))
+        );
+
+        // No sample within the freshness window: unknown.
+        let stale = start + FRESH + Duration::from_secs(1);
+        assert_eq!(table.effective("a", stale, FRESH), None);
+    }
+}