add drop guard to client tasks

Author: jstuczyn

common/client-core/src/client/base_client/mod.rs

@@ -337,8 +337,14 @@ where
             debug_config.cover_traffic,
             stats_tx,
         );
-        shutdown_tracker
-            .try_spawn_named_with_shutdown(async move { stream.run().await }, "CoverTrafficStream");
+        let drop_guard = shutdown_tracker.clone_shutdown_token().drop_guard();
+        shutdown_tracker.try_spawn_named_with_shutdown(
+            async move {
+                let _ = drop_guard;
+                stream.run().await
+            },
+            "CoverTrafficStream",
+        );
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -415,8 +421,10 @@ where
             "AcknowledgementController::RetransmissionRequestListener",
         );
 
+        let drop_guard = shutdown_tracker.clone_shutdown_token().drop_guard();
         shutdown_tracker.try_spawn_named_with_shutdown(
             async move {
+                let _ = drop_guard;
                 sent_notification_listener.run().await;
             },
             "AcknowledgementController::SentNotificationListener",
@@ -427,8 +435,6 @@ where
             async move { ack_action_controller.run(shutdown_token).await },
             "AcknowledgementController::ActionController",
         );
-
-        // .start(packet_type);
     }
 
     // buffer controlling all messages fetched from provider
@@ -701,8 +707,12 @@ where
             // don't spawn the refresher if we don't want to be refreshing the topology.
             // only use the initial values obtained
             info!("Starting topology refresher...");
+            let drop_guard = shutdown_tracker.clone_shutdown_token().drop_guard();
             shutdown_tracker.try_spawn_named_with_shutdown(
-                async move { topology_refresher.run().await },
+                async move {
+                    let _ = drop_guard;
+                    topology_refresher.run().await
+                },
                 "TopologyRefresher",
             );
         }
@@ -917,7 +927,7 @@ where
             self.user_agent.clone(),
             generate_client_stats_id(*self_address.identity()),
             input_sender.clone(),
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         );
 
         // needs to be started as the first thing to block if required waiting for the gateway
@@ -927,7 +937,7 @@ where
             shared_topology_accessor.clone(),
             self_address.gateway(),
             self.wait_for_gateway,
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         )
         .await?;
 
@@ -947,15 +957,15 @@ where
             stats_reporter.clone(),
             #[cfg(unix)]
             self.connection_fd_callback,
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         )
         .await?;
         let gateway_ws_fd = gateway_transceiver.ws_fd();
 
         let reply_storage = Self::setup_persistent_reply_storage(
             reply_storage_backend,
             key_rotation_config,
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         )
         .await?;
 
@@ -966,7 +976,7 @@ where
             reply_storage.key_storage(),
             reply_controller_sender.clone(),
             stats_reporter.clone(),
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         );
 
         // The message_sender is the transmitter for any component generating sphinx packets
@@ -975,7 +985,7 @@ where
         // The MixTrafficController then sends the actual traffic
 
         let (message_sender, client_request_sender) =
-            Self::start_mix_traffic_controller(gateway_transceiver, &shutdown_tracker.clone());
+            Self::start_mix_traffic_controller(gateway_transceiver, &shutdown_tracker);
 
         // Channels that the websocket listener can use to signal downstream to the real traffic
         // controller that connections are closed.
@@ -1004,7 +1014,7 @@ where
             shared_lane_queue_lengths.clone(),
             client_connection_rx,
             stats_reporter.clone(),
-            &shutdown_tracker.clone(),
+            &shutdown_tracker,
         );
 
         if !self
@@ -1020,7 +1030,7 @@ where
                 shared_topology_accessor.clone(),
                 message_sender,
                 stats_reporter.clone(),
-                &shutdown_tracker.clone(),
+                &shutdown_tracker,
             );
         }
 

common/client-core/src/client/mix_traffic/mod.rs

@@ -138,6 +138,7 @@ impl MixTrafficController {
 
     pub async fn run(&mut self) {
         debug!("Started MixTrafficController with graceful shutdown support");
+        let _drop_guard = self.shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs

@@ -80,6 +80,7 @@ impl AcknowledgementListener {
     pub(crate) async fn run(&mut self, shutdown_token: ShutdownToken) {
         debug!("Started AcknowledgementListener with graceful shutdown support");
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs

@@ -245,6 +245,7 @@ impl ActionController {
     pub(crate) async fn run(&mut self, shutdown_token: ShutdownToken) {
         debug!("Started ActionController with graceful shutdown support");
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs

@@ -216,6 +216,7 @@ where
     pub(crate) async fn run(&mut self, shutdown_token: ShutdownToken) {
         debug!("Started InputMessageListener with graceful shutdown support");
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs

@@ -167,6 +167,7 @@ where
     pub(crate) async fn run(&mut self, shutdown_token: ShutdownToken) {
         debug!("Started RetransmissionRequestListener with graceful shutdown support");
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/real_messages_control/real_traffic_stream.rs

@@ -579,6 +579,7 @@ where
 
         // avoid borrow on self
         let shutdown_token = self.shutdown_token.clone();
+        let _drop_guard = shutdown_token.clone().drop_guard();
         #[cfg(not(target_arch = "wasm32"))]
         {
             let mut status_timer = tokio::time::interval(Duration::from_secs(5));

common/client-core/src/client/received_buffer.rs

@@ -497,6 +497,8 @@ impl<R: MessageReceiver> RequestReceiver<R> {
 
     pub(crate) async fn run(&mut self) {
         debug!("Started RequestReceiver with graceful shutdown support");
+
+        let _drop_guard = self.shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;
@@ -540,6 +542,8 @@ impl<R: MessageReceiver> FragmentedMessageReceiver<R> {
 
     pub(crate) async fn run(&mut self) -> Result<(), MessageRecoveryError> {
         debug!("Started FragmentedMessageReceiver with graceful shutdown support");
+
+        let _drop_guard = self.shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/replies/reply_controller/mod.rs

@@ -152,6 +152,7 @@ where
         let polling_rate = self.config.key_rotation.epoch_duration / 8;
         let mut invalidation_inspection = new_interval_stream(polling_rate);
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
         loop {
             tokio::select! {
                 biased;

common/client-core/src/client/statistics_control.rs

@@ -119,6 +119,8 @@ impl StatisticsControl {
         let mut snapshot_interval =
             gloo_timers::future::IntervalStream::new(SNAPSHOT_INTERVAL.as_millis() as u32);
 
+        let _drop_guard = shutdown_token.clone().drop_guard();
+
         loop {
             tokio::select! {
                 biased;