Move TaskClient to Self in few tasks

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

@@ -310,7 +310,7 @@ where
         topology_accessor: TopologyAccessor,
         mix_tx: BatchMixMessageSender,
         stats_tx: ClientStatsSender,
-        shutdown: TaskClient,
+        task_client: TaskClient,
     ) {
         info!("Starting loop cover traffic stream...");
 
@@ -323,9 +323,10 @@ where
             debug_config.traffic,
             debug_config.cover_traffic,
             stats_tx,
+            task_client,
         );
 
-        stream.start_with_shutdown(shutdown);
+        stream.start();
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -340,7 +341,7 @@ where
         reply_controller_receiver: ReplyControllerReceiver,
         lane_queue_lengths: LaneQueueLengths,
         client_connection_rx: ConnectionCommandReceiver,
-        shutdown: TaskClient,
+        task_client: TaskClient,
         packet_type: PacketType,
         stats_tx: ClientStatsSender,
     ) {
@@ -358,8 +359,9 @@ where
             lane_queue_lengths,
             client_connection_rx,
             stats_tx,
+            task_client,
         )
-        .start_with_shutdown(shutdown, packet_type);
+        .start(packet_type);
     }
 
     // buffer controlling all messages fetched from provider

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

@@ -13,6 +13,7 @@ use nym_sphinx::cover::generate_loop_cover_packet;
 use nym_sphinx::params::{PacketSize, PacketType};
 use nym_sphinx::utils::sample_poisson_duration;
 use nym_statistics_common::clients::{packet_statistics::PacketStatisticsEvent, ClientStatsSender};
+use nym_task::TaskClient;
 use rand::{rngs::OsRng, CryptoRng, Rng};
 use std::pin::Pin;
 use std::sync::Arc;
@@ -64,6 +65,8 @@ where
     packet_type: PacketType,
 
     stats_tx: ClientStatsSender,
+
+    task_client: TaskClient,
 }
 
 impl<R> Stream for LoopCoverTrafficStream<R>
@@ -110,6 +113,7 @@ impl LoopCoverTrafficStream<OsRng> {
         traffic_config: config::Traffic,
         cover_config: config::CoverTraffic,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         let rng = OsRng;
 
@@ -128,6 +132,7 @@ impl LoopCoverTrafficStream<OsRng> {
             secondary_packet_size: traffic_config.secondary_packet_size,
             packet_type: traffic_config.packet_type,
             stats_tx,
+            task_client,
         }
     }
 
@@ -224,7 +229,7 @@ impl LoopCoverTrafficStream<OsRng> {
         tokio::task::yield_now().await;
     }
 
-    pub fn start_with_shutdown(mut self, mut shutdown: nym_task::TaskClient) {
+    pub fn start(mut self) {
         if self.cover_traffic.disable_loop_cover_traffic_stream {
             // we should have never got here in the first place - the task should have never been created to begin with
             // so panic and review the code that lead to this branch
@@ -238,6 +243,8 @@ impl LoopCoverTrafficStream<OsRng> {
         );
         self.set_next_delay(sampled);
 
+        let mut shutdown = self.task_client.fork("select");
+
         spawn_future(async move {
             debug!("Started LoopCoverTrafficStream with graceful shutdown support");
 

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

@@ -24,6 +24,7 @@ use nym_sphinx::{
     Delay as SphinxDelay,
 };
 use nym_statistics_common::clients::ClientStatsSender;
+use nym_task::TaskClient;
 use rand::{CryptoRng, Rng};
 use std::{
     sync::{Arc, Weak},
@@ -203,6 +204,7 @@ where
     retransmission_request_listener: RetransmissionRequestListener<R>,
     sent_notification_listener: SentNotificationListener,
     action_controller: ActionController,
+    task_client: TaskClient,
 }
 
 impl<R> AcknowledgementController<R>
@@ -216,6 +218,7 @@ where
         message_handler: MessageHandler<R>,
         reply_controller_sender: ReplyControllerSender,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         let (retransmission_tx, retransmission_rx) = mpsc::unbounded();
 
@@ -262,21 +265,18 @@ where
             retransmission_request_listener,
             sent_notification_listener,
             action_controller,
+            task_client,
         }
     }
 
-    pub(super) fn start_with_shutdown(
-        self,
-        shutdown: nym_task::TaskClient,
-        packet_type: PacketType,
-    ) {
+    pub(super) fn start(self, packet_type: PacketType) {
         let mut acknowledgement_listener = self.acknowledgement_listener;
         let mut input_message_listener = self.input_message_listener;
         let mut retransmission_request_listener = self.retransmission_request_listener;
         let mut sent_notification_listener = self.sent_notification_listener;
         let mut action_controller = self.action_controller;
 
-        let shutdown_handle = shutdown.fork("acknowledgement_listener");
+        let shutdown_handle = self.task_client.fork("acknowledgement_listener");
         spawn_future(async move {
             acknowledgement_listener
                 .run_with_shutdown(shutdown_handle)
@@ -284,7 +284,7 @@ where
             debug!("The acknowledgement listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("input_message_listener");
+        let shutdown_handle = self.task_client.fork("input_message_listener");
         spawn_future(async move {
             input_message_listener
                 .run_with_shutdown(shutdown_handle)
@@ -292,7 +292,7 @@ where
             debug!("The input listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("retransmission_request_listener");
+        let shutdown_handle = self.task_client.fork("retransmission_request_listener");
         spawn_future(async move {
             retransmission_request_listener
                 .run_with_shutdown(shutdown_handle, packet_type)
@@ -300,7 +300,7 @@ where
             debug!("The retransmission request listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("sent_notification_listener");
+        let shutdown_handle = self.task_client.fork("sent_notification_listener");
         spawn_future(async move {
             sent_notification_listener
                 .run_with_shutdown(shutdown_handle)
@@ -310,7 +310,7 @@ where
 
         spawn_future(async move {
             action_controller
-                .run_with_shutdown(shutdown.with_suffix("action_controller"))
+                .run_with_shutdown(self.task_client.with_suffix("action_controller"))
                 .await;
             debug!("The controller has finished execution!");
         });

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

@@ -31,6 +31,7 @@ use nym_sphinx::addressing::clients::Recipient;
 use nym_sphinx::params::PacketType;
 use nym_statistics_common::clients::ClientStatsSender;
 use nym_task::connections::{ConnectionCommandReceiver, LaneQueueLengths};
+use nym_task::TaskClient;
 use rand::{rngs::OsRng, CryptoRng, Rng};
 use std::sync::Arc;
 
@@ -147,6 +148,7 @@ impl RealMessagesController<OsRng> {
         lane_queue_lengths: LaneQueueLengths,
         client_connection_rx: ConnectionCommandReceiver,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         let rng = OsRng;
 
@@ -186,6 +188,7 @@ impl RealMessagesController<OsRng> {
             message_handler.clone(),
             reply_controller_sender,
             stats_tx.clone(),
+            task_client.fork("ack_control"),
         );
 
         let reply_control = ReplyController::new(
@@ -193,6 +196,7 @@ impl RealMessagesController<OsRng> {
             message_handler,
             reply_storage,
             reply_controller_receiver,
+            task_client.fork("reply_controller"),
         );
 
         let out_queue_control = OutQueueControl::new(
@@ -205,6 +209,7 @@ impl RealMessagesController<OsRng> {
             lane_queue_lengths,
             client_connection_rx,
             stats_tx,
+            task_client.with_suffix("out_queue_control"),
         );
 
         RealMessagesController {
@@ -214,22 +219,20 @@ impl RealMessagesController<OsRng> {
         }
     }
 
-    pub fn start_with_shutdown(self, shutdown: nym_task::TaskClient, packet_type: PacketType) {
+    pub fn start(self, packet_type: PacketType) {
         let mut out_queue_control = self.out_queue_control;
         let ack_control = self.ack_control;
         let mut reply_control = self.reply_control;
 
-        let shutdown_handle = shutdown.fork("out_queue_control");
         spawn_future(async move {
-            out_queue_control.run_with_shutdown(shutdown_handle).await;
+            out_queue_control.run().await;
             debug!("The out queue controller has finished execution!");
         });
-        let shutdown_handle = shutdown.fork("reply_control");
         spawn_future(async move {
-            reply_control.run_with_shutdown(shutdown_handle).await;
+            reply_control.run().await;
             debug!("The reply controller has finished execution!");
         });
 
-        ack_control.start_with_shutdown(shutdown.with_suffix("ack_control"), packet_type);
+        ack_control.start(packet_type);
     }
 }

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

@@ -22,6 +22,7 @@ use nym_statistics_common::clients::{packet_statistics::PacketStatisticsEvent, C
 use nym_task::connections::{
     ConnectionCommand, ConnectionCommandReceiver, ConnectionId, LaneQueueLengths, TransmissionLane,
 };
+use nym_task::TaskClient;
 use rand::{CryptoRng, Rng};
 use std::pin::Pin;
 use std::sync::Arc;
@@ -117,6 +118,8 @@ where
 
     /// Channel used for sending metrics events (specifically `PacketStatistics` events) to the metrics tracker.
     stats_tx: ClientStatsSender,
+
+    task_client: TaskClient,
 }
 
 #[derive(Debug)]
@@ -176,6 +179,7 @@ where
         lane_queue_lengths: LaneQueueLengths,
         client_connection_rx: ConnectionCommandReceiver,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         OutQueueControl {
             config,
@@ -190,6 +194,7 @@ where
             client_connection_rx,
             lane_queue_lengths,
             stats_tx,
+            task_client,
         }
     }
 
@@ -222,7 +227,6 @@ where
     async fn on_message(
         &mut self,
         next_message: StreamMessage,
-        task_client: &mut nym_task::TaskClient,
     ) {
         trace!("created new message");
 
@@ -277,7 +281,7 @@ where
         };
 
         if let Err(err) = self.mix_tx.send(vec![next_message]).await {
-            if !task_client.is_shutdown_poll() {
+            if !self.task_client.is_shutdown_poll() {
                 log::error!("Failed to send: {err}");
             }
         } else {
@@ -512,7 +516,7 @@ where
     }
 
     #[cfg(not(target_arch = "wasm32"))]
-    fn log_status(&self, shutdown: &mut nym_task::TaskClient) {
+    fn log_status(&self, shutdown: &mut TaskClient) {
         use crate::error::ClientCoreStatusMessage;
 
         let packets = self.transmission_buffer.total_size();
@@ -543,9 +547,11 @@ where
         }
     }
 
-    pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(super) async fn run(&mut self) {
         debug!("Started OutQueueControl with graceful shutdown support");
 
+        let mut shutdown = self.task_client.fork("select");
+
         #[cfg(not(target_arch = "wasm32"))]
         {
             let mut status_timer = tokio::time::interval(Duration::from_secs(5));
@@ -561,7 +567,7 @@ where
                         self.log_status(&mut shutdown);
                     }
                     next_message = self.next() => if let Some(next_message) = next_message {
-                        self.on_message(next_message, &mut shutdown).await;
+                        self.on_message(next_message).await;
                     } else {
                         log::trace!("OutQueueControl: Stopping since channel closed");
                         break;

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

@@ -12,6 +12,7 @@ use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag;
 use nym_sphinx::anonymous_replies::ReplySurb;
 use nym_sphinx::chunking::fragment::{Fragment, FragmentIdentifier};
 use nym_task::connections::{ConnectionId, TransmissionLane};
+use nym_task::TaskClient;
 use rand::{CryptoRng, Rng};
 use std::cmp::{max, min};
 use std::collections::btree_map::Entry;
@@ -68,6 +69,9 @@ pub struct ReplyController<R> {
 
     message_handler: MessageHandler<R>,
     full_reply_storage: CombinedReplyStorage,
+
+    // Listen for shutdown signals
+    task_client: TaskClient,
 }
 
 impl<R> ReplyController<R>
@@ -79,6 +83,7 @@ where
         message_handler: MessageHandler<R>,
         full_reply_storage: CombinedReplyStorage,
         request_receiver: ReplyControllerReceiver,
+        task_client: TaskClient,
     ) -> Self {
         ReplyController {
             config,
@@ -87,6 +92,7 @@ where
             pending_retransmissions: HashMap::new(),
             message_handler,
             full_reply_storage,
+            task_client,
         }
     }
 
@@ -846,9 +852,11 @@ where
     //     todo!()
     // }
 
-    pub(crate) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(crate) async fn run(&mut self) {
         debug!("Started ReplyController with graceful shutdown support");
 
+        let mut shutdown = self.task_client.fork("select");
+
         let polling_rate = Duration::from_secs(5);
         let mut stale_inspection = new_interval_stream(polling_rate);
 

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

@@ -74,7 +74,7 @@ impl StatisticsControl {
                 stats_rx,
                 report_tx,
                 reporting_config,
-                task_client: task_client.fork("statistics-control"),
+                task_client: task_client.fork("statistics_control"),
             },
             ClientStatsSender::new(Some(stats_tx), task_client),
         )

common/task/src/manager.rs

@@ -299,6 +299,8 @@ impl Clone for TaskClient {
             None
         };
 
+        log::debug!("Cloned task client: {name:?}");
+
         TaskClient {
             name,
             shutdown: AtomicBool::new(self.shutdown.load(Ordering::Relaxed)),
@@ -344,6 +346,7 @@ impl TaskClient {
             format!("unknown-{suffix}")
         };
 
+        log::debug!("Forked task client: {child_name}");
         child.name = Some(child_name);
         child
     }
@@ -377,6 +380,7 @@ impl TaskClient {
         } else {
             format!("unknown-{suffix}")
         };
+        log::debug!("Renamed task client: {name}");
         self.named(name)
     }