feat(handshake): add ping latency tracking to MetaAddr

Author: syszery

zebra-network/src/address_book/tests/vectors.rs

@@ -1,6 +1,6 @@
 //! Fixed test vectors for the address book.
 
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 use chrono::Utc;
 use tracing::Span;
@@ -128,7 +128,9 @@ fn address_book_peer_order() {
 fn reconnection_peers_skips_recently_updated_ip() {
     // tests that reconnection_peers() skips addresses where there's a connection at that IP with a recent:
     // - `last_response`
-    test_reconnection_peers_skips_recently_updated_ip(true, MetaAddr::new_responded);
+    test_reconnection_peers_skips_recently_updated_ip(true, |addr| {
+        MetaAddr::new_responded(addr, Duration::ZERO)
+    });
 
     // tests that reconnection_peers() *does not* skip addresses where there's a connection at that IP with a recent:
     // - `last_attempt`

zebra-network/src/address_book_peers/mock.rs

@@ -2,31 +2,35 @@
 
 use std::{
     sync::{Arc, Mutex},
-    time::Instant,
+    time::{Duration, Instant},
 };
 
 use crate::{meta_addr::MetaAddr, AddressBookPeers, PeerSocketAddr};
 
+const DEFAULT_RTT: Duration = Duration::from_millis(100);
+
 /// A mock [`AddressBookPeers`] implementation that's always empty.
 #[derive(Debug, Default, Clone)]
 pub struct MockAddressBookPeers {
     /// Return value for mock `recently_live_peers` method.
     recently_live_peers: Vec<MetaAddr>,
+    default_rtt: Duration,
 }
 
 impl MockAddressBookPeers {
     /// Creates a new [`MockAddressBookPeers`]
     pub fn new(recently_live_peers: Vec<MetaAddr>) -> Self {
         Self {
             recently_live_peers,
+            default_rtt: DEFAULT_RTT,
         }
     }
 
     /// Adds a peer to the mock address book.
     pub fn add_peer(&mut self, peer: PeerSocketAddr) -> bool {
         // The real add peer will use `MetaAddr::new_initial_peer` but we just want to get a `MetaAddr` for the mock.
         self.recently_live_peers.push(
-            MetaAddr::new_responded(peer).into_new_meta_addr(
+            MetaAddr::new_responded(peer, self.default_rtt).into_new_meta_addr(
                 Instant::now(),
                 chrono::Utc::now()
                     .try_into()

zebra-network/src/meta_addr.rs

@@ -2,7 +2,7 @@
 
 use std::{
     cmp::{max, Ordering},
-    time::Instant,
+    time::{Duration, Instant},
 };
 
 use chrono::Utc;
@@ -189,6 +189,12 @@ pub struct MetaAddr {
     /// See the [`MetaAddr::last_seen`] method for details.
     last_response: Option<DateTime32>,
 
+    /// The last measured round-trip time (RTT) for this peer, if available.
+    ///
+    /// This value is updated when the peer responds to a ping (Pong).
+    #[allow(dead_code)]
+    pub(crate) rtt: Option<Duration>,
+
     /// The last time we tried to open an outbound connection to this peer.
     ///
     /// See the [`MetaAddr::last_attempt`] method for details.
@@ -278,6 +284,7 @@ pub enum MetaAddrChange {
             proptest(strategy = "canonical_peer_addr_strategy()")
         )]
         addr: PeerSocketAddr,
+        rtt: Option<Duration>,
     },
 
     /// Updates an existing `MetaAddr` when a peer fails.
@@ -320,6 +327,7 @@ impl MetaAddr {
             services: Some(untrusted_services),
             untrusted_last_seen: Some(untrusted_last_seen),
             last_response: None,
+            rtt: None,
             last_attempt: None,
             last_failure: None,
             last_connection_state: NeverAttemptedGossiped,
@@ -380,9 +388,10 @@ impl MetaAddr {
     /// - malicious peers could interfere with other peers' [`AddressBook`](crate::AddressBook)
     ///   state, or
     /// - Zebra could advertise unreachable addresses to its own peers.
-    pub fn new_responded(addr: PeerSocketAddr) -> MetaAddrChange {
+    pub fn new_responded(addr: PeerSocketAddr, rtt: Duration) -> MetaAddrChange {
         UpdateResponded {
             addr: canonical_peer_addr(*addr),
+            rtt: Some(rtt),
         }
     }
 
@@ -691,6 +700,7 @@ impl MetaAddr {
             untrusted_last_seen: Some(last_seen),
             last_response: None,
             // these fields aren't sent to the remote peer, but sanitize them anyway
+            rtt: None,
             last_attempt: None,
             last_failure: None,
             last_connection_state: NeverAttemptedGossiped,
@@ -826,6 +836,14 @@ impl MetaAddrChange {
         }
     }
 
+    /// Return the RTT for this change, if available
+    pub fn rtt(&self) -> Option<Duration> {
+        match self {
+            UpdateResponded { rtt, .. } => *rtt,
+            _ => None,
+        }
+    }
+
     /// Return the last failure for this change, if available.
     pub fn last_failure(&self, now: Instant) -> Option<Instant> {
         match self {
@@ -864,6 +882,7 @@ impl MetaAddrChange {
             services: self.untrusted_services(),
             untrusted_last_seen: self.untrusted_last_seen(local_now),
             last_response: self.last_response(local_now),
+            rtt: self.rtt(),
             last_attempt: self.last_attempt(instant_now),
             last_failure: self.last_failure(instant_now),
             last_connection_state: self.peer_addr_state(),
@@ -907,6 +926,7 @@ impl MetaAddrChange {
             services: self.untrusted_services(),
             untrusted_last_seen: self.untrusted_last_seen(local_now),
             last_response: self.last_response(local_now),
+            rtt: None,
             last_attempt: None,
             last_failure: None,
             last_connection_state: self.peer_addr_state(),
@@ -1061,6 +1081,7 @@ impl MetaAddrChange {
                     .or_else(|| self.untrusted_last_seen(local_now)),
                 // The peer has not been attempted, so these fields must be None
                 last_response: None,
+                rtt: None,
                 last_attempt: None,
                 last_failure: None,
                 last_connection_state: self.peer_addr_state(),
@@ -1083,6 +1104,7 @@ impl MetaAddrChange {
                 // This is a wall clock time, but we already checked that responses are in order.
                 // Even if the wall clock time has jumped, we want to use the latest time.
                 last_response: self.last_response(local_now).or(previous.last_response),
+                rtt: self.rtt(),
                 // These are monotonic times, we already checked the responses are in order.
                 last_attempt: self.last_attempt(instant_now).or(previous.last_attempt),
                 last_failure: self.last_failure(instant_now).or(previous.last_failure),

zebra-network/src/meta_addr/tests/vectors.rs

@@ -1,6 +1,6 @@
 //! Fixed test cases for MetaAddr and MetaAddrChange.
 
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 use chrono::Utc;
 
@@ -33,6 +33,7 @@ fn sanitize_extremes() {
         services: Default::default(),
         untrusted_last_seen: Some(u32::MIN.into()),
         last_response: Some(u32::MIN.into()),
+        rtt: Some(Duration::ZERO),
         last_attempt: None,
         last_failure: None,
         last_connection_state: Default::default(),
@@ -45,6 +46,7 @@ fn sanitize_extremes() {
         services: Default::default(),
         untrusted_last_seen: Some(u32::MAX.into()),
         last_response: Some(u32::MAX.into()),
+        rtt: Some(Duration::ZERO),
         last_attempt: None,
         last_failure: None,
         last_connection_state: Default::default(),
@@ -154,7 +156,7 @@ fn recently_responded_peer_is_gossipable() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let peer = MetaAddr::new_responded(address)
+    let peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -174,7 +176,7 @@ fn not_so_recently_responded_peer_is_still_gossipable() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let mut peer = MetaAddr::new_responded(address)
+    let mut peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -204,7 +206,7 @@ fn responded_long_ago_peer_is_not_gossipable() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let mut peer = MetaAddr::new_responded(address)
+    let mut peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -234,7 +236,7 @@ fn long_delayed_change_is_not_applied() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let peer = MetaAddr::new_responded(address)
+    let peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -277,7 +279,7 @@ fn later_revert_change_is_applied() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let peer = MetaAddr::new_responded(address)
+    let peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -319,7 +321,7 @@ fn concurrent_state_revert_change_is_not_applied() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let peer = MetaAddr::new_responded(address)
+    let peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -378,7 +380,7 @@ fn concurrent_state_progress_change_is_applied() {
     let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
 
     // Create a peer that has responded
-    let peer = MetaAddr::new_responded(address)
+    let peer = MetaAddr::new_responded(address, Duration::ZERO)
         .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
         .expect("Failed to create MetaAddr for responded peer");
 
@@ -423,3 +425,24 @@ fn concurrent_state_progress_change_is_applied() {
          peer: {peer:?}"
     );
 }
+
+#[test]
+fn rtt_is_stored_correctly_in_meta_addr() {
+    let _init_guard = zebra_test::init();
+
+    let instant_now = Instant::now();
+    let chrono_now = Utc::now();
+    let local_now: DateTime32 = chrono_now.try_into().expect("will succeed until 2038");
+
+    let address = PeerSocketAddr::from(([192, 168, 180, 9], 10_000));
+    let peer_seed = MetaAddr::new_initial_peer(address).into_new_meta_addr(instant_now, local_now);
+
+    let rtt = Duration::from_millis(128);
+
+    // Create a peer that has responded
+    let peer = MetaAddr::new_responded(address, rtt)
+        .apply_to_meta_addr(peer_seed, instant_now, chrono_now)
+        .expect("Failed to create MetaAddr for responded peer");
+
+    assert_eq!(peer.rtt, Some(rtt));
+}

zebra-network/src/peer/handshake.rs

@@ -9,6 +9,7 @@ use std::{
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
+    time::Duration,
 };
 
 use chrono::{TimeZone, Utc};
@@ -1313,20 +1314,22 @@ async fn send_periodic_heartbeats_run_loop(
         // We've reached another heartbeat interval without
         // shutting down, so do a heartbeat request.
         let heartbeat = send_one_heartbeat(&mut server_tx);
-        heartbeat_timeout(heartbeat, &heartbeat_ts_collector, &connected_addr).await?;
+        let rtt = heartbeat_timeout(heartbeat, &heartbeat_ts_collector, &connected_addr).await?;
 
-        // # Security
-        //
-        // Peer heartbeats are rate-limited because:
-        // - opening connections is rate-limited
-        // - the number of connections is limited
-        // - Zebra initiates each heartbeat using a timer
-        if let Some(book_addr) = connected_addr.get_address_book_addr() {
-            // the collector doesn't depend on network activity,
-            // so this await should not hang
-            let _ = heartbeat_ts_collector
-                .send(MetaAddr::new_responded(book_addr))
-                .await;
+        if let Some(rtt) = rtt {
+            // # Security
+            //
+            // Peer heartbeats are rate-limited because:
+            // - opening connections is rate-limited
+            // - the number of connections is limited
+            // - Zebra initiates each heartbeat using a timer
+            if let Some(book_addr) = connected_addr.get_address_book_addr() {
+                // the collector doesn't depend on network activity,
+                // so this await should not hang
+                let _ = heartbeat_ts_collector
+                    .send(MetaAddr::new_responded(book_addr, rtt))
+                    .await;
+            }
         }
     }
 
@@ -1336,7 +1339,7 @@ async fn send_periodic_heartbeats_run_loop(
 /// Send one heartbeat using `server_tx`.
 async fn send_one_heartbeat(
     server_tx: &mut futures::channel::mpsc::Sender<ClientRequest>,
-) -> Result<(), BoxError> {
+) -> Result<Response, BoxError> {
     // We just reached a heartbeat interval, so start sending
     // a heartbeat.
     let (tx, rx) = oneshot::channel();
@@ -1376,23 +1379,20 @@ async fn send_one_heartbeat(
     // Heartbeats are checked internally to the
     // connection logic, but we need to wait on the
     // response to avoid canceling the request.
-    rx.await??;
-    tracing::trace!("got heartbeat response");
+    let response = rx.await??;
+    tracing::trace!(?response, "got heartbeat response");
 
-    Ok(())
+    Ok(response)
 }
 
 /// Wrap `fut` in a timeout, handing any inner or outer errors using
 /// `handle_heartbeat_error`.
-async fn heartbeat_timeout<F, T>(
-    fut: F,
+async fn heartbeat_timeout(
+    fut: impl Future<Output = Result<Response, BoxError>>,
     address_book_updater: &tokio::sync::mpsc::Sender<MetaAddrChange>,
     connected_addr: &ConnectedAddr,
-) -> Result<T, BoxError>
-where
-    F: Future<Output = Result<T, BoxError>>,
-{
-    let t = match timeout(constants::HEARTBEAT_INTERVAL, fut).await {
+) -> Result<Option<Duration>, BoxError> {
+    let response = match timeout(constants::HEARTBEAT_INTERVAL, fut).await {
         Ok(inner_result) => {
             handle_heartbeat_error(inner_result, address_book_updater, connected_addr).await?
         }
@@ -1401,7 +1401,12 @@ where
         }
     };
 
-    Ok(t)
+    let rtt = match response {
+        Response::Pong(rtt) => Some(rtt),
+        _ => None,
+    };
+
+    Ok(rtt)
 }
 
 /// If `result.is_err()`, mark `connected_addr` as failed using `address_book_updater`.