Merge pull request #5435 from nymtech/jon/task-all-stop

Remove all recv_with_delay and add shutdown condition to loops in client-core

clients/native/src/client/mod.rs

@@ -49,7 +49,7 @@ impl SocketClient {
         client_output: ClientOutput,
         client_state: ClientState,
         self_address: &Recipient,
-        shutdown: nym_task::TaskClient,
+        task_client: nym_task::TaskClient,
         packet_type: PacketType,
     ) {
         info!("Starting websocket listener...");
@@ -77,10 +77,15 @@ impl SocketClient {
             shared_lane_queue_lengths,
             reply_controller_sender,
             Some(packet_type),
+            task_client.fork("websocket_handler"),
         );
 
-        websocket::Listener::new(config.socket.host, config.socket.listening_port)
-            .start(websocket_handler, shutdown);
+        websocket::Listener::new(
+            config.socket.host,
+            config.socket.listening_port,
+            task_client.with_suffix("websocket_listener"),
+        )
+        .start(websocket_handler);
     }
 
     /// blocking version of `start_socket` method. Will run forever (or until SIGINT is sent)

clients/native/src/websocket/handler.rs

@@ -19,6 +19,7 @@ use nym_sphinx::receiver::ReconstructedMessage;
 use nym_task::connections::{
     ConnectionCommand, ConnectionCommandSender, ConnectionId, LaneQueueLengths, TransmissionLane,
 };
+use nym_task::TaskClient;
 use std::time::Duration;
 use tokio::net::TcpStream;
 use tokio::time::Instant;
@@ -43,9 +44,11 @@ pub(crate) struct HandlerBuilder {
     lane_queue_lengths: LaneQueueLengths,
     reply_controller_sender: ReplyControllerSender,
     packet_type: Option<PacketType>,
+    task_client: TaskClient,
 }
 
 impl HandlerBuilder {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         msg_input: InputMessageSender,
         client_connection_tx: ConnectionCommandSender,
@@ -54,6 +57,7 @@ impl HandlerBuilder {
         lane_queue_lengths: LaneQueueLengths,
         reply_controller_sender: ReplyControllerSender,
         packet_type: Option<PacketType>,
+        task_client: TaskClient,
     ) -> Self {
         Self {
             msg_input,
@@ -63,11 +67,14 @@ impl HandlerBuilder {
             lane_queue_lengths,
             reply_controller_sender,
             packet_type,
+            task_client,
         }
     }
 
     // TODO: make sure we only ever have one active handler
     pub fn create_active_handler(&self) -> Handler {
+        let mut task_client = self.task_client.fork("active_handler");
+        task_client.disarm();
         Handler {
             msg_input: self.msg_input.clone(),
             client_connection_tx: self.client_connection_tx.clone(),
@@ -78,6 +85,7 @@ impl HandlerBuilder {
             lane_queue_lengths: self.lane_queue_lengths.clone(),
             reply_controller_sender: self.reply_controller_sender.clone(),
             packet_type: self.packet_type,
+            task_client,
         }
     }
 }
@@ -92,16 +100,18 @@ pub(crate) struct Handler {
     lane_queue_lengths: LaneQueueLengths,
     reply_controller_sender: ReplyControllerSender,
     packet_type: Option<PacketType>,
+    task_client: TaskClient,
 }
 
 impl Drop for Handler {
     fn drop(&mut self) {
-        if self
+        if let Err(err) = self
             .buffer_requester
             .unbounded_send(ReceivedBufferMessage::ReceiverDisconnect)
-            .is_err()
         {
-            error!("we failed to disconnect the receiver from the buffer! presumably the shutdown procedure has been initiated!")
+            if !self.task_client.is_shutdown_poll() {
+                error!("failed to disconnect the receiver from the buffer: {err}");
+            }
         }
     }
 }
@@ -125,10 +135,23 @@ impl Handler {
         };
 
         // get the number of pending replies waiting for reply surbs
-        let reply_queue_length = self
+        let reply_queue_length = match self
             .reply_controller_sender
             .get_lane_queue_length(connection_id)
-            .await;
+            .await
+        {
+            Ok(length) => length,
+            Err(err) => {
+                if !self.task_client.is_shutdown_poll() {
+                    error!(
+                        "Failed to get reply queue length for connection {connection_id}: {err}"
+                    );
+                }
+                // We're just going to assume that the queue is empty, and I think that's okay
+                // during shutdown.
+                0
+            }
+        };
 
         let queue_length = base_length + reply_queue_length;
 
@@ -168,10 +191,11 @@ impl Handler {
 
         // the ack control is now responsible for chunking, etc.
         let input_msg = InputMessage::new_regular(recipient, message, lane, self.packet_type);
-        self.msg_input
-            .send(input_msg)
-            .await
-            .expect("InputMessageReceiver has stopped receiving!");
+        if let Err(err) = self.msg_input.send(input_msg).await {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send message to the input buffer: {err}");
+            }
+        }
 
         // Only reply back with a `LaneQueueLength` if the sender providided a connection id
         let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -200,10 +224,11 @@ impl Handler {
 
         let input_msg =
             InputMessage::new_anonymous(recipient, message, reply_surbs, lane, self.packet_type);
-        self.msg_input
-            .send(input_msg)
-            .await
-            .expect("InputMessageReceiver has stopped receiving!");
+        if let Err(err) = self.msg_input.send(input_msg).await {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send anonymous message to the input buffer: {err}");
+            }
+        }
 
         // Only reply back with a `LaneQueueLength` if the sender providided a connection id
         let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -227,10 +252,11 @@ impl Handler {
         });
 
         let input_msg = InputMessage::new_reply(recipient_tag, message, lane, self.packet_type);
-        self.msg_input
-            .send(input_msg)
-            .await
-            .expect("InputMessageReceiver has stopped receiving!");
+        if let Err(err) = self.msg_input.send(input_msg).await {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send reply message to the input buffer: {err}");
+            }
+        }
 
         // Only reply back with a `LaneQueueLength` if the sender providided a connection id
         let TransmissionLane::ConnectionId(connection_id) = lane else {
@@ -245,9 +271,14 @@ impl Handler {
     }
 
     fn handle_closed_connection(&self, connection_id: u64) -> Option<ServerResponse> {
-        self.client_connection_tx
+        if let Err(err) = self
+            .client_connection_tx
             .unbounded_send(ConnectionCommand::Close(connection_id))
-            .unwrap();
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send close connection command: {err}");
+            }
+        }
         None
     }
 
@@ -362,11 +393,10 @@ impl Handler {
         }
     }
 
-    async fn listen_for_requests(
-        &mut self,
-        mut msg_receiver: ReconstructedMessagesReceiver,
-        mut task_client: nym_task::TaskClient,
-    ) {
+    async fn listen_for_requests(&mut self, mut msg_receiver: ReconstructedMessagesReceiver) {
+        let mut task_client = self.task_client.fork("select");
+        task_client.disarm();
+
         while !task_client.is_shutdown() {
             tokio::select! {
                 // we can either get a client request from the websocket
@@ -415,15 +445,7 @@ impl Handler {
     }
 
     // consume self to make sure `drop` is called after this is done
-    pub(crate) async fn handle_connection(
-        mut self,
-        socket: TcpStream,
-        mut task_client: nym_task::TaskClient,
-    ) {
-        // We don't want a crash in the connection handler to trigger a shutdown of the whole
-        // process.
-        task_client.disarm();
-
+    pub(crate) async fn handle_connection(mut self, socket: TcpStream) {
         let ws_stream = match accept_async(socket).await {
             Ok(ws_stream) => ws_stream,
             Err(err) => {
@@ -436,14 +458,18 @@ impl Handler {
         let (reconstructed_sender, reconstructed_receiver) = mpsc::unbounded();
 
         // tell the buffer to start sending stuff to us
-        self.buffer_requester
-            .unbounded_send(ReceivedBufferMessage::ReceiverAnnounce(
-                reconstructed_sender,
-            ))
-            .expect("the buffer request failed!");
+        if let Err(err) =
+            self.buffer_requester
+                .unbounded_send(ReceivedBufferMessage::ReceiverAnnounce(
+                    reconstructed_sender,
+                ))
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("failed to announce the receiver to the buffer: {err}");
+            }
+        }
 
-        self.listen_for_requests(reconstructed_receiver, task_client)
-            .await;
+        self.listen_for_requests(reconstructed_receiver).await;
     }
 }
 

clients/native/src/websocket/listener.rs

@@ -3,6 +3,7 @@
 
 use super::handler::HandlerBuilder;
 use log::*;
+use nym_task::TaskClient;
 use std::net::IpAddr;
 use std::{net::SocketAddr, process, sync::Arc};
 use tokio::io::AsyncWriteExt;
@@ -22,21 +23,19 @@ impl State {
 pub(crate) struct Listener {
     address: SocketAddr,
     state: State,
+    task_client: TaskClient,
 }
 
 impl Listener {
-    pub(crate) fn new(host: IpAddr, port: u16) -> Self {
+    pub(crate) fn new(host: IpAddr, port: u16, task_client: TaskClient) -> Self {
         Listener {
             address: SocketAddr::new(host, port),
             state: State::AwaitingConnection,
+            task_client,
         }
     }
 
-    pub(crate) async fn run(
-        &mut self,
-        handler: HandlerBuilder,
-        mut task_client: nym_task::TaskClient,
-    ) {
+    pub(crate) async fn run(&mut self, handler: HandlerBuilder) {
         let tcp_listener = match tokio::net::TcpListener::bind(self.address).await {
             Ok(listener) => listener,
             Err(err) => {
@@ -47,11 +46,11 @@ impl Listener {
 
         let notify = Arc::new(Notify::new());
 
-        loop {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 // When the handler finishes we check if shutdown is signalled
                 _ = notify.notified() => {
-                    if task_client.is_shutdown() {
+                    if self.task_client.is_shutdown() {
                         log::trace!("Websocket listener: detected shutdown after connection closed");
                         break;
                     }
@@ -60,7 +59,7 @@ impl Listener {
                 }
                 // ... but when there is no connected client at the time of shutdown being
                 // signalled, we handle it here.
-                _ = task_client.recv() => {
+                _ = self.task_client.recv() => {
                     if !self.state.is_connected() {
                         log::trace!("Not connected: shutting down");
                         break;
@@ -88,9 +87,8 @@ impl Listener {
                                 // hanging because the executor doesn't come back here
                                 let notify_clone = Arc::clone(&notify);
                                 let fresh_handler = handler.create_active_handler();
-                                let task_client_handler = task_client.clone();
                                 tokio::spawn(async move {
-                                    fresh_handler.handle_connection(socket, task_client_handler).await;
+                                    fresh_handler.handle_connection(socket).await;
                                     notify_clone.notify_one();
                                 });
                                 self.state = State::Connected;
@@ -104,13 +102,9 @@ impl Listener {
         log::debug!("Websocket listener: Exiting");
     }
 
-    pub(crate) fn start(
-        mut self,
-        handler: HandlerBuilder,
-        shutdown: nym_task::TaskClient,
-    ) -> JoinHandle<()> {
+    pub(crate) fn start(mut self, handler: HandlerBuilder) -> JoinHandle<()> {
         info!("Running websocket on {:?}", self.address.to_string());
 
-        tokio::spawn(async move { self.run(handler, shutdown).await })
+        tokio::spawn(async move { self.run(handler).await })
     }
 }

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
@@ -382,8 +384,9 @@ where
                 reply_key_storage,
                 reply_controller_sender,
                 metrics_reporter,
+                shutdown,
             );
-        controller.start_with_shutdown(shutdown)
+        controller.start()
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -572,15 +575,22 @@ where
         topology_accessor: TopologyAccessor,
         local_gateway: NodeIdentity,
         wait_for_gateway: bool,
-        mut shutdown: TaskClient,
+        mut task_client: TaskClient,
     ) -> Result<(), ClientCoreError> {
         let topology_refresher_config =
             TopologyRefresherConfig::new(topology_config.topology_refresh_rate);
 
+        if topology_config.disable_refreshing {
+            // if we're not spawning the refresher, don't cause shutdown immediately
+            info!("The background topology refesher is not going to be started");
+            task_client.disarm();
+        }
+
         let mut topology_refresher = TopologyRefresher::new(
             topology_refresher_config,
             topology_accessor,
             topology_provider,
+            task_client,
         );
         // before returning, block entire runtime to refresh the current network view so that any
         // components depending on topology would see a non-empty view
@@ -621,15 +631,11 @@ where
             }
         }
 
-        if topology_config.disable_refreshing {
-            // if we're not spawning the refresher, don't cause shutdown immediately
-            info!("The topology refesher is not going to be started");
-            shutdown.disarm();
-        } else {
+        if !topology_config.disable_refreshing {
             // 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...");
-            topology_refresher.start_with_shutdown(shutdown);
+            topology_refresher.start();
         }
 
         Ok(())
@@ -640,17 +646,17 @@ where
         user_agent: Option<UserAgent>,
         client_stats_id: String,
         input_sender: Sender<InputMessage>,
-        shutdown: TaskClient,
+        task_client: TaskClient,
     ) -> ClientStatsSender {
         info!("Starting statistics control...");
-        StatisticsControl::create_and_start_with_shutdown(
+        StatisticsControl::create_and_start(
             config.debug.stats_reporting,
             user_agent
                 .map(|u| u.application)
                 .unwrap_or("unknown".to_string()),
             client_stats_id,
             input_sender.clone(),
-            shutdown.with_suffix("controller"),
+            task_client,
         )
     }
 
@@ -660,8 +666,8 @@ where
     ) -> (BatchMixMessageSender, ClientRequestSender) {
         info!("Starting mix traffic controller...");
         let (mix_traffic_controller, mix_tx, client_tx) =
-            MixTrafficController::new(gateway_transceiver);
-        mix_traffic_controller.start_with_shutdown(shutdown);
+            MixTrafficController::new(gateway_transceiver, shutdown);
+        mix_traffic_controller.start();
         (mix_tx, client_tx)
     }
 

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,
         }
     }
 
@@ -175,7 +180,7 @@ impl LoopCoverTrafficStream<OsRng> {
             }
         };
 
-        let cover_message = generate_loop_cover_packet(
+        let cover_message = match generate_loop_cover_packet(
             &mut self.rng,
             topology_ref,
             &self.ack_key,
@@ -184,8 +189,15 @@ impl LoopCoverTrafficStream<OsRng> {
             self.cover_traffic.loop_cover_traffic_average_delay,
             cover_traffic_packet_size,
             self.packet_type,
-        )
-        .expect("Somehow failed to generate a loop cover message with a valid topology");
+        ) {
+            Ok(cover_message) => cover_message,
+            Err(err) => {
+                warn!(
+                    "Somehow failed to generate a loop cover message with a valid topology: {err}"
+                );
+                return;
+            }
+        };
 
         if let Err(err) = self.mix_tx.try_send(vec![cover_message]) {
             match err {
@@ -217,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
@@ -231,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/mix_traffic/mod.rs

@@ -7,6 +7,7 @@ use crate::spawn_future;
 use log::*;
 use nym_gateway_requests::ClientRequest;
 use nym_sphinx::forwarding::packet::MixPacket;
+use nym_task::TaskClient;
 use transceiver::ErasedGatewayError;
 
 pub type BatchMixMessageSender = tokio::sync::mpsc::Sender<Vec<MixPacket>>;
@@ -32,11 +33,14 @@ pub struct MixTrafficController {
     // TODO: this is temporary work-around.
     // in long run `gateway_client` will be moved away from `MixTrafficController` anyway.
     consecutive_gateway_failure_count: usize,
+
+    task_client: TaskClient,
 }
 
 impl MixTrafficController {
     pub fn new<T>(
         gateway_transceiver: T,
+        task_client: TaskClient,
     ) -> (
         MixTrafficController,
         BatchMixMessageSender,
@@ -56,6 +60,7 @@ impl MixTrafficController {
                 mix_rx: message_receiver,
                 client_rx: client_receiver,
                 consecutive_gateway_failure_count: 0,
+                task_client,
             },
             message_sender,
             client_sender,
@@ -64,6 +69,7 @@ impl MixTrafficController {
 
     pub fn new_dynamic(
         gateway_transceiver: Box<dyn GatewayTransceiver + Send>,
+        task_client: TaskClient,
     ) -> (
         MixTrafficController,
         BatchMixMessageSender,
@@ -78,6 +84,7 @@ impl MixTrafficController {
                 mix_rx: message_receiver,
                 client_rx: client_receiver,
                 consecutive_gateway_failure_count: 0,
+                task_client,
             },
             message_sender,
             client_sender,
@@ -109,11 +116,11 @@ impl MixTrafficController {
         result
     }
 
-    pub fn start_with_shutdown(mut self, mut shutdown: nym_task::TaskClient) {
+    pub fn start(mut self) {
         spawn_future(async move {
             debug!("Started MixTrafficController with graceful shutdown support");
 
-            loop {
+            while !self.task_client.is_shutdown() {
                 tokio::select! {
                     mix_packets = self.mix_rx.recv() => match mix_packets {
                         Some(mix_packets) => {
@@ -125,7 +132,7 @@ impl MixTrafficController {
                                     error!("Failed to send sphinx packet to the gateway {MAX_FAILURE_COUNT} times in a row - assuming the gateway is dead");
                                     // Do we need to handle the embedded mixnet client case
                                     // separately?
-                                    shutdown.send_we_stopped(Box::new(ClientCoreError::GatewayFailedToForwardMessages));
+                                    self.task_client.send_we_stopped(Box::new(ClientCoreError::GatewayFailedToForwardMessages));
                                     break;
                                 }
                             }
@@ -146,13 +153,13 @@ impl MixTrafficController {
                             log::trace!("MixTrafficController, client request channel closed");
                         }
                     },
-                    _ = shutdown.recv_with_delay() => {
+                    _ = self.task_client.recv() => {
                         log::trace!("MixTrafficController: Received shutdown");
                         break;
                     }
                 }
             }
-            shutdown.recv_timeout().await;
+            self.task_client.recv_timeout().await;
 
             log::debug!("MixTrafficController: Exiting");
         });

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

@@ -11,6 +11,7 @@ use nym_sphinx::{
     acknowledgements::{identifier::recover_identifier, AckKey},
     chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID},
 };
+use nym_task::TaskClient;
 use std::sync::Arc;
 
 /// Module responsible for listening for any data resembling acknowledgements from the network
@@ -20,6 +21,7 @@ pub(super) struct AcknowledgementListener {
     ack_receiver: AcknowledgementReceiver,
     action_sender: AckActionSender,
     stats_tx: ClientStatsSender,
+    task_client: TaskClient,
 }
 
 impl AcknowledgementListener {
@@ -28,12 +30,14 @@ impl AcknowledgementListener {
         ack_receiver: AcknowledgementReceiver,
         action_sender: AckActionSender,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         AcknowledgementListener {
             ack_key,
             ack_receiver,
             action_sender,
             stats_tx,
+            task_client,
         }
     }
 
@@ -64,9 +68,14 @@ impl AcknowledgementListener {
         trace!("Received {} from the mix network", frag_id);
         self.stats_tx
             .report(PacketStatisticsEvent::RealAckReceived(ack_content.len()).into());
-        self.action_sender
+        if let Err(err) = self
+            .action_sender
             .unbounded_send(Action::new_remove(frag_id))
-            .unwrap();
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send remove action to action controller: {err}");
+            }
+        }
     }
 
     async fn handle_ack_receiver_item(&mut self, item: Vec<Vec<u8>>) {
@@ -76,10 +85,10 @@ impl AcknowledgementListener {
         }
     }
 
-    pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(super) async fn run(&mut self) {
         debug!("Started AcknowledgementListener with graceful shutdown support");
 
-        while !shutdown.is_shutdown() {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 acks = self.ack_receiver.next() => match acks {
                     Some(acks) => self.handle_ack_receiver_item(acks).await,
@@ -88,12 +97,12 @@ impl AcknowledgementListener {
                         break;
                     }
                 },
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv() => {
                     log::trace!("AcknowledgementListener: Received shutdown");
                 }
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv_timeout().await;
         log::debug!("AcknowledgementListener: Exiting");
     }
 }

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

@@ -9,6 +9,7 @@ use log::*;
 use nym_nonexhaustive_delayqueue::{Expired, NonExhaustiveDelayQueue, QueueKey};
 use nym_sphinx::chunking::fragment::FragmentIdentifier;
 use nym_sphinx::Delay as SphinxDelay;
+use nym_task::TaskClient;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
@@ -101,6 +102,8 @@ pub(super) struct ActionController {
 
     /// Channel for notifying `RetransmissionRequestListener` about expired acknowledgements.
     retransmission_sender: RetransmissionRequestSender,
+
+    task_client: TaskClient,
 }
 
 impl ActionController {
@@ -108,6 +111,7 @@ impl ActionController {
         config: Config,
         retransmission_sender: RetransmissionRequestSender,
         incoming_actions: AckActionReceiver,
+        task_client: TaskClient,
     ) -> Self {
         ActionController {
             config,
@@ -115,6 +119,7 @@ impl ActionController {
             pending_acks_timers: NonExhaustiveDelayQueue::new(),
             incoming_actions,
             retransmission_sender,
+            task_client,
         }
     }
 
@@ -216,11 +221,7 @@ impl ActionController {
     }
 
     // note: when the entry expires it's automatically removed from pending_acks_timers
-    fn handle_expired_ack_timer(
-        &mut self,
-        expired_ack: Expired<FragmentIdentifier>,
-        task_client: &mut nym_task::TaskClient,
-    ) {
+    fn handle_expired_ack_timer(&mut self, expired_ack: Expired<FragmentIdentifier>) {
         // I'm honestly not sure how to handle it, because getting it means other things in our
         // system are already misbehaving. If we ever see this panic, then I guess we should worry
         // about it. Perhaps just reschedule it at later point?
@@ -238,15 +239,13 @@ impl ActionController {
             // downgrading an arc and then upgrading vs cloning is difference of 30ns vs 15ns
             // so it's literally a NO difference while it might prevent us from unnecessarily
             // resending data (in maybe 1 in 1 million cases, but it's something)
-            if self
+            if let Err(err) = self
                 .retransmission_sender
                 .unbounded_send(Arc::downgrade(pending_ack_data))
-                .is_err()
             {
-                assert!(
-                    task_client.is_shutdown_poll(),
-                    "Failed to send pending ack for retransmission"
-                );
+                if !self.task_client.is_shutdown_poll() {
+                    log::error!("Failed to send pending ack for retransmission: {err}");
+                }
             }
         } else {
             // this shouldn't cause any issues but shouldn't have happened to begin with!
@@ -265,10 +264,10 @@ impl ActionController {
         }
     }
 
-    pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(super) async fn run(&mut self) {
         debug!("Started ActionController with graceful shutdown support");
 
-        loop {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 action = self.incoming_actions.next() => match action {
                     Some(action) => self.process_action(action),
@@ -280,19 +279,19 @@ impl ActionController {
                     }
                 },
                 expired_ack = self.pending_acks_timers.next() => match expired_ack {
-                    Some(expired_ack) => self.handle_expired_ack_timer(expired_ack, &mut shutdown),
+                    Some(expired_ack) => self.handle_expired_ack_timer(expired_ack),
                     None => {
                         log::trace!("ActionController: Stopping since ack channel closed");
                         break;
                     }
                 },
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv() => {
                     log::trace!("ActionController: Received shutdown");
                     break;
                 }
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv_timeout().await;
         log::debug!("ActionController: Exiting");
     }
 }

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

@@ -11,6 +11,7 @@ use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag;
 use nym_sphinx::forwarding::packet::MixPacket;
 use nym_sphinx::params::PacketType;
 use nym_task::connections::TransmissionLane;
+use nym_task::TaskClient;
 use rand::{CryptoRng, Rng};
 
 /// Module responsible for dealing with the received messages: splitting them, creating acknowledgements,
@@ -23,6 +24,7 @@ where
     input_receiver: InputMessageReceiver,
     message_handler: MessageHandler<R>,
     reply_controller_sender: ReplyControllerSender,
+    task_client: TaskClient,
 }
 
 impl<R> InputMessageListener<R>
@@ -36,11 +38,13 @@ where
         input_receiver: InputMessageReceiver,
         message_handler: MessageHandler<R>,
         reply_controller_sender: ReplyControllerSender,
+        task_client: TaskClient,
     ) -> Self {
         InputMessageListener {
             input_receiver,
             message_handler,
             reply_controller_sender,
+            task_client,
         }
     }
 
@@ -63,8 +67,14 @@ where
         lane: TransmissionLane,
     ) {
         // offload reply handling to the dedicated task
-        self.reply_controller_sender
+        if let Err(err) = self
+            .reply_controller_sender
             .send_reply(recipient_tag, data, lane)
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("failed to send a reply - {err}");
+            }
+        }
     }
 
     async fn handle_plain_message(
@@ -164,10 +174,10 @@ where
         };
     }
 
-    pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(super) async fn run(&mut self) {
         debug!("Started InputMessageListener with graceful shutdown support");
 
-        while !shutdown.is_shutdown() {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 input_msg = self.input_receiver.recv() => match input_msg {
                     Some(input_msg) => {
@@ -178,12 +188,12 @@ where
                         break;
                     }
                 },
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv() => {
                     log::trace!("InputMessageListener: Received shutdown");
                 }
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv_timeout().await;
         log::debug!("InputMessageListener: Exiting");
     }
 }

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},
@@ -66,7 +67,7 @@ pub(crate) enum PacketDestination {
 
 /// Structure representing a data `Fragment` that is on-route to the specified `Recipient`
 #[derive(Debug)]
-pub(crate) struct PendingAcknowledgement {
+pub struct PendingAcknowledgement {
     message_chunk: Fragment,
     delay: SphinxDelay,
     destination: PacketDestination,
@@ -216,6 +217,7 @@ where
         message_handler: MessageHandler<R>,
         reply_controller_sender: ReplyControllerSender,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         let (retransmission_tx, retransmission_rx) = mpsc::unbounded();
 
@@ -225,6 +227,7 @@ where
             action_config,
             retransmission_tx,
             connectors.ack_action_receiver,
+            task_client.fork("action_controller"),
         );
 
         // will listen for any acks coming from the network
@@ -233,6 +236,7 @@ where
             connectors.ack_receiver,
             connectors.ack_action_sender.clone(),
             stats_tx,
+            task_client.fork("acknowledgement_listener"),
         );
 
         // will listen for any new messages from the client
@@ -240,6 +244,7 @@ where
             connectors.input_receiver,
             message_handler.clone(),
             reply_controller_sender.clone(),
+            task_client.fork("input_message_listener"),
         );
 
         // will listen for any ack timeouts and trigger retransmission
@@ -249,12 +254,16 @@ where
             message_handler,
             retransmission_rx,
             reply_controller_sender,
+            task_client.fork("retransmission_request_listener"),
         );
 
         // will listen for events indicating the packet was sent through the network so that
         // the retransmission timer should be started.
-        let sent_notification_listener =
-            SentNotificationListener::new(connectors.sent_notifier, connectors.ack_action_sender);
+        let sent_notification_listener = SentNotificationListener::new(
+            connectors.sent_notifier,
+            connectors.ack_action_sender,
+            task_client.with_suffix("sent_notification_listener"),
+        );
 
         AcknowledgementController {
             acknowledgement_listener,
@@ -265,53 +274,35 @@ where
         }
     }
 
-    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");
         spawn_future(async move {
-            acknowledgement_listener
-                .run_with_shutdown(shutdown_handle)
-                .await;
+            acknowledgement_listener.run().await;
             debug!("The acknowledgement listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("input_message_listener");
         spawn_future(async move {
-            input_message_listener
-                .run_with_shutdown(shutdown_handle)
-                .await;
+            input_message_listener.run().await;
             debug!("The input listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("retransmission_request_listener");
         spawn_future(async move {
-            retransmission_request_listener
-                .run_with_shutdown(shutdown_handle, packet_type)
-                .await;
+            retransmission_request_listener.run(packet_type).await;
             debug!("The retransmission request listener has finished execution!");
         });
 
-        let shutdown_handle = shutdown.fork("sent_notification_listener");
         spawn_future(async move {
-            sent_notification_listener
-                .run_with_shutdown(shutdown_handle)
-                .await;
+            sent_notification_listener.run().await;
             debug!("The sent notification listener has finished execution!");
         });
 
         spawn_future(async move {
-            action_controller
-                .run_with_shutdown(shutdown.with_suffix("action_controller"))
-                .await;
+            action_controller.run().await;
             debug!("The controller has finished execution!");
         });
     }

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

@@ -14,7 +14,7 @@ use log::*;
 use nym_sphinx::chunking::fragment::Fragment;
 use nym_sphinx::preparer::PreparedFragment;
 use nym_sphinx::{addressing::clients::Recipient, params::PacketType};
-use nym_task::connections::TransmissionLane;
+use nym_task::{connections::TransmissionLane, TaskClient};
 use rand::{CryptoRng, Rng};
 use std::sync::{Arc, Weak};
 
@@ -25,6 +25,7 @@ pub(super) struct RetransmissionRequestListener<R> {
     message_handler: MessageHandler<R>,
     request_receiver: RetransmissionRequestReceiver,
     reply_controller_sender: ReplyControllerSender,
+    task_client: TaskClient,
 }
 
 impl<R> RetransmissionRequestListener<R>
@@ -37,6 +38,7 @@ where
         message_handler: MessageHandler<R>,
         request_receiver: RetransmissionRequestReceiver,
         reply_controller_sender: ReplyControllerSender,
+        task_client: TaskClient,
     ) -> Self {
         RetransmissionRequestListener {
             maximum_retransmissions,
@@ -44,6 +46,7 @@ where
             message_handler,
             request_receiver,
             reply_controller_sender,
+            task_client,
         }
     }
 
@@ -79,9 +82,12 @@ where
         if let Some(limit) = self.maximum_retransmissions {
             if timed_out_ack.retransmissions >= limit {
                 warn!("reached maximum number of allowed retransmissions for the packet");
-                self.action_sender
+                if let Err(err) = self
+                    .action_sender
                     .unbounded_send(Action::new_remove(frag_id))
-                    .unwrap();
+                {
+                    error!("Failed to send remove action to the controller: {err}");
+                }
                 return;
             }
         }
@@ -93,11 +99,16 @@ where
             } => {
                 // if this is retransmission for reply, offload it to the dedicated task
                 // that deals with all the surbs
-                return self.reply_controller_sender.send_retransmission_data(
+                if let Err(err) = self.reply_controller_sender.send_retransmission_data(
                     *recipient_tag,
                     weak_timed_out_ack,
                     *extra_surb_request,
-                );
+                ) {
+                    if !self.task_client.is_shutdown_poll() {
+                        error!("Failed to send retransmission data to the reply controller: {err}");
+                    }
+                }
+                return;
             }
             PacketDestination::KnownRecipient(recipient) => {
                 self.prepare_normal_retransmission_chunk(
@@ -114,9 +125,12 @@ where
             Err(err) => {
                 warn!("Could not retransmit the packet - {err}");
                 // we NEED to start timer here otherwise we will have this guy permanently stuck in memory
-                self.action_sender
+                if let Err(err) = self
+                    .action_sender
                     .unbounded_send(Action::new_start_timer(frag_id))
-                    .unwrap();
+                {
+                    error!("Failed to send start timer action to the controller: {err}");
+                }
                 return;
             }
         };
@@ -141,9 +155,14 @@ where
         // is sent to the `OutQueueControl` and has gone through its internal queue
         // with the additional poisson delay.
         // And since Actions are executed in order `UpdateTimer` will HAVE TO be executed before `StartTimer`
-        self.action_sender
+        if let Err(err) = self
+            .action_sender
             .unbounded_send(Action::new_update_pending_ack(frag_id, new_delay))
-            .unwrap();
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send update pending ack action to the controller: {err}");
+            }
+        }
 
         // send to `OutQueueControl` to eventually send to the mix network
         self.message_handler
@@ -157,14 +176,10 @@ where
             .await
     }
 
-    pub(super) async fn run_with_shutdown(
-        &mut self,
-        mut shutdown: nym_task::TaskClient,
-        packet_type: PacketType,
-    ) {
+    pub(super) async fn run(&mut self, packet_type: PacketType) {
         debug!("Started RetransmissionRequestListener with graceful shutdown support");
 
-        while !shutdown.is_shutdown() {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 timed_out_ack = self.request_receiver.next() => match timed_out_ack {
                     Some(timed_out_ack) => self.on_retransmission_request(timed_out_ack, packet_type).await,
@@ -173,12 +188,12 @@ where
                         break;
                     }
                 },
-                _ = shutdown.recv() => {
+                _ = self.task_client.recv() => {
                     log::trace!("RetransmissionRequestListener: Received shutdown");
                 }
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv_timeout().await;
         log::debug!("RetransmissionRequestListener: Exiting");
     }
 }

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

@@ -6,6 +6,7 @@ use super::SentPacketNotificationReceiver;
 use futures::StreamExt;
 use log::*;
 use nym_sphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID};
+use nym_task::TaskClient;
 
 /// Module responsible for starting up retransmission timers.
 /// It is required because when we send our packet to the `real traffic stream` controlled
@@ -14,16 +15,19 @@ use nym_sphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID};
 pub(super) struct SentNotificationListener {
     sent_notifier: SentPacketNotificationReceiver,
     action_sender: AckActionSender,
+    task_client: TaskClient,
 }
 
 impl SentNotificationListener {
     pub(super) fn new(
         sent_notifier: SentPacketNotificationReceiver,
         action_sender: AckActionSender,
+        task_client: TaskClient,
     ) -> Self {
         SentNotificationListener {
             sent_notifier,
             action_sender,
+            task_client,
         }
     }
 
@@ -32,15 +36,20 @@ impl SentNotificationListener {
             trace!("sent off a cover message - no need to start retransmission timer!");
             return;
         }
-        self.action_sender
+        if let Err(err) = self
+            .action_sender
             .unbounded_send(Action::new_start_timer(frag_id))
-            .unwrap();
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send start timer action to action controller: {err}");
+            }
+        }
     }
 
-    pub(super) async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    pub(super) async fn run(&mut self) {
         debug!("Started SentNotificationListener with graceful shutdown support");
 
-        loop {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 frag_id = self.sent_notifier.next() => match frag_id {
                     Some(frag_id) => {
@@ -51,13 +60,13 @@ impl SentNotificationListener {
                         break;
                     }
                 },
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv() => {
                     log::trace!("SentNotificationListener: Received shutdown");
                     break;
                 }
             }
         }
-        assert!(shutdown.is_shutdown_poll());
+        assert!(self.task_client.is_shutdown_poll());
         log::debug!("SentNotificationListener: Exiting");
     }
 }

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

@@ -19,6 +19,7 @@ use nym_sphinx::params::{PacketSize, PacketType};
 use nym_sphinx::preparer::{MessagePreparer, PreparedFragment};
 use nym_sphinx::Delay;
 use nym_task::connections::TransmissionLane;
+use nym_task::TaskClient;
 use nym_topology::{NymRouteProvider, NymTopologyError};
 use rand::{CryptoRng, Rng};
 use std::collections::HashMap;
@@ -149,12 +150,14 @@ pub(crate) struct MessageHandler<R> {
     topology_access: TopologyAccessor,
     reply_key_storage: SentReplyKeys,
     tag_storage: UsedSenderTags,
+    task_client: TaskClient,
 }
 
 impl<R> MessageHandler<R>
 where
     R: CryptoRng + Rng,
 {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         config: Config,
         rng: R,
@@ -163,6 +166,7 @@ where
         topology_access: TopologyAccessor,
         reply_key_storage: SentReplyKeys,
         tag_storage: UsedSenderTags,
+        task_client: TaskClient,
     ) -> Self
     where
         R: Copy,
@@ -183,6 +187,7 @@ where
             topology_access,
             reply_key_storage,
             tag_storage,
+            task_client,
         }
     }
 
@@ -609,15 +614,25 @@ where
     }
 
     pub(crate) fn update_ack_delay(&self, id: FragmentIdentifier, new_delay: Delay) {
-        self.action_sender
+        if let Err(err) = self
+            .action_sender
             .unbounded_send(Action::UpdatePendingAck(id, new_delay))
-            .expect("action control task has died")
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send update action to the controller: {err}");
+            }
+        }
     }
 
     pub(crate) fn insert_pending_acks(&self, pending_acks: Vec<PendingAcknowledgement>) {
-        self.action_sender
+        if let Err(err) = self
+            .action_sender
             .unbounded_send(Action::new_insert(pending_acks))
-            .expect("action control task has died")
+        {
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to send insert action to the controller: {err}");
+            }
+        }
     }
 
     // tells real message sender (with the poisson timer) to send this to the mix network
@@ -631,9 +646,9 @@ where
             .send((messages, transmission_lane))
             .await
         {
-            error!(
-                "Failed to forward messages to the real message sender (OutQueueControl): {err}"
-            );
+            if !self.task_client.is_shutdown_poll() {
+                error!("Failed to forward messages to the real message sender: {err}");
+            }
         }
     }
 }

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;
 
@@ -177,6 +179,7 @@ impl RealMessagesController<OsRng> {
             topology_access.clone(),
             reply_storage.key_storage(),
             reply_storage.tags_storage(),
+            task_client.fork("message_handler"),
         );
 
         let ack_control = AcknowledgementController::new(
@@ -186,6 +189,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 +197,7 @@ impl RealMessagesController<OsRng> {
             message_handler,
             reply_storage,
             reply_controller_receiver,
+            task_client.fork("reply_controller"),
         );
 
         let out_queue_control = OutQueueControl::new(
@@ -205,6 +210,7 @@ impl RealMessagesController<OsRng> {
             lane_queue_lengths,
             client_connection_rx,
             stats_tx,
+            task_client.with_suffix("out_queue_control"),
         );
 
         RealMessagesController {
@@ -214,22 +220,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,
         }
     }
 
@@ -198,7 +203,9 @@ where
         // queues and client load rather than the required delay. So realistically we can treat
         // whatever is about to happen as negligible additional delay.
         trace!("{} is about to get sent to the mixnet", frag_id);
-        self.sent_notifier.unbounded_send(frag_id).unwrap();
+        if let Err(err) = self.sent_notifier.unbounded_send(frag_id) {
+            error!("Failed to notify about sent message: {err}");
+        }
     }
 
     fn loop_cover_message_size(&mut self) -> PacketSize {
@@ -271,7 +278,9 @@ where
         };
 
         if let Err(err) = self.mix_tx.send(vec![next_message]).await {
-            log::error!("Failed to send: {err}");
+            if !self.task_client.is_shutdown_poll() {
+                log::error!("Failed to send: {err}");
+            }
         } else {
             let event = if fragment_id.is_some() {
                 PacketStatisticsEvent::RealPacketSent(packet_size)
@@ -504,7 +513,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();
@@ -535,14 +544,16 @@ 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));
 
-            loop {
+            while !shutdown.is_shutdown() {
                 tokio::select! {
                     biased;
                     _ = shutdown.recv() => {

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

@@ -20,6 +20,7 @@ use nym_sphinx::message::{NymMessage, PlainMessage};
 use nym_sphinx::params::ReplySurbKeyDigestAlgorithm;
 use nym_sphinx::receiver::{MessageReceiver, MessageRecoveryError, ReconstructedMessage};
 use nym_statistics_common::clients::{packet_statistics::PacketStatisticsEvent, ClientStatsSender};
+use nym_task::TaskClient;
 use std::collections::HashSet;
 use std::sync::Arc;
 
@@ -152,6 +153,7 @@ struct ReceivedMessagesBuffer<R: MessageReceiver> {
     inner: Arc<Mutex<ReceivedMessagesBufferInner<R>>>,
     reply_key_storage: SentReplyKeys,
     reply_controller_sender: ReplyControllerSender,
+    task_client: TaskClient,
 }
 
 impl<R: MessageReceiver> ReceivedMessagesBuffer<R> {
@@ -160,6 +162,7 @@ impl<R: MessageReceiver> ReceivedMessagesBuffer<R> {
         reply_key_storage: SentReplyKeys,
         reply_controller_sender: ReplyControllerSender,
         stats_tx: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         ReceivedMessagesBuffer {
             inner: Arc::new(Mutex::new(ReceivedMessagesBufferInner {
@@ -172,6 +175,7 @@ impl<R: MessageReceiver> ReceivedMessagesBuffer<R> {
             })),
             reply_key_storage,
             reply_controller_sender,
+            task_client,
         }
     }
 
@@ -257,11 +261,15 @@ impl<R: MessageReceiver> ReceivedMessagesBuffer<R> {
                 }
             };
 
-            self.reply_controller_sender.send_additional_surbs(
+            if let Err(err) = self.reply_controller_sender.send_additional_surbs(
                 msg.sender_tag,
                 reply_surbs,
                 from_surb_request,
-            )
+            ) {
+                if !self.task_client.is_shutdown_poll() {
+                    error!("{err}");
+                }
+            }
         }
         reconstructed
     }
@@ -276,8 +284,14 @@ impl<R: MessageReceiver> ReceivedMessagesBuffer<R> {
                 ReplyMessageContent::Data { message } => reconstructed.push(message.into()),
                 ReplyMessageContent::SurbRequest { recipient, amount } => {
                     debug!("received request for {amount} additional reply SURBs from {recipient}");
-                    self.reply_controller_sender
-                        .send_additional_surbs_request(*recipient, amount);
+                    if let Err(err) = self
+                        .reply_controller_sender
+                        .send_additional_surbs_request(*recipient, amount)
+                    {
+                        if !self.task_client.is_shutdown_poll() {
+                            error!("{err}");
+                        }
+                    }
                 }
             }
         }
@@ -399,16 +413,19 @@ pub enum ReceivedBufferMessage {
 struct RequestReceiver<R: MessageReceiver> {
     received_buffer: ReceivedMessagesBuffer<R>,
     query_receiver: ReceivedBufferRequestReceiver,
+    task_client: TaskClient,
 }
 
 impl<R: MessageReceiver> RequestReceiver<R> {
     fn new(
         received_buffer: ReceivedMessagesBuffer<R>,
         query_receiver: ReceivedBufferRequestReceiver,
+        task_client: TaskClient,
     ) -> Self {
         RequestReceiver {
             received_buffer,
             query_receiver,
+            task_client,
         }
     }
 
@@ -423,12 +440,12 @@ impl<R: MessageReceiver> RequestReceiver<R> {
         }
     }
 
-    async fn run_with_shutdown(&mut self, mut shutdown: nym_task::TaskClient) {
+    async fn run(&mut self) {
         debug!("Started RequestReceiver with graceful shutdown support");
-        while !shutdown.is_shutdown() {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 biased;
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv() => {
                     log::trace!("RequestReceiver: Received shutdown");
                 }
                 request = self.query_receiver.next() => {
@@ -441,7 +458,7 @@ impl<R: MessageReceiver> RequestReceiver<R> {
                 },
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv().await;
         log::debug!("RequestReceiver: Exiting");
     }
 }
@@ -449,25 +466,25 @@ impl<R: MessageReceiver> RequestReceiver<R> {
 struct FragmentedMessageReceiver<R: MessageReceiver> {
     received_buffer: ReceivedMessagesBuffer<R>,
     mixnet_packet_receiver: MixnetMessageReceiver,
+    task_client: TaskClient,
 }
 
 impl<R: MessageReceiver> FragmentedMessageReceiver<R> {
     fn new(
         received_buffer: ReceivedMessagesBuffer<R>,
         mixnet_packet_receiver: MixnetMessageReceiver,
+        task_client: TaskClient,
     ) -> Self {
         FragmentedMessageReceiver {
             received_buffer,
             mixnet_packet_receiver,
+            task_client,
         }
     }
 
-    async fn run_with_shutdown(
-        &mut self,
-        mut shutdown: nym_task::TaskClient,
-    ) -> Result<(), MessageRecoveryError> {
+    async fn run(&mut self) -> Result<(), MessageRecoveryError> {
         debug!("Started FragmentedMessageReceiver with graceful shutdown support");
-        while !shutdown.is_shutdown() {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 new_messages = self.mixnet_packet_receiver.next() => {
                     if let Some(new_messages) = new_messages {
@@ -477,12 +494,12 @@ impl<R: MessageReceiver> FragmentedMessageReceiver<R> {
                         break;
                     }
                 },
-                _ = shutdown.recv_with_delay() => {
+                _ = self.task_client.recv_with_delay() => {
                     log::trace!("FragmentedMessageReceiver: Received shutdown");
                 }
             }
         }
-        shutdown.recv_timeout().await;
+        self.task_client.recv_timeout().await;
         log::debug!("FragmentedMessageReceiver: Exiting");
         Ok(())
     }
@@ -501,41 +518,42 @@ impl<R: MessageReceiver + Clone + Send + 'static> ReceivedMessagesBufferControll
         reply_key_storage: SentReplyKeys,
         reply_controller_sender: ReplyControllerSender,
         metrics_reporter: ClientStatsSender,
+        task_client: TaskClient,
     ) -> Self {
         let received_buffer = ReceivedMessagesBuffer::new(
             local_encryption_keypair,
             reply_key_storage,
             reply_controller_sender,
             metrics_reporter,
+            task_client.fork("received_messages_buffer"),
         );
 
         ReceivedMessagesBufferController {
             fragmented_message_receiver: FragmentedMessageReceiver::new(
                 received_buffer.clone(),
                 mixnet_packet_receiver,
+                task_client.fork("fragmented_message_receiver"),
+            ),
+            request_receiver: RequestReceiver::new(
+                received_buffer,
+                query_receiver,
+                task_client.with_suffix("request_receiver"),
             ),
-            request_receiver: RequestReceiver::new(received_buffer, query_receiver),
         }
     }
 
-    pub fn start_with_shutdown(self, shutdown: nym_task::TaskClient) {
+    pub fn start(self) {
         let mut fragmented_message_receiver = self.fragmented_message_receiver;
         let mut request_receiver = self.request_receiver;
 
-        let shutdown_handle = shutdown.fork("fragmented_message_receiver");
         spawn_future(async move {
-            match fragmented_message_receiver
-                .run_with_shutdown(shutdown_handle)
-                .await
-            {
+            match fragmented_message_receiver.run().await {
                 Ok(_) => {}
                 Err(e) => error!("{e}"),
             }
         });
         spawn_future(async move {
-            request_receiver
-                .run_with_shutdown(shutdown.with_suffix("request_receiver"))
-                .await;
+            request_receiver.run().await;
         });
     }
 }

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);
 
@@ -860,7 +868,7 @@ where
         while !shutdown.is_shutdown() {
             tokio::select! {
                 biased;
-                _ = shutdown.recv_with_delay() => {
+                _ = shutdown.recv() => {
                     log::trace!("ReplyController: Received shutdown");
                 },
                 req = self.request_receiver.next() => match req {

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

@@ -15,6 +15,27 @@ pub(crate) fn new_control_channels() -> (ReplyControllerSender, ReplyControllerR
     (tx.into(), rx)
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum ReplyControllerSenderError {
+    #[error("failed to send retransmission data to reply controller")]
+    SendRetransmissionData(#[source] mpsc::TrySendError<ReplyControllerMessage>),
+
+    #[error("failed to send reply to reply controller")]
+    SendReply(#[source] mpsc::TrySendError<ReplyControllerMessage>),
+
+    #[error("failed to send additional surbs to reply controller")]
+    AdditionalSurbs(#[source] mpsc::TrySendError<ReplyControllerMessage>),
+
+    #[error("failed to send additional surbs request to reply controller")]
+    AdditionalSurbsRequest(#[source] mpsc::TrySendError<ReplyControllerMessage>),
+
+    #[error("failed to request lane queue length from reply controller")]
+    LaneQueueLength(#[source] mpsc::TrySendError<ReplyControllerMessage>),
+
+    #[error("response channel was dropped before we could receive the response")]
+    ResponseChannelDropped(#[source] oneshot::Canceled),
+}
+
 #[derive(Debug, Clone)]
 pub struct ReplyControllerSender(mpsc::UnboundedSender<ReplyControllerMessage>);
 
@@ -30,14 +51,14 @@ impl ReplyControllerSender {
         recipient: AnonymousSenderTag,
         timed_out_ack: Weak<PendingAcknowledgement>,
         extra_surb_request: bool,
-    ) {
+    ) -> Result<(), ReplyControllerSenderError> {
         self.0
             .unbounded_send(ReplyControllerMessage::RetransmitReply {
                 recipient,
                 timed_out_ack,
                 extra_surb_request,
             })
-            .expect("ReplyControllerReceiver has died!")
+            .map_err(ReplyControllerSenderError::SendRetransmissionData)
     }
 
     pub(crate) fn send_reply(
@@ -45,14 +66,14 @@ impl ReplyControllerSender {
         recipient: AnonymousSenderTag,
         message: Vec<u8>,
         lane: TransmissionLane,
-    ) {
+    ) -> Result<(), ReplyControllerSenderError> {
         self.0
             .unbounded_send(ReplyControllerMessage::SendReply {
                 recipient,
                 message,
                 lane,
             })
-            .expect("ReplyControllerReceiver has died!")
+            .map_err(ReplyControllerSenderError::SendReply)
     }
 
     pub(crate) fn send_additional_surbs(
@@ -60,42 +81,47 @@ impl ReplyControllerSender {
         sender_tag: AnonymousSenderTag,
         reply_surbs: Vec<ReplySurb>,
         from_surb_request: bool,
-    ) {
+    ) -> Result<(), ReplyControllerSenderError> {
         self.0
             .unbounded_send(ReplyControllerMessage::AdditionalSurbs {
                 sender_tag,
                 reply_surbs,
                 from_surb_request,
             })
-            .expect("ReplyControllerReceiver has died!")
+            .map_err(ReplyControllerSenderError::AdditionalSurbs)
     }
 
-    pub(crate) fn send_additional_surbs_request(&self, recipient: Recipient, amount: u32) {
+    pub(crate) fn send_additional_surbs_request(
+        &self,
+        recipient: Recipient,
+        amount: u32,
+    ) -> Result<(), ReplyControllerSenderError> {
         self.0
             .unbounded_send(ReplyControllerMessage::AdditionalSurbsRequest {
                 recipient: Box::new(recipient),
                 amount,
             })
-            .expect("ReplyControllerReceiver has died!")
+            .map_err(ReplyControllerSenderError::AdditionalSurbsRequest)
     }
 
-    pub async fn get_lane_queue_length(&self, connection_id: ConnectionId) -> usize {
+    pub async fn get_lane_queue_length(
+        &self,
+        connection_id: ConnectionId,
+    ) -> Result<usize, ReplyControllerSenderError> {
         let (response_tx, response_rx) = oneshot::channel();
-        self.0
+        if let Err(err) = self
+            .0
             .unbounded_send(ReplyControllerMessage::LaneQueueLength {
                 connection_id,
                 response_channel: response_tx,
             })
-            .expect("ReplyControllerReceiver has died!");
-
-        match response_rx.await {
-            Ok(length) => length,
-            Err(_) => {
-                error!("The reply controller has dropped our response channel!");
-                // TODO: should we panic here instead? this message implies something weird and unrecoverable has happened
-                0
-            }
+        {
+            return Err(ReplyControllerSenderError::LaneQueueLength(err));
         }
+
+        response_rx
+            .await
+            .map_err(ReplyControllerSenderError::ResponseChannelDropped)
     }
 }
 
@@ -110,7 +136,10 @@ impl ReplyQueueLengths {
         }
     }
 
-    pub async fn get_lane_queue_length(&self, connection_id: ConnectionId) -> usize {
+    pub async fn get_lane_queue_length(
+        &self,
+        connection_id: ConnectionId,
+    ) -> Result<usize, ReplyControllerSenderError> {
         self.reply_controller_sender
             .get_lane_queue_length(connection_id)
             .await
@@ -120,7 +149,7 @@ impl ReplyQueueLengths {
 pub(crate) type ReplyControllerReceiver = mpsc::UnboundedReceiver<ReplyControllerMessage>;
 
 #[derive(Debug)]
-pub(crate) enum ReplyControllerMessage {
+pub enum ReplyControllerMessage {
     RetransmitReply {
         recipient: AnonymousSenderTag,
         timed_out_ack: Weak<PendingAcknowledgement>,

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

@@ -22,7 +22,7 @@ use nym_sphinx::addressing::Recipient;
 use nym_statistics_common::clients::{
     ClientStatsController, ClientStatsReceiver, ClientStatsSender,
 };
-use nym_task::connections::TransmissionLane;
+use nym_task::{connections::TransmissionLane, TaskClient};
 use std::time::Duration;
 
 use crate::{
@@ -51,6 +51,9 @@ pub(crate) struct StatisticsControl {
 
     /// Config for stats reporting (enabled, address, interval)
     reporting_config: StatsReporting,
+
+    /// Task client for listening for shutdown
+    task_client: TaskClient,
 }
 
 impl StatisticsControl {
@@ -59,19 +62,24 @@ impl StatisticsControl {
         client_type: String,
         client_stats_id: String,
         report_tx: InputMessageSender,
+        task_client: TaskClient,
     ) -> (Self, ClientStatsSender) {
         let (stats_tx, stats_rx) = tokio::sync::mpsc::unbounded_channel();
 
         let stats = ClientStatsController::new(client_stats_id, client_type);
 
+        let mut task_client_stats_sender = task_client.fork("stats_sender");
+        task_client_stats_sender.disarm();
+
         (
             StatisticsControl {
                 stats,
                 stats_rx,
                 report_tx,
                 reporting_config,
+                task_client,
             },
-            ClientStatsSender::new(Some(stats_tx)),
+            ClientStatsSender::new(Some(stats_tx), task_client_stats_sender),
         )
     }
 
@@ -91,7 +99,7 @@ impl StatisticsControl {
         }
     }
 
-    async fn run_with_shutdown(&mut self, mut task_client: nym_task::TaskClient) {
+    async fn run(&mut self) {
         log::debug!("Started StatisticsControl with graceful shutdown support");
 
         #[cfg(not(target_arch = "wasm32"))]
@@ -121,10 +129,10 @@ impl StatisticsControl {
         let mut snapshot_interval =
             gloo_timers::future::IntervalStream::new(SNAPSHOT_INTERVAL.as_millis() as u32);
 
-        loop {
+        while !self.task_client.is_shutdown() {
             tokio::select! {
                 biased;
-                _ = task_client.recv() => {
+                _ = self.task_client.recv() => {
                     log::trace!("StatisticsControl: Received shutdown");
                     break;
                 },
@@ -149,29 +157,34 @@ impl StatisticsControl {
                 }
 
                 _ = local_report_interval.next() => {
-                    self.stats.local_report(&mut task_client);
+                    self.stats.local_report(&mut self.task_client);
                 }
             }
         }
         log::debug!("StatisticsControl: Exiting");
     }
 
-    pub(crate) fn start_with_shutdown(mut self, task_client: nym_task::TaskClient) {
+    pub(crate) fn start(mut self) {
         spawn_future(async move {
-            self.run_with_shutdown(task_client).await;
+            self.run().await;
         })
     }
 
-    pub(crate) fn create_and_start_with_shutdown(
+    pub(crate) fn create_and_start(
         reporting_config: StatsReporting,
         client_type: String,
         client_stats_id: String,
         report_tx: InputMessageSender,
-        task_client: nym_task::TaskClient,
+        task_client: TaskClient,
     ) -> ClientStatsSender {
-        let (controller, sender) =
-            Self::create(reporting_config, client_type, client_stats_id, report_tx);
-        controller.start_with_shutdown(task_client);
+        let (controller, sender) = Self::create(
+            reporting_config,
+            client_type,
+            client_stats_id,
+            report_tx,
+            task_client,
+        );
+        controller.start();
         sender
     }
 }

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

@@ -6,6 +6,7 @@ pub(crate) use accessor::{TopologyAccessor, TopologyReadPermit};
 use futures::StreamExt;
 use log::*;
 use nym_sphinx::addressing::nodes::NodeIdentity;
+use nym_task::TaskClient;
 use nym_topology::NymTopologyError;
 use std::time::Duration;
 
@@ -43,6 +44,8 @@ pub struct TopologyRefresher {
 
     refresh_rate: Duration,
     consecutive_failure_count: usize,
+
+    task_client: TaskClient,
 }
 
 impl TopologyRefresher {
@@ -50,12 +53,14 @@ impl TopologyRefresher {
         cfg: TopologyRefresherConfig,
         topology_accessor: TopologyAccessor,
         topology_provider: Box<dyn TopologyProvider + Send + Sync>,
+        task_client: TaskClient,
     ) -> Self {
         TopologyRefresher {
             topology_provider,
             topology_accessor,
             refresh_rate: cfg.refresh_rate,
             consecutive_failure_count: 0,
+            task_client,
         }
     }
 
@@ -142,7 +147,7 @@ impl TopologyRefresher {
         }
     }
 
-    pub fn start_with_shutdown(mut self, mut shutdown: nym_task::TaskClient) {
+    pub fn start(mut self) {
         spawn_future(async move {
             debug!("Started TopologyRefresher with graceful shutdown support");
 
@@ -155,17 +160,17 @@ impl TopologyRefresher {
             let mut interval =
                 gloo_timers::future::IntervalStream::new(self.refresh_rate.as_millis() as u32);
 
-            while !shutdown.is_shutdown() {
+            while !self.task_client.is_shutdown() {
                 tokio::select! {
                     _ = interval.next() => {
                         self.try_refresh().await;
                     },
-                    _ = shutdown.recv() => {
+                    _ = self.task_client.recv() => {
                         log::trace!("TopologyRefresher: Received shutdown");
                     },
                 }
             }
-            shutdown.recv_timeout().await;
+            self.task_client.recv_timeout().await;
             log::debug!("TopologyRefresher: Exiting");
         })
     }

common/client-libs/gateway-client/src/client/mod.rs

@@ -1053,7 +1053,7 @@ impl GatewayClient<InitOnly, EphemeralCredentialStorage> {
             connection: SocketState::NotConnected,
             packet_router,
             bandwidth_controller: None,
-            stats_reporter: ClientStatsSender::new(None),
+            stats_reporter: ClientStatsSender::new(None, task_client.clone()),
             negotiated_protocol: None,
             #[cfg(unix)]
             connection_fd_callback: None,

common/statistics/src/clients/mod.rs

@@ -25,19 +25,28 @@ pub type ClientStatsReceiver = tokio::sync::mpsc::UnboundedReceiver<ClientStatsE
 #[derive(Clone)]
 pub struct ClientStatsSender {
     stats_tx: Option<UnboundedSender<ClientStatsEvents>>,
+    task_client: TaskClient,
 }
 
 impl ClientStatsSender {
     /// Create a new statistics Sender
-    pub fn new(stats_tx: Option<UnboundedSender<ClientStatsEvents>>) -> Self {
-        ClientStatsSender { stats_tx }
+    pub fn new(
+        stats_tx: Option<UnboundedSender<ClientStatsEvents>>,
+        task_client: TaskClient,
+    ) -> Self {
+        ClientStatsSender {
+            stats_tx,
+            task_client,
+        }
     }
 
     /// Report a statistics event using the sender.
     pub fn report(&self, event: ClientStatsEvents) {
         if let Some(tx) = &self.stats_tx {
             if let Err(err) = tx.send(event) {
-                log::error!("Failed to send stats event: {:?}", err);
+                if !self.task_client.is_shutdown_poll() {
+                    log::error!("Failed to send stats event: {err}");
+                }
             }
         }
     }

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)
     }
 

common/wireguard/src/peer_handle.rs

@@ -43,7 +43,7 @@ impl PeerHandle {
         let timeout_check_interval = tokio_stream::wrappers::IntervalStream::new(
             tokio::time::interval(DEFAULT_PEER_TIMEOUT_CHECK),
         );
-        let mut task_client = task_client.fork(format!("peer-{public_key}"));
+        let mut task_client = task_client.fork(format!("peer_{public_key}"));
         task_client.disarm();
         PeerHandle {
             public_key,

gateway/src/node/client_handling/websocket/connection_handler/mod.rs

@@ -95,7 +95,7 @@ where
     S: AsyncRead + AsyncWrite + Unpin + Send,
 {
     // don't accept any new requests if we have already received shutdown
-    if handle.shutdown.is_shutdown() {
+    if handle.shutdown.is_shutdown_poll() {
         debug!("stopping the handle as we have received a shutdown");
         return;
     }

gateway/src/node/client_handling/websocket/listener.rs

@@ -50,7 +50,7 @@ impl Listener {
                 connection = tcp_listener.accept() => {
                     match connection {
                         Ok((socket, remote_addr)) => {
-                            let shutdown = self.shutdown.fork(format!("websocket-handler-{remote_addr}"));
+                            let shutdown = self.shutdown.fork(format!("websocket_handler_{remote_addr}"));
                             trace!("received a socket connection from {remote_addr}");
                             // TODO: I think we *REALLY* need a mechanism for having a maximum number of connected
                             // clients or spawned tokio tasks -> perhaps a worker system?

gateway/src/node/mod.rs

@@ -225,7 +225,7 @@ impl GatewayTasksBuilder {
                 handler_config,
                 nyxd_client,
                 self.identity_keypair.public_key().to_bytes(),
-                self.shutdown.fork("ecash-manager"),
+                self.shutdown.fork("ecash_manager"),
                 self.storage.clone(),
             )
             .await?,
@@ -279,13 +279,13 @@ impl GatewayTasksBuilder {
         let mut message_router_builder = SpMessageRouterBuilder::new(
             *self.identity_keypair.public_key(),
             self.mix_packet_sender.clone(),
-            self.shutdown.fork("network-requester-message-router"),
+            self.shutdown.fork("network_requester_message_router"),
         );
         let transceiver = message_router_builder.gateway_transceiver();
 
         let (on_start_tx, on_start_rx) = oneshot::channel();
         let mut nr_builder = NRServiceProviderBuilder::new(nr_opts.config.clone())
-            .with_shutdown(self.shutdown.fork("network-requester-sp"))
+            .with_shutdown(self.shutdown.fork("network_requester_sp"))
             .with_custom_gateway_transceiver(transceiver)
             .with_wait_for_gateway(true)
             .with_minimum_gateway_performance(0)
@@ -314,13 +314,13 @@ impl GatewayTasksBuilder {
         let mut message_router_builder = SpMessageRouterBuilder::new(
             *self.identity_keypair.public_key(),
             self.mix_packet_sender.clone(),
-            self.shutdown.fork("ipr-message-router"),
+            self.shutdown.fork("ipr_message_router"),
         );
         let transceiver = message_router_builder.gateway_transceiver();
 
         let (on_start_tx, on_start_rx) = oneshot::channel();
         let mut ip_packet_router = IpPacketRouter::new(ip_opts.config.clone())
-            .with_shutdown(self.shutdown.fork("ipr-sp"))
+            .with_shutdown(self.shutdown.fork("ipr_sp"))
             .with_custom_gateway_transceiver(Box::new(transceiver))
             .with_wait_for_gateway(true)
             .with_minimum_gateway_performance(0)
@@ -418,7 +418,7 @@ impl GatewayTasksBuilder {
         let mut message_router_builder = SpMessageRouterBuilder::new(
             *self.identity_keypair.public_key(),
             self.mix_packet_sender.clone(),
-            self.shutdown.fork("authenticator-message-router"),
+            self.shutdown.fork("authenticator_message_router"),
         );
         let transceiver = message_router_builder.gateway_transceiver();
 
@@ -431,7 +431,7 @@ impl GatewayTasksBuilder {
         )
         .with_ecash_verifier(ecash_manager)
         .with_custom_gateway_transceiver(transceiver)
-        .with_shutdown(self.shutdown.fork("authenticator-sp"))
+        .with_shutdown(self.shutdown.fork("authenticator_sp"))
         .with_wait_for_gateway(true)
         .with_minimum_gateway_performance(0)
         .with_custom_topology_provider(topology_provider)
@@ -451,7 +451,7 @@ impl GatewayTasksBuilder {
     pub fn build_stale_messages_cleaner(&self) -> StaleMessagesCleaner {
         StaleMessagesCleaner::new(
             &self.storage,
-            self.shutdown.fork("stale-messages-cleaner"),
+            self.shutdown.fork("stale_messages_cleaner"),
             self.config.debug.stale_messages_max_age,
             self.config.debug.stale_messages_cleaner_run_interval,
         )

nym-api/src/network_monitor/monitor/sender.rs

@@ -167,7 +167,7 @@ impl PacketSender {
         let gateway_packet_router = PacketRouter::new(
             ack_sender,
             message_sender,
-            task_client.fork("packet-router"),
+            task_client.fork("packet_router"),
         );
 
         let shared_keys = fresh_gateway_client_data
@@ -184,7 +184,10 @@ impl PacketSender {
             shared_keys,
             gateway_packet_router,
             Some(fresh_gateway_client_data.bandwidth_controller.clone()),
-            nym_statistics_common::clients::ClientStatsSender::new(None),
+            nym_statistics_common::clients::ClientStatsSender::new(
+                None,
+                task_client.fork("client_stats_sender"),
+            ),
             #[cfg(unix)]
             None,
             task_client,

service-providers/authenticator/src/cli/run.rs

@@ -41,7 +41,7 @@ pub(crate) async fn execute(args: &Run) -> Result<(), AuthenticatorError> {
         Arc::new(KeyPair::new(&mut OsRng)),
     );
     let task_handler = TaskHandle::default();
-    let handler = DummyHandler::new(peer_rx, task_handler.fork("peer-handler"));
+    let handler = DummyHandler::new(peer_rx, task_handler.fork("peer_handler"));
     tokio::spawn(async move {
         handler.run().await;
     });

wasm/node-tester/src/tester.rs

@@ -190,12 +190,15 @@ impl NymNodeTesterBuilder {
 
         let gateway_identity = gateway_info.gateway_id;
 
+        let mut stats_sender_task = task_manager.subscribe().named("stats_sender");
+        stats_sender_task.disarm();
+
         let mut gateway_client =
             if let Some(existing_client) = initialisation_result.authenticated_ephemeral_client {
                 existing_client.upgrade(
                     packet_router,
                     self.bandwidth_controller.take(),
-                    ClientStatsSender::new(None),
+                    ClientStatsSender::new(None, stats_sender_task),
                     gateway_task,
                 )
             } else {
@@ -211,7 +214,7 @@ impl NymNodeTesterBuilder {
                     Some(gateway_info.shared_key),
                     packet_router,
                     self.bandwidth_controller.take(),
-                    ClientStatsSender::new(None),
+                    ClientStatsSender::new(None, stats_sender_task),
                     gateway_task,
                 )
             };