fix(iroh-net)!: improve magicsock's shutdown story

iroh-net/src/magic_endpoint.rs

@@ -4,7 +4,7 @@ use std::{net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
 
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use derive_more::Debug;
-use futures::StreamExt;
+use futures::{future, StreamExt};
 use quinn_proto::VarInt;
 use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
 use tracing::{debug, trace};
@@ -560,10 +560,40 @@ impl MagicEndpoint {
     ///
     /// Returns an error if closing the magic socket failed.
     /// TODO: Document error cases.
-    pub async fn close(&self, error_code: VarInt, reason: &[u8]) -> Result<()> {
-        self.cancel_token.cancel();
-        self.endpoint.close(error_code, reason);
-        self.msock.close().await?;
+    pub async fn close(self, error_code: VarInt, reason: &[u8]) -> Result<()> {
+        let MagicEndpoint {
+            msock,
+            endpoint,
+            cancel_token,
+            ..
+        } = self;
+        cancel_token.cancel();
+        endpoint.close(error_code, reason);
+
+        {
+            // waiting for connections to close can be slow: this timeout decide whethers we should
+            // inform the user of a slow operation
+            let logging_timeout = tokio::time::sleep(Duration::from_millis(500));
+            futures::pin_mut!(logging_timeout);
+
+            let wait_idle = endpoint.wait_idle();
+            futures::pin_mut!(wait_idle);
+
+            // race the timeout and closing. Nothing else needs to be done if closing finishes
+            // first. If the timeouts is done first, give the user some logs
+            if let future::Either::Right(((), wait_idle_fut)) =
+                future::select(wait_idle, logging_timeout).await
+            {
+                tracing::info!("Closing connections");
+                wait_idle_fut.await;
+                tracing::info!("Connections closed");
+            }
+        }
+        // In case this is the last clone of `MagicEndpoint`, dropping the `quinn::Endpoint` will
+        // make it more likely that the underlying socket is not polled by quinn anymore after this
+        drop(endpoint);
+
+        msock.close().await?;
         Ok(())
     }
 
@@ -582,8 +612,8 @@ impl MagicEndpoint {
     }
 
     #[cfg(test)]
-    pub(crate) fn magic_sock(&self) -> &MagicSock {
-        &self.msock
+    pub(crate) fn magic_sock(&self) -> MagicSock {
+        self.msock.clone()
     }
     #[cfg(test)]
     pub(crate) fn endpoint(&self) -> &quinn::Endpoint {

iroh-net/src/magicsock.rs

@@ -287,6 +287,11 @@ impl Inner {
         let bytes_total: usize = transmits.iter().map(|t| t.contents.len()).sum();
         inc_by!(MagicsockMetrics, send_data, bytes_total as _);
 
+        let mut n = 0;
+        if transmits.is_empty() {
+            return Poll::Ready(Ok(n));
+        }
+
         if self.is_closed() {
             inc_by!(MagicsockMetrics, send_data_network_down, bytes_total as _);
             return Poll::Ready(Err(io::Error::new(
@@ -295,10 +300,6 @@ impl Inner {
             )));
         }
 
-        let mut n = 0;
-        if transmits.is_empty() {
-            return Poll::Ready(Ok(n));
-        }
         trace!(
             "sending:\n{}",
             transmits.iter().fold(
@@ -484,6 +485,7 @@ impl Inner {
         Ok(sock)
     }
 
+    // NOTE: Receiving on a [`Self::closed`] socket will return [`Poll::Pending`] undefinitely.
     #[instrument(skip_all, fields(me = %self.me))]
     fn poll_recv(
         &self,
@@ -494,10 +496,7 @@ impl Inner {
         // FIXME: currently ipv4 load results in ipv6 traffic being ignored
         debug_assert_eq!(bufs.len(), metas.len(), "non matching bufs & metas");
         if self.is_closed() {
-            return Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::NotConnected,
-                "connection closed",
-            )));
+            return Poll::Pending;
         }
 
         // order of polling is: UDPv4, UDPv6, relay
@@ -2982,11 +2981,13 @@ pub(crate) mod tests {
             let _guard = mesh_stacks(vec![m1.clone(), m2.clone()], url.clone()).await?;
 
             println!("closing endpoints");
+            let msock1 = m1.endpoint.magic_sock();
+            let msock2 = m2.endpoint.magic_sock();
             m1.endpoint.close(0u32.into(), b"done").await?;
             m2.endpoint.close(0u32.into(), b"done").await?;
 
-            assert!(m1.endpoint.magic_sock().inner.is_closed());
-            assert!(m2.endpoint.magic_sock().inner.is_closed());
+            assert!(msock1.inner.is_closed());
+            assert!(msock2.inner.is_closed());
         }
         Ok(())
     }