fix(iroh): fix very slow initial connection establishment (#3434)

## Description

Fixes very slow (~1s) initial connection establishment

Also adds a test for initial connection delay.

Initial connection was *very* slow even with relays disabled and static
discovery. There was a fixed delay.

## Breaking Changes

None

## Notes & open questions

~~Question: not sure how we can turn this into a proper test without
making it super flaky, since it is a timing test. Maybe measure the
ratio between first and second round and have some very conservative
bounds?~~

Changed the test name and put the timings in the assertion. I think the
test is good to go as is. The ratio is extremely conservative, so it
should rarely flake out.

Note: I do 32 connections as an attempt to make the timings more
reliable, but strictly speaking it would also work with 1 connection.

```
thread 'endpoint::tests::initial_connection_time' panicked at iroh/src/endpoint.rs:3269:9:
First round: 32.230450542s, second round 0.063144708s
```

## Change checklist
<!-- Remove any that are not relevant. -->
- [ ] 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.
- [ ] List all breaking changes in the above "Breaking Changes" section.
- [ ] Open an issue or PR on any number0 repos that are affected by this
breaking change. Give guidance on how the updates should be handled or
do the actual updates themselves. The major ones are:
    - [ ] [`quic-rpc`](https://github.com/n0-computer/quic-rpc)
    - [ ] [`iroh-gossip`](https://github.com/n0-computer/iroh-gossip)
    - [ ] [`iroh-blobs`](https://github.com/n0-computer/iroh-blobs)
    - [ ] [`dumbpipe`](https://github.com/n0-computer/dumbpipe)
    - [ ] [`sendme`](https://github.com/n0-computer/sendme)

---------

Co-authored-by: dignifiedquire <[email protected]>

Author: rklaehn

iroh/src/endpoint.rs

@@ -2253,7 +2253,7 @@ mod tests {
     };
 
     use iroh_base::{NodeAddr, NodeId, SecretKey};
-    use n0_future::{StreamExt, task::AbortOnDropHandle};
+    use n0_future::{BufferedStreamExt, StreamExt, stream, task::AbortOnDropHandle};
     use n0_snafu::{Error, Result, ResultExt};
     use n0_watcher::Watcher;
     use quinn::ConnectionError;
@@ -2264,7 +2264,9 @@ mod tests {
     use super::Endpoint;
     use crate::{
         RelayMode,
+        discovery::static_provider::StaticProvider,
         endpoint::{ConnectOptions, Connection, ConnectionType, RemoteInfo},
+        protocol::{AcceptError, ProtocolHandler, Router},
         test_utils::{run_relay_server, run_relay_server_with},
     };
 
@@ -3194,4 +3196,79 @@ mod tests {
 
         Ok(())
     }
+
+    /// Tests that initial connection establishment isn't extremely slow compared
+    /// to subsequent connections.
+    ///
+    /// This is a time based test, but uses a very large ratio to reduce flakiness.
+    /// It also does a number of connections to average out any anomalies.
+    #[tokio::test]
+    #[traced_test]
+    async fn connect_multi_time() -> Result {
+        let n = 32;
+
+        const NOOP_ALPN: &[u8] = b"noop";
+
+        #[derive(Debug, Clone)]
+        struct Noop;
+
+        impl ProtocolHandler for Noop {
+            async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
+                connection.closed().await;
+                Ok(())
+            }
+        }
+
+        async fn noop_server() -> Result<(Router, NodeAddr)> {
+            let endpoint = Endpoint::builder()
+                .relay_mode(RelayMode::Disabled)
+                .bind()
+                .await
+                .e()?;
+            let addr = endpoint.node_addr().initialized().await;
+            let router = Router::builder(endpoint).accept(NOOP_ALPN, Noop).spawn();
+            Ok((router, addr))
+        }
+
+        let routers = stream::iter(0..n)
+            .map(|_| noop_server())
+            .buffered_unordered(32)
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<Result<Vec<_>, _>>()
+            .e()?;
+
+        let addrs = routers
+            .iter()
+            .map(|(_, addr)| addr.clone())
+            .collect::<Vec<_>>();
+        let ids = addrs.iter().map(|addr| addr.node_id).collect::<Vec<_>>();
+        let discovery = StaticProvider::from_node_info(addrs);
+        let endpoint = Endpoint::builder()
+            .relay_mode(RelayMode::Disabled)
+            .discovery(discovery)
+            .bind()
+            .await
+            .e()?;
+        // wait for the endpoint to be initialized. This should not be needed,
+        // but we don't want to measure endpoint init time but connection time
+        // from a fully initialized endpoint.
+        endpoint.node_addr().initialized().await;
+        let t0 = Instant::now();
+        for id in &ids {
+            let conn = endpoint.connect(*id, NOOP_ALPN).await?;
+            conn.close(0u32.into(), b"done");
+        }
+        let dt0 = t0.elapsed().as_secs_f64();
+        let t1 = Instant::now();
+        for id in &ids {
+            let conn = endpoint.connect(*id, NOOP_ALPN).await?;
+            conn.close(0u32.into(), b"done");
+        }
+        let dt1 = t1.elapsed().as_secs_f64();
+
+        assert!(dt0 / dt1 < 20.0, "First round: {dt0}s, second round {dt1}s");
+        Ok(())
+    }
 }

iroh/src/magicsock.rs

@@ -485,6 +485,7 @@ impl MagicSock {
     #[instrument(skip_all)]
     fn prepare_send(
         &self,
+        udp_sender: &UdpSender,
         transmit: &quinn_udp::Transmit,
     ) -> io::Result<SmallVec<[transports::Addr; 3]>> {
         self.metrics
@@ -526,9 +527,7 @@ impl MagicSock {
                 ) {
                     Some((node_id, udp_addr, relay_url, ping_actions)) => {
                         if !ping_actions.is_empty() {
-                            self.actor_sender
-                                .try_send(ActorMessage::PingActions(ping_actions))
-                                .ok();
+                            self.try_send_ping_actions(udp_sender, ping_actions).ok();
                         }
                         if let Some(addr) = udp_addr {
                             active_paths.push(transports::Addr::from(addr));
@@ -1018,6 +1017,118 @@ impl MagicSock {
             }
         }
     }
+    /// Tries to send out the given ping actions out.
+    fn try_send_ping_actions(&self, sender: &UdpSender, msgs: Vec<PingAction>) -> io::Result<()> {
+        for msg in msgs {
+            // Abort sending as soon as we know we are shutting down.
+            if self.is_closing() || self.is_closed() {
+                return Ok(());
+            }
+            match msg {
+                PingAction::SendCallMeMaybe {
+                    relay_url,
+                    dst_node,
+                } => {
+                    // Sends the call-me-maybe DISCO message, queuing if addresses are too stale.
+                    //
+                    // To send the call-me-maybe message, we need to know our current direct addresses.  If
+                    // this information is too stale, the call-me-maybe is queued while a net_report run is
+                    // scheduled.  Once this run finishes, the call-me-maybe will be sent.
+                    match self.direct_addrs.fresh_enough() {
+                        Ok(()) => {
+                            let msg = disco::Message::CallMeMaybe(
+                                self.direct_addrs.to_call_me_maybe_message(),
+                            );
+                            if !self.disco.try_send(
+                                SendAddr::Relay(relay_url.clone()),
+                                dst_node,
+                                msg.clone(),
+                            ) {
+                                warn!(dstkey = %dst_node.fmt_short(), %relay_url, "relay channel full, dropping call-me-maybe");
+                            } else {
+                                debug!(dstkey = %dst_node.fmt_short(), %relay_url, "call-me-maybe sent");
+                            }
+                        }
+                        Err(last_refresh_ago) => {
+                            debug!(
+                                ?last_refresh_ago,
+                                "want call-me-maybe but direct addrs stale; queuing after restun",
+                            );
+                            self.actor_sender
+                                .try_send(ActorMessage::ScheduleDirectAddrUpdate(
+                                    UpdateReason::RefreshForPeering,
+                                    Some((dst_node, relay_url)),
+                                ))
+                                .ok();
+                        }
+                    }
+                }
+                PingAction::SendPing(SendPing {
+                    id,
+                    dst,
+                    dst_node,
+                    tx_id,
+                    purpose,
+                }) => {
+                    let msg = disco::Message::Ping(disco::Ping {
+                        tx_id,
+                        node_key: self.public_key,
+                    });
+
+                    self.try_send_disco_message(sender, dst.clone(), dst_node, msg)?;
+                    debug!(%dst, tx = %HEXLOWER.encode(&tx_id), ?purpose, "ping sent");
+                    let msg_sender = self.actor_sender.clone();
+                    self.node_map
+                        .notify_ping_sent(id, dst, tx_id, purpose, msg_sender);
+                }
+            }
+        }
+        Ok(())
+    }
+
+    /// Tries to send out a disco message.
+    fn try_send_disco_message(
+        &self,
+        sender: &UdpSender,
+        dst: SendAddr,
+        dst_key: PublicKey,
+        msg: disco::Message,
+    ) -> io::Result<()> {
+        let dst = match dst {
+            SendAddr::Udp(addr) => transports::Addr::Ip(addr),
+            SendAddr::Relay(url) => transports::Addr::Relay(url, dst_key),
+        };
+
+        trace!(?dst, %msg, "send disco message (UDP)");
+        if self.is_closed() {
+            return Err(io::Error::new(
+                io::ErrorKind::NotConnected,
+                "connection closed",
+            ));
+        }
+
+        let pkt = self.disco.encode_and_seal(self.public_key, dst_key, &msg);
+
+        let transmit = transports::Transmit {
+            contents: &pkt,
+            ecn: None,
+            segment_size: None,
+        };
+
+        let dst2 = dst.clone();
+        match sender.inner_try_send(&dst2, None, &transmit) {
+            Ok(()) => {
+                trace!(?dst, %msg, "sent disco message");
+                self.metrics.magicsock.sent_disco_udp.inc();
+                disco_message_sent(&msg, &self.metrics.magicsock);
+                Ok(())
+            }
+            Err(err) => {
+                warn!(?dst, ?msg, ?err, "failed to send disco message");
+                Err(err)
+            }
+        }
+    }
 
     /// Publishes our address to a discovery service, if configured.
     ///
@@ -1649,7 +1760,6 @@ impl AsyncUdpSocket for MagicUdpSocket {
 
 #[derive(Debug)]
 enum ActorMessage {
-    PingActions(Vec<PingAction>),
     EndpointPingExpired(usize, stun_rs::TransactionId),
     NetworkChange,
     ScheduleDirectAddrUpdate(UpdateReason, Option<(NodeId, RelayUrl)>),
@@ -1790,7 +1900,7 @@ impl Actor {
 
                     trace!(?msg, "tick: msg");
                     self.msock.metrics.magicsock.actor_tick_msg.inc();
-                    self.handle_actor_message(msg, &sender).await;
+                    self.handle_actor_message(msg).await;
                 }
                 tick = self.periodic_re_stun_timer.tick() => {
                     trace!("tick: re_stun {:?}", tick);
@@ -1949,7 +2059,7 @@ impl Actor {
     /// Processes an incoming actor message.
     ///
     /// Returns `true` if it was a shutdown.
-    async fn handle_actor_message(&mut self, msg: ActorMessage, sender: &UdpSender) {
+    async fn handle_actor_message(&mut self, msg: ActorMessage) {
         match msg {
             ActorMessage::EndpointPingExpired(id, txid) => {
                 self.msock.node_map.notify_ping_timeout(id, txid);
@@ -1969,9 +2079,6 @@ impl Actor {
             ActorMessage::ForceNetworkChange(is_major) => {
                 self.handle_network_change(is_major).await;
             }
-            ActorMessage::PingActions(ping_actions) => {
-                self.handle_ping_actions(sender, ping_actions).await;
-            }
         }
     }
 

iroh/src/magicsock/transports.rs

@@ -454,7 +454,7 @@ impl UdpSender {
     }
 
     /// Best effort sending
-    fn inner_try_send(
+    pub(crate) fn inner_try_send(
         &self,
         destination: &Addr,
         src: Option<IpAddr>,
@@ -504,7 +504,7 @@ impl quinn::UdpSender for UdpSender {
         transmit: &quinn_udp::Transmit,
         cx: &mut Context,
     ) -> Poll<io::Result<()>> {
-        let active_paths = self.msock.prepare_send(transmit)?;
+        let active_paths = self.msock.prepare_send(&self, transmit)?;
 
         if active_paths.is_empty() {
             // Returning Ok here means we let QUIC timeout.
@@ -556,7 +556,7 @@ impl quinn::UdpSender for UdpSender {
     }
 
     fn try_send(self: Pin<&mut Self>, transmit: &quinn_udp::Transmit) -> io::Result<()> {
-        let active_paths = self.msock.prepare_send(transmit)?;
+        let active_paths = self.msock.prepare_send(&self, transmit)?;
         if active_paths.is_empty() {
             // Returning Ok here means we let QUIC timeout.
             // Returning an error would immediately fail a connection.