diff --git a/common/client-core/src/client/base_client/mod.rs b/common/client-core/src/client/base_client/mod.rs index 38d9329f4a..61c40d9a16 100644 --- a/common/client-core/src/client/base_client/mod.rs +++ b/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 diff --git a/common/client-core/src/client/cover_traffic_stream.rs b/common/client-core/src/client/cover_traffic_stream.rs index a7b84cd3a4..74154d08b4 100644 --- a/common/client-core/src/client/cover_traffic_stream.rs +++ b/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 Stream for LoopCoverTrafficStream @@ -110,6 +113,7 @@ impl LoopCoverTrafficStream { traffic_config: config::Traffic, cover_config: config::CoverTraffic, stats_tx: ClientStatsSender, + task_client: TaskClient, ) -> Self { let rng = OsRng; @@ -128,6 +132,7 @@ impl LoopCoverTrafficStream { secondary_packet_size: traffic_config.secondary_packet_size, packet_type: traffic_config.packet_type, stats_tx, + task_client, } } @@ -224,7 +229,7 @@ impl LoopCoverTrafficStream { 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 { ); self.set_next_delay(sampled); + let mut shutdown = self.task_client.fork("select"); + spawn_future(async move { debug!("Started LoopCoverTrafficStream with graceful shutdown support"); diff --git a/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs b/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs index 24efc7a560..ab2f4e86cb 100644 --- a/common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs +++ b/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, sent_notification_listener: SentNotificationListener, action_controller: ActionController, + task_client: TaskClient, } impl AcknowledgementController @@ -216,6 +218,7 @@ where message_handler: MessageHandler, 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!"); }); diff --git a/common/client-core/src/client/real_messages_control/mod.rs b/common/client-core/src/client/real_messages_control/mod.rs index 71e649c69a..e3fc9ccd87 100644 --- a/common/client-core/src/client/real_messages_control/mod.rs +++ b/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 { lane_queue_lengths: LaneQueueLengths, client_connection_rx: ConnectionCommandReceiver, stats_tx: ClientStatsSender, + task_client: TaskClient, ) -> Self { let rng = OsRng; @@ -186,6 +188,7 @@ impl RealMessagesController { 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 { message_handler, reply_storage, reply_controller_receiver, + task_client.fork("reply_controller"), ); let out_queue_control = OutQueueControl::new( @@ -205,6 +209,7 @@ impl RealMessagesController { lane_queue_lengths, client_connection_rx, stats_tx, + task_client.with_suffix("out_queue_control"), ); RealMessagesController { @@ -214,22 +219,20 @@ impl RealMessagesController { } } - 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); } } diff --git a/common/client-core/src/client/real_messages_control/real_traffic_stream.rs b/common/client-core/src/client/real_messages_control/real_traffic_stream.rs index 9b85e25eca..e2b3045521 100644 --- a/common/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/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; diff --git a/common/client-core/src/client/replies/reply_controller/mod.rs b/common/client-core/src/client/replies/reply_controller/mod.rs index 9716b25b7e..085fbfa6e4 100644 --- a/common/client-core/src/client/replies/reply_controller/mod.rs +++ b/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 { message_handler: MessageHandler, full_reply_storage: CombinedReplyStorage, + + // Listen for shutdown signals + task_client: TaskClient, } impl ReplyController @@ -79,6 +83,7 @@ where message_handler: MessageHandler, 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); diff --git a/common/client-core/src/client/statistics_control.rs b/common/client-core/src/client/statistics_control.rs index 3e6ed230bb..a002348fab 100644 --- a/common/client-core/src/client/statistics_control.rs +++ b/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), ) diff --git a/common/task/src/manager.rs b/common/task/src/manager.rs index 55c55f040b..a9ebab36c5 100644 --- a/common/task/src/manager.rs +++ b/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) }