feat(iroh): improve shutdown interactions (#2980)

## Description

- ties the cancellation tokens between `Router` and `Endpoint` closer
together
- closing `Endpoint` is now idempotent

## Breaking Changes

- added `iroh::Endpoint::close` that takes no arguments anymore, it
default to using code `0` and an empty message
- added `iroh::Endpoint::is_closed`

## Change checklist

- [x] Self-review.
- [x] 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.
- [x] Tests if relevant.
- [x] All breaking changes documented.

Author: dignifiedquire

README.md

@@ -77,7 +77,7 @@ let response = recv.read_to_end(1000).await?;
 assert_eq!(&response, b"Hello, world!");
 
 // Close the endpoint and all its connections
-endpoint.close(0u32.into(), b"bye!").await?;
+endpoint.close().await?;
 ```
 
 And on the accepting side:

iroh/bench/src/lib.rs

@@ -87,7 +87,7 @@ impl EndpointSelector {
     pub async fn close(self) -> Result<()> {
         match self {
             EndpointSelector::Iroh(endpoint) => {
-                endpoint.close(0u32.into(), b"").await?;
+                endpoint.close().await?;
             }
             #[cfg(not(any(target_os = "freebsd", target_os = "openbsd", target_os = "netbsd")))]
             EndpointSelector::Quinn(endpoint) => {

iroh/examples/connect.rs

@@ -90,6 +90,6 @@ async fn main() -> anyhow::Result<()> {
 
     // We received the last message: close all connections and allow for the close
     // message to be sent.
-    endpoint.close(0u8.into(), b"bye").await?;
+    endpoint.close().await?;
     Ok(())
 }

iroh/examples/echo.rs

@@ -53,9 +53,6 @@ async fn connect_side(addr: NodeAddr) -> Result<()> {
     let response = recv.read_to_end(1000).await?;
     assert_eq!(&response, b"Hello, world!");
 
-    // Close the endpoint (and all its connections) in one:
-    endpoint.close(0u32.into(), b"bye!").await?;
-
     Ok(())
 }
 

iroh/examples/transfer.rs

@@ -206,7 +206,7 @@ async fn fetch(ticket: &str, relay_url: Option<String>) -> anyhow::Result<()> {
     // We received the last message: close all connections and allow for the close
     // message to be sent.
     tokio::time::timeout(Duration::from_secs(3), async move {
-        let res = endpoint.close(0u8.into(), b"bye").await;
+        let res = endpoint.close().await;
         if res.is_err() {
             println!("failed to close connection: {res:#?}");
         }

iroh/src/discovery.rs

@@ -405,7 +405,7 @@ impl DiscoveryTask {
         debug!("discovery: start");
         loop {
             let next = tokio::select! {
-                _ = ep.cancelled() => break,
+                _ = ep.cancel_token().cancelled() => break,
                 next = stream.next() => next
             };
             match next {

iroh/src/endpoint.rs

@@ -26,7 +26,7 @@ use derive_more::Debug;
 use futures_lite::{Stream, StreamExt};
 use iroh_base::relay_map::RelayMap;
 use pin_project::pin_project;
-use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, instrument, trace, warn};
 use url::Url;
 
@@ -954,38 +954,46 @@ impl Endpoint {
 
     /// Closes the QUIC endpoint and the magic socket.
     ///
-    /// This will close all open QUIC connections with the provided error_code and
-    /// reason. See [`quinn::Connection`] for details on how these are interpreted.
+    /// This will close any remaining open [`Connection`]s with an error code
+    /// of `0` and an empty reason.  Though it is best practice to close those
+    /// explicitly before with a custom error code and reason.
     ///
-    /// It will then wait for all connections to actually be shutdown, and afterwards close
-    /// the magic socket.  Be aware however that the underlying UDP sockets are only closed
+    /// It will then make a best effort to wait for all close notifications to be
+    /// acknowledged by the peers, re-transmitting them if needed. This ensures the
+    /// peers are aware of the closed connections instead of having to wait for a timeout
+    /// on the connection. Once all connections are closed or timed out, the magic socket is closed.
+    ///
+    /// Be aware however that the underlying UDP sockets are only closed
     /// on [`Drop`], bearing in mind the [`Endpoint`] is only dropped once all the clones
     /// are dropped.
     ///
     /// Returns an error if closing the magic socket failed.
     /// TODO: Document error cases.
-    pub async fn close(&self, error_code: VarInt, reason: &[u8]) -> Result<()> {
-        let Endpoint {
-            msock,
-            endpoint,
-            cancel_token,
-            ..
-        } = self;
-        cancel_token.cancel();
+    pub async fn close(&self) -> Result<()> {
+        if self.is_closed() {
+            return Ok(());
+        }
+
+        self.cancel_token.cancel();
         tracing::debug!("Closing connections");
-        endpoint.close(error_code, reason);
-        endpoint.wait_idle().await;
+        self.endpoint.close(0u16.into(), b"");
+        self.endpoint.wait_idle().await;
 
         tracing::debug!("Connections closed");
-
-        msock.close().await?;
+        self.msock.close().await?;
         Ok(())
     }
 
+    /// Check if this endpoint is still alive, or already closed.
+    pub fn is_closed(&self) -> bool {
+        self.cancel_token.is_cancelled() && self.msock.is_closed()
+    }
+
     // # Remaining private methods
 
-    pub(crate) fn cancelled(&self) -> WaitForCancellationFuture<'_> {
-        self.cancel_token.cancelled()
+    /// Expose the internal [`CancellationToken`] to link shutdowns.
+    pub(crate) fn cancel_token(&self) -> &CancellationToken {
+        &self.cancel_token
     }
 
     /// Return the quic mapped address for this `node_id` and possibly start discovery
@@ -1600,7 +1608,7 @@ mod tests {
 
         info!("closing endpoint");
         // close the endpoint and restart it
-        endpoint.close(0u32.into(), b"done").await.unwrap();
+        endpoint.close().await.unwrap();
 
         info!("restarting endpoint");
         // now restart it and check the addressing info of the peer
@@ -1699,7 +1707,7 @@ mod tests {
                 send.stopped().await.unwrap();
                 recv.read_to_end(0).await.unwrap();
                 info!("client finished");
-                ep.close(0u32.into(), &[]).await.unwrap();
+                ep.close().await.unwrap();
                 info!("client closed");
             }
             .instrument(error_span!("client", %i))

iroh/src/lib.rs

@@ -173,7 +173,7 @@
 //!
 //!     // Gracefully close the connection and endpoint.
 //!     conn.close(1u8.into(), b"done");
-//!     ep.close(0u8.into(), b"ep closing").await?;
+//!     ep.close().await?;
 //!     println!("Client closed");
 //!     Ok(())
 //! }
@@ -202,7 +202,7 @@
 //!
 //!     // Wait for the client to close the connection and gracefully close the endpoint.
 //!     conn.closed().await;
-//!     ep.close(0u8.into(), b"ep closing").await?;
+//!     ep.close().await?;
 //!     Ok(())
 //! }
 //! ```

iroh/src/magicsock.rs

@@ -274,7 +274,7 @@ impl MagicSock {
         self.closing.load(Ordering::Relaxed)
     }
 
-    fn is_closed(&self) -> bool {
+    pub(crate) fn is_closed(&self) -> bool {
         self.closed.load(Ordering::SeqCst)
     }
 
@@ -3238,8 +3238,8 @@ mod tests {
             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?;
+            m1.endpoint.close().await?;
+            m2.endpoint.close().await?;
 
             assert!(msock1.msock.is_closed());
             assert!(msock2.msock.is_closed());

iroh/src/protocol.rs

@@ -268,7 +268,10 @@ impl RouterBuilder {
         let mut join_set = JoinSet::new();
         let endpoint = self.endpoint.clone();
         let protos = protocols.clone();
-        let cancel = CancellationToken::new();
+
+        // We use a child token of the endpoint, to ensure that this is shutdown
+        // when the endpoint is shutdown, but that we can shutdown ourselves independently.
+        let cancel = endpoint.cancel_token().child_token();
         let cancel_token = cancel.clone();
 
         let run_loop_fut = async move {
@@ -341,15 +344,10 @@ impl RouterBuilder {
 
 /// Shutdown the different parts of the router concurrently.
 async fn shutdown(endpoint: &Endpoint, protocols: Arc<ProtocolMap>) {
-    let error_code = 1u16;
-
     // We ignore all errors during shutdown.
     let _ = tokio::join!(
         // Close the endpoint.
-        // Closing the Endpoint is the equivalent of calling Connection::close on all
-        // connections: Operations will immediately fail with ConnectionError::LocallyClosed.
-        // All streams are interrupted, this is not graceful.
-        endpoint.close(error_code.into(), b"provider terminating"),
+        endpoint.close(),
         // Shutdown protocol handlers.
         protocols.shutdown(),
     );
@@ -378,3 +376,24 @@ async fn handle_connection(incoming: crate::endpoint::Incoming, protocols: Arc<P
         warn!("Handling incoming connection ended with error: {err}");
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_shutdown() -> Result<()> {
+        let endpoint = Endpoint::builder().bind().await?;
+        let router = Router::builder(endpoint.clone()).spawn().await?;
+
+        assert!(!router.is_shutdown());
+        assert!(!endpoint.is_closed());
+
+        router.shutdown().await?;
+
+        assert!(router.is_shutdown());
+        assert!(endpoint.is_closed());
+
+        Ok(())
+    }
+}