feat(iroh-relay): send regular pings to check the connection (#3113)

## Description

Because of issues with TCP, we now send regular ping-pongs from client
to server and vice versa.

## Breaking Changes

<!-- Optional, if there are any breaking changes document them,
including how to migrate older code. -->

## Notes & open questions

<!-- Any notes, remarks or open questions you have to make about the PR.
-->

## Change checklist

- [ ] Self-review.
- [ ] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [ ] Tests if relevant.
- [ ] All breaking changes documented.

---------

Co-authored-by: Floris Bruynooghe <[email protected]>

Author: dignifiedquire

iroh-relay/src/lib.rs

@@ -38,6 +38,8 @@ pub mod quic;
 #[cfg(feature = "server")]
 pub mod server;
 
+mod ping_tracker;
+
 mod key_cache;
 mod relay_map;
 pub(crate) use key_cache::KeyCache;
@@ -47,4 +49,7 @@ mod dns;
 
 pub use protos::relay::MAX_PACKET_SIZE;
 
-pub use self::relay_map::{RelayMap, RelayNode, RelayQuicConfig};
+pub use self::{
+    ping_tracker::PingTracker,
+    relay_map::{RelayMap, RelayNode, RelayQuicConfig},
+};

iroh-relay/src/ping_tracker.rs

@@ -0,0 +1,78 @@
+use std::time::{Duration, Instant};
+
+use tracing::debug;
+
+/// Maximum time for a ping response in the relay protocol.
+pub const PING_TIMEOUT: Duration = Duration::from_secs(5);
+
+/// Tracks pings on a single relay connection.
+///
+/// Only the last ping needs is useful, any previously sent ping is forgotten and ignored.
+#[derive(Debug)]
+pub struct PingTracker {
+    inner: Option<PingInner>,
+    default_timeout: Duration,
+}
+
+#[derive(Debug)]
+struct PingInner {
+    data: [u8; 8],
+    deadline: Instant,
+}
+
+impl Default for PingTracker {
+    fn default() -> Self {
+        Self::new(PING_TIMEOUT)
+    }
+}
+
+impl PingTracker {
+    /// Creates a new ping tracker, setting the ping timeout for pings.
+    pub fn new(default_timeout: Duration) -> Self {
+        Self {
+            inner: None,
+            default_timeout,
+        }
+    }
+
+    /// Returns the current timeout set for pings.
+    pub fn default_timeout(&self) -> Duration {
+        self.default_timeout
+    }
+
+    /// Starts a new ping.
+    pub fn new_ping(&mut self) -> [u8; 8] {
+        let ping_data = rand::random();
+        debug!(data = ?ping_data, "Sending ping to relay server.");
+        self.inner = Some(PingInner {
+            data: ping_data,
+            deadline: Instant::now() + self.default_timeout,
+        });
+        ping_data
+    }
+
+    /// Updates the ping tracker with a received pong.
+    ///
+    /// Only the pong of the most recent ping will do anything.  There is no harm feeding
+    /// any pong however.
+    pub fn pong_received(&mut self, data: [u8; 8]) {
+        if self.inner.as_ref().map(|inner| inner.data) == Some(data) {
+            debug!(?data, "Pong received from relay server");
+            self.inner = None;
+        }
+    }
+
+    /// Cancel-safe waiting for a ping timeout.
+    ///
+    /// Unless the most recent sent ping times out, this will never return.
+    pub async fn timeout(&mut self) {
+        match self.inner {
+            Some(PingInner { deadline, data }) => {
+                tokio::time::sleep_until(deadline.into()).await;
+                debug!(?data, "Ping timeout.");
+                self.inner = None;
+            }
+            None => std::future::pending().await,
+        }
+    }
+}

iroh-relay/src/protos/relay.rs

@@ -41,8 +41,13 @@ const MAX_FRAME_SIZE: usize = 1024 * 1024;
 /// The Relay magic number, sent in the FrameType::ClientInfo frame upon initial connection.
 const MAGIC: &str = "RELAY🔑";
 
+/// Interval in which we ping the relay server to ensure the connection is alive.
+///
+/// The default QUIC max_idle_timeout is 30s, so setting that to half this time gives some
+/// chance of recovering.
 #[cfg(feature = "server")]
-pub(crate) const KEEP_ALIVE: Duration = Duration::from_secs(60);
+pub(crate) const PING_INTERVAL: Duration = Duration::from_secs(15);
+
 /// The number of packets buffered for sending per client
 #[cfg(feature = "server")]
 pub(crate) const PER_CLIENT_SEND_QUEUE_DEPTH: usize = 512; //32;

iroh-relay/src/server/client.rs

@@ -9,16 +9,21 @@ use futures_sink::Sink;
 use futures_util::{SinkExt, Stream, StreamExt};
 use iroh_base::NodeId;
 use iroh_metrics::{inc, inc_by};
-use tokio::sync::mpsc::{self, error::TrySendError};
+use rand::Rng;
+use tokio::{
+    sync::mpsc::{self, error::TrySendError},
+    time::MissedTickBehavior,
+};
 use tokio_util::{sync::CancellationToken, task::AbortOnDropHandle};
 use tracing::{debug, error, instrument, trace, warn, Instrument};
 
 use crate::{
     protos::{
         disco,
-        relay::{write_frame, Frame, KEEP_ALIVE},
+        relay::{write_frame, Frame, PING_INTERVAL},
     },
     server::{clients::Clients, metrics::Metrics, streams::RelayedStream, ClientRateLimit},
+    PingTracker,
 };
 
 /// A request to write a dataframe to a Client
@@ -102,6 +107,7 @@ impl Client {
             node_id,
             connection_id,
             clients: clients.clone(),
+            ping_tracker: PingTracker::default(),
         };
 
         // start io loop
@@ -201,6 +207,7 @@ struct Actor {
     connection_id: u64,
     /// Reference to the other connected clients.
     clients: Clients,
+    ping_tracker: PingTracker,
 }
 
 impl Actor {
@@ -220,10 +227,16 @@ impl Actor {
     }
 
     async fn run_inner(&mut self, done: CancellationToken) -> Result<()> {
-        let jitter = Duration::from_secs(5);
-        let mut keep_alive = tokio::time::interval(KEEP_ALIVE + jitter);
+        // Add some jitter to ping pong interactions, to avoid all pings being sent at the same time
+        let next_interval = || {
+            let random_secs = rand::rngs::OsRng.gen_range(1..=5);
+            Duration::from_secs(random_secs) + PING_INTERVAL
+        };
+
+        let mut ping_interval = tokio::time::interval(next_interval());
         // ticks immediately
-        keep_alive.tick().await;
+        ping_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
+        ping_interval.tick().await;
 
         loop {
             tokio::select! {
@@ -237,6 +250,8 @@ impl Actor {
                 }
                 maybe_frame = self.stream.next() => {
                     self.handle_frame(maybe_frame).await.context("handle read")?;
+                    // reset the ping interval, we just received a message
+                    ping_interval.reset();
                 }
                 // First priority, disco packets
                 packet = self.disco_send_queue.recv() => {
@@ -254,11 +269,19 @@ impl Actor {
                     trace!("node_id gone: {:?}", node_id);
                     self.write_frame(Frame::NodeGone { node_id }).await?;
                 }
-                _ = keep_alive.tick() => {
-                    trace!("keep alive");
-                    self.write_frame(Frame::KeepAlive).await?;
+                _ = self.ping_tracker.timeout() => {
+                    trace!("pong timed out");
+                    break;
+                }
+                _ = ping_interval.tick() => {
+                    trace!("keep alive ping");
+                    // new interval
+                    ping_interval.reset_after(next_interval());
+                    let data = self.ping_tracker.new_ping();
+                    self.write_frame(Frame::Ping { data }).await?;
                 }
             }
+
             self.stream.flush().await.context("tick flush")?;
         }
         Ok(())
@@ -321,6 +344,7 @@ impl Actor {
             Some(frame) => frame?,
             None => anyhow::bail!("stream terminated"),
         };
+
         match frame {
             Frame::SendPacket { dst_key, packet } => {
                 let packet_len = packet.len();
@@ -333,6 +357,9 @@ impl Actor {
                 self.write_frame(Frame::Pong { data }).await?;
                 inc!(Metrics, sent_pong);
             }
+            Frame::Pong { data } => {
+                self.ping_tracker.pong_received(data);
+            }
             Frame::Health { problem } => {
                 bail!("server issue: {:?}", problem);
             }
@@ -567,6 +594,7 @@ mod tests {
             connection_id: 0,
             node_id,
             clients: clients.clone(),
+            ping_tracker: PingTracker::default(),
         };
 
         let done = CancellationToken::new();