record RTT from keepalive messages

1 files changed, 14 insertions, 6 deletions

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;
             }