1 files changed, 112 insertions, 0 deletions
diff --git a/src/swarm/rtt.rs b/src/swarm/rtt.rs
new file mode 100644
index 0000000..04d0722
--- /dev/null
+++ b/src/swarm/rtt.rs
@@ -0,0 +1,112 @@
+//! 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.
+
+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, Duration>,
+}
+
+#[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);
+        }
+    }
+
+    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);
+        }
+    }
+
+    /// The smoothed round-trip time, once measured.
+    pub(crate) fn get(&self, peer: &str) -> Option<Duration> {
+        self.smoothed.get(peer).copied()
+    }
+
+    /// All measured peers.
+    pub(crate) fn all(&self) -> impl Iterator<Item = (&str, Duration)> {
+        self.smoothed
+            .iter()
+            .map(|(peer, 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) {
+        let smoothed = match self.smoothed.get(peer) {
+            Some(old) => (*old * 3 + sample) / 4,
+            None => sample,
+        };
+        self.smoothed.insert(peer.to_owned(), smoothed);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn samples_only_unambiguous_answers_and_smooths() {
+        let mut table = RttTable::default();
+        let start = Instant::now();
+        table.ping_sent("a", start);
+        table.pong_received("a", start + Duration::from_millis(100));
+        assert_eq!(table.get("a"), Some(Duration::from_millis(100)));
+
+        table.ping_sent("a", start);
+        table.pong_received("a", start + Duration::from_millis(200));
+        assert_eq!(table.get("a"), Some(Duration::from_millis(125)));
+
+        // Two PINGs out: the PONG can't be attributed.
+        table.ping_sent("a", start);
+        table.ping_sent("a", start + Duration::from_secs(1));
+        table.pong_received("a", start + Duration::from_secs(1));
+        assert_eq!(table.get("a"), Some(Duration::from_millis(125)));
+
+        // An unsolicited PONG changes nothing.
+        table.pong_received("b", start);
+        assert_eq!(table.get("b"), None);
+    }
+}