From aaa2d69e885b4f88ceb387a4c286cd4229ccc2e6 Mon Sep 17 00:00:00 2001 From: dignifiedquire Date: Fri, 29 Nov 2024 16:30:04 +0100 Subject: [PATCH] feat(iroh): simplify close API call --- iroh/bench/src/lib.rs | 2 +- iroh/examples/connect.rs | 2 +- iroh/examples/echo.rs | 3 --- iroh/examples/transfer.rs | 2 +- iroh/src/endpoint.rs | 22 +++++++++++++++++++--- iroh/src/magicsock.rs | 4 ++-- iroh/src/protocol.rs | 7 +------ 7 files changed, 25 insertions(+), 17 deletions(-) diff --git a/iroh/bench/src/lib.rs b/iroh/bench/src/lib.rs index ce890e257ba..9de2f0ac727 100644 --- a/iroh/bench/src/lib.rs +++ b/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) => { diff --git a/iroh/examples/connect.rs b/iroh/examples/connect.rs index 01fd710fbbd..ec287903b27 100644 --- a/iroh/examples/connect.rs +++ b/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(()) } diff --git a/iroh/examples/echo.rs b/iroh/examples/echo.rs index d8bef70698f..bcbdfaa2ec9 100644 --- a/iroh/examples/echo.rs +++ b/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(()) } diff --git a/iroh/examples/transfer.rs b/iroh/examples/transfer.rs index 45aec76ebe0..5f7c2e22c25 100644 --- a/iroh/examples/transfer.rs +++ b/iroh/examples/transfer.rs @@ -206,7 +206,7 @@ async fn fetch(ticket: &str, relay_url: Option) -> 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:#?}"); } diff --git a/iroh/src/endpoint.rs b/iroh/src/endpoint.rs index a257192aa40..d0bd6801ee9 100644 --- a/iroh/src/endpoint.rs +++ b/iroh/src/endpoint.rs @@ -952,6 +952,22 @@ impl Endpoint { // # Methods for terminating the endpoint. + /// Closes the QUIC endpoint and the magic socket. + /// + /// This will close all open QUIC connections. + /// + /// 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 + /// 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) -> Result<()> { + self.close_with_code(1u16.into(), b"shutting down").await?; + Ok(()) + } + /// Closes the QUIC endpoint and the magic socket. /// /// This will close all open QUIC connections with the provided error_code and @@ -964,7 +980,7 @@ impl Endpoint { /// /// Returns an error if closing the magic socket failed. /// TODO: Document error cases. - pub async fn close(&self, error_code: VarInt, reason: &[u8]) -> Result<()> { + pub async fn close_with_code(&self, error_code: VarInt, reason: &[u8]) -> Result<()> { if self.is_closed() { return Ok(()); } @@ -1603,7 +1619,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 @@ -1702,7 +1718,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)) diff --git a/iroh/src/magicsock.rs b/iroh/src/magicsock.rs index 162dcbba075..35f6c70a331 100644 --- a/iroh/src/magicsock.rs +++ b/iroh/src/magicsock.rs @@ -3266,8 +3266,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()); diff --git a/iroh/src/protocol.rs b/iroh/src/protocol.rs index fa6ba672194..89d599d3326 100644 --- a/iroh/src/protocol.rs +++ b/iroh/src/protocol.rs @@ -344,15 +344,10 @@ impl RouterBuilder { /// Shutdown the different parts of the router concurrently. async fn shutdown(endpoint: &Endpoint, protocols: Arc) { - 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(), );