record RTT from keepalive messages

4 files changed, 130 insertions, 14 deletions

diff --git a/src/swarm/mod.rs b/src/swarm/mod.rs
index 5616fe8..6de7bf3 100644
--- a/src/swarm/mod.rs
+++ b/src/swarm/mod.rs
@@ -7,6 +7,7 @@ pub(crate) mod members;
 pub(crate) mod pair;
 pub(crate) mod peer;
 pub(crate) mod relay;
+pub(crate) mod rtt;
 pub(crate) mod source;
 pub(crate) mod sync;
 pub(crate) mod topology;
diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 2fd248e..55d0022 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -96,6 +96,16 @@ async fn apply_forget_records(state: &AppState) {
 
 // --- keepalive and self-healing ----------------------------------------------
 
+/// Sends `peer` a PING with our keepalive, noting the time for its round trip.
+pub(crate) async fn send_ping(state: &AppState, peer: &str) -> Result<()> {
+    let keepalive = local_keepalive(state, peer).await;
+    state.lock_rtt().ping_sent(peer, Instant::now());
+    state
+        .transport
+        .send_message(peer, &ParsedMessage::Ping(keepalive))
+        .await
+}
+
 /// Our keepalive for `peer`.
 pub(crate) async fn local_keepalive(state: &AppState, peer: &str) -> transport::Keepalive {
     // AutoCommit::get_heads takes &mut self (it closes the pending transaction).
@@ -266,6 +276,7 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
         .as_ref()
         .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
         .unwrap_or_default();
+    state.lock_rtt().hello_sent(peer, Instant::now());
     let result = state
         .transport
         .send_message(
@@ -460,12 +471,7 @@ pub(crate) async fn run_peer_worker(state: AppState) {
             let probe_due =
                 !ping_due && crate::swarm::source::probe_due(&state, peer, Instant::now());
             if ping_due || probe_due {
-                let local = local_keepalive(&state, peer).await;
-                if let Err(error) = state
-                    .transport
-                    .send_message(peer, &ParsedMessage::Ping(local))
-                    .await
-                {
+                if let Err(error) = send_ping(&state, peer).await {
                     debug!(peer = %peer, error = %format!("{error:#}"), "Keepalive PING failed");
                 }
             }
@@ -761,6 +767,7 @@ pub(crate) async fn handle_message(
                 .send(SyncEvent::PeerValidated(peer.clone()));
         }
         transport::ParsedMessage::HelloReply(reply) => {
+            state.lock_rtt().hello_reply_received(&peer, Instant::now());
             establish_peer(state, &peer, reply.swarm_id.clone(), false, reply.ready).await?;
 
             discover_peers(state, &peer, &reply.peers).await;
@@ -796,6 +803,7 @@ pub(crate) async fn handle_message(
             }
         }
         transport::ParsedMessage::Pong(remote) => {
+            state.lock_rtt().pong_received(&peer, Instant::now());
             if reactivated || peer_swarm(state, &peer).await.is_none() {
                 ensure_handshake(state, &peer).await;
             }
diff --git a/src/swarm/relay.rs b/src/swarm/relay.rs
index a0fa6e1..454a46a 100644
--- a/src/swarm/relay.rs
+++ b/src/swarm/relay.rs
@@ -18,10 +18,10 @@ use tokio::time::Instant;
 use tracing::debug;
 
 use crate::state::{AppState, Standing};
-use crate::swarm::peer::local_keepalive;
+use crate::swarm::peer::send_ping;
 use crate::swarm::sync::SyncEvent;
 use crate::swarm::topology;
-use crate::transport::{self, HeadsSummary, ParsedMessage};
+use crate::transport::{self, HeadsSummary};
 
 /// A pull that brought nothing for this long no longer blocks pulling from another peer.
 const PULL_TIMEOUT: Duration = Duration::from_secs(15);
@@ -41,12 +41,7 @@ pub(crate) async fn announce(state: &AppState, from: &str) {
         .cloned()
         .collect();
     for peer in neighbours {
-        let keepalive = local_keepalive(state, &peer).await;
-        if let Err(error) = state
-            .transport
-            .send_message(&peer, &ParsedMessage::Ping(keepalive))
-            .await
-        {
+        if let Err(error) = send_ping(state, &peer).await {
             debug!(peer = %peer, error = %format!("{error:#}"), "Announcing our heads failed");
         }
     }
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);
+    }
+}