diff --git a/network/src/network.rs b/network/src/network.rs index d7ce369897..d198450dd8 100644 --- a/network/src/network.rs +++ b/network/src/network.rs @@ -1095,8 +1095,12 @@ impl NetworkService { .unzip(); let receiver: CancellationToken = new_tokio_exit_rx(); + + let server_canceller: CancellationToken = CancellationToken::new(); + let (start_sender, start_receiver) = mpsc::channel(); { + let server_canceller_clone = server_canceller.clone(); let network_state = Arc::clone(&network_state); let p2p_control: ServiceAsyncControl = p2p_control.clone().into(); handle.spawn_task(async move { @@ -1134,6 +1138,15 @@ impl NetworkService { break; }, + _ = server_canceller_clone.cancelled() => { + + debug!("NetworkController dropped, start shutdown..."); + let _ = p2p_control.shutdown().await; + // Drop senders to stop all corresponding background task + drop(bg_signals); + + break; + }, else => { let _ = p2p_control.shutdown().await; // Drop senders to stop all corresponding background task @@ -1165,6 +1178,7 @@ impl NetworkService { network_state, p2p_control, ping_controller, + server_canceller, }) } } @@ -1176,6 +1190,13 @@ pub struct NetworkController { network_state: Arc, p2p_control: ServiceControl, ping_controller: Option>, + server_canceller: CancellationToken, +} + +impl Drop for NetworkController { + fn drop(&mut self) { + self.server_canceller.cancel(); + } } impl NetworkController { diff --git a/util/stop-handler/src/stop_register.rs b/util/stop-handler/src/stop_register.rs index d242fc6034..e122f4b8e0 100644 --- a/util/stop-handler/src/stop_register.rs +++ b/util/stop-handler/src/stop_register.rs @@ -1,5 +1,5 @@ use ckb_channel::TrySendError; -use ckb_logger::{debug, error, info, trace, warn}; +use ckb_logger::{debug, info, trace, warn}; use ckb_util::Mutex; use tokio_util::sync::CancellationToken;