From d5bdc1c8c050ca8f48f87f1178938fc6ffc54932 Mon Sep 17 00:00:00 2001 From: garwah <14845405+garwahl@users.noreply.github.com> Date: Wed, 2 Oct 2024 19:15:12 +1000 Subject: [PATCH] feat(net): Perform graceful network shutdown before dropping network in NetworkManager (#11404) Co-authored-by: garwah --- crates/net/network/src/manager.rs | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs index 1f33f618b1..e8ffb81c5e 100644 --- a/crates/net/network/src/manager.rs +++ b/crates/net/network/src/manager.rs @@ -600,14 +600,7 @@ impl NetworkManager { } NetworkHandleMessage::Shutdown(tx) => { - // Set connection status to `Shutdown`. Stops node to accept - // new incoming connections as well as sending connection requests to newly - // discovered nodes. - self.swarm.on_shutdown_requested(); - // Disconnect all active connections - self.swarm.sessions_mut().disconnect_all(Some(DisconnectReason::ClientQuitting)); - // drop pending connections - self.swarm.sessions_mut().disconnect_all_pending(); + self.perform_network_shutdown(); let _ = tx.send(()); } NetworkHandleMessage::ReputationChange(peer_id, kind) => { @@ -966,10 +959,24 @@ impl NetworkManager { }, } + self.perform_network_shutdown(); let res = shutdown_hook(self); drop(graceful_guard); res } + + /// Performs a graceful network shutdown by stopping new connections from being accepted while + /// draining current and pending connections. + fn perform_network_shutdown(&mut self) { + // Set connection status to `Shutdown`. Stops node from accepting + // new incoming connections as well as sending connection requests to newly + // discovered nodes. + self.swarm.on_shutdown_requested(); + // Disconnect all active connections + self.swarm.sessions_mut().disconnect_all(Some(DisconnectReason::ClientQuitting)); + // drop pending connections + self.swarm.sessions_mut().disconnect_all_pending(); + } } impl Future for NetworkManager {