Add drop_guard to NetworkController

Signed-off-by: Eval EXEC <[email protected]>
nervosnetwork · Oct 7, 2023 · bee38a2 · bee38a2
1 parent 53e9f03
commit bee38a2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git 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<NetworkState>,
     p2p_control: ServiceControl,
     ping_controller: Option<Sender<()>>,
+    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
@@ -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;