record RTT from keepalive messages

7 files changed, 148 insertions, 20 deletions

diff --git a/src/api/peers.rs b/src/api/peers.rs
index d381395..3893e0f 100644
--- a/src/api/peers.rs
+++ b/src/api/peers.rs
@@ -62,6 +62,9 @@ pub(crate) struct PeerStatus {
     heads: Option<usize>,
     /// We keep it alive and sync with it on our own; other peers we only answer.
     neighbour: bool,
+    /// Smoothed round-trip time in milliseconds, from PING/PONG and HELLO/HELLO_REPLY;
+    /// absent until measured.
+    rtt_ms: Option<f64>,
 }
 
 #[derive(Debug, Serialize, ToSchema)]
@@ -109,6 +112,10 @@ pub(crate) async fn list_peers(
 
     for (public_key, peer_state_mutex) in snapshot {
         let neighbour = crate::swarm::topology::is_neighbour(&state, &public_key);
+        let rtt_ms = state
+            .lock_rtt()
+            .get(&public_key)
+            .map(|rtt| rtt.as_secs_f64() * 1000.0);
         let ps = peer_state_mutex.lock().await;
 
         let (status, dormant_for_secs) = match ps.liveness {
@@ -133,6 +140,7 @@ pub(crate) async fn list_peers(
             same_epoch: ps.epoch.map(|epoch| epoch == local_epoch),
             heads: ps.reported_heads,
             neighbour,
+            rtt_ms,
         });
     }
 
diff --git a/src/main.rs b/src/main.rs
index 527a730..f5f21fe 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -194,6 +194,7 @@ async fn main() -> Result<()> {
         neighbours_k: args.neighbours.max(2),
         neighbours: Default::default(),
         pulling: Default::default(),
+        rtt: Default::default(),
         stale: Default::default(),
         catch_up_source: Default::default(),
     };
diff --git a/src/state.rs b/src/state.rs
index 10471df..2bc9264 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -536,6 +536,8 @@ pub(crate) struct AppState {
     /// Target number of neighbours (`--neighbours`); up to this many + 1 members the
     /// swarm is a full mesh.
     pub(crate) neighbours_k: usize,
+    /// Round-trip times to peers (`swarm::rtt`).
+    pub(crate) rtt: Arc<std::sync::Mutex<crate::swarm::rtt::RttTable>>,
     /// The neighbour we are pulling changes from, and since when (`swarm::relay`).
     pub(crate) pulling: Arc<std::sync::Mutex<Option<(String, Instant)>>>,
     /// Our neighbours, once known (`swarm::topology`).
@@ -566,6 +568,12 @@ impl AppState {
         self.transport.local_id()
     }
 
+    pub(crate) fn lock_rtt(&self) -> std::sync::MutexGuard<'_, crate::swarm::rtt::RttTable> {
+        self.rtt
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
     /// The local epoch. Code that labels or applies sync messages must use the epoch
     /// of the locked `Replica` instead, so label and document can't disagree.
     pub(crate) fn local_epoch(&self) -> Epoch {
@@ -814,12 +822,7 @@ impl AppState {
         }
 
         for id in announce {
-            let keepalive = crate::swarm::peer::local_keepalive(self, id).await;
-            if let Err(error) = self
-                .transport
-                .send_message(id, &transport::ParsedMessage::Ping(keepalive))
-                .await
-            {
+            if let Err(error) = crate::swarm::peer::send_ping(self, id).await {
                 debug!(peer = %id, error = %format!("{error:#}"), "Epoch announcement failed");
             }
         }
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);
+    }
+}