Transparent ShutdownManager with cascading ShutdownTrackers (#6040)

* Idea for transparent ShutdownManager use

* Tracker hierarchies

* Fix wasm shutdown, convinience shutdown method

Author: durch

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

@@ -1,9 +1,7 @@
 // Copyright 2023 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::error::ClientCoreError;
 use crate::{client::replies::reply_storage, config::DebugConfig};
-use nym_task::{ShutdownManager, ShutdownToken, ShutdownTracker};
 
 pub fn setup_empty_reply_surb_backend(debug_config: &DebugConfig) -> reply_storage::Empty {
     reply_storage::Empty {
@@ -15,49 +13,3 @@ pub fn setup_empty_reply_surb_backend(debug_config: &DebugConfig) -> reply_stora
             .maximum_reply_surb_storage_threshold,
     }
 }
-
-// old 'TaskHandle'
-pub(crate) enum ShutdownHelper {
-    Internal(ShutdownManager),
-    External(ShutdownTracker),
-}
-
-fn new_shutdown_manager() -> Result<ShutdownManager, ClientCoreError> {
-    cfg_if::cfg_if! {
-        if #[cfg(not(target_arch = "wasm32"))] {
-            Ok(ShutdownManager::build_new_default()?)
-        } else {
-            Ok(ShutdownManager::new_without_signals())
-        }
-    }
-}
-
-impl ShutdownHelper {
-    pub(crate) fn new(shutdown_tracker: Option<ShutdownTracker>) -> Result<Self, ClientCoreError> {
-        match shutdown_tracker {
-            None => Ok(ShutdownHelper::Internal(new_shutdown_manager()?)),
-            Some(shutdown_tracker) => Ok(ShutdownHelper::External(shutdown_tracker)),
-        }
-    }
-
-    pub(crate) fn into_internal(self) -> Option<ShutdownManager> {
-        match self {
-            ShutdownHelper::Internal(manager) => Some(manager),
-            ShutdownHelper::External(_) => None,
-        }
-    }
-
-    pub(crate) fn shutdown_token(&self) -> ShutdownToken {
-        match self {
-            ShutdownHelper::External(shutdown) => shutdown.clone_shutdown_token(),
-            ShutdownHelper::Internal(shutdown) => shutdown.clone_shutdown_token(),
-        }
-    }
-
-    pub(crate) fn tracker(&self) -> &ShutdownTracker {
-        match self {
-            ShutdownHelper::External(shutdown) => shutdown,
-            ShutdownHelper::Internal(shutdown) => shutdown.shutdown_tracker(),
-        }
-    }
-}

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

@@ -4,7 +4,6 @@
 use super::mix_traffic::ClientRequestSender;
 use super::received_buffer::ReceivedBufferMessage;
 use super::statistics_control::StatisticsControl;
-use crate::client::base_client::helpers::ShutdownHelper;
 use crate::client::base_client::storage::helpers::store_client_keys;
 use crate::client::base_client::storage::MixnetClientStorage;
 use crate::client::cover_traffic_stream::LoopCoverTrafficStream;
@@ -53,7 +52,7 @@ use nym_sphinx::receiver::{ReconstructedMessage, SphinxMessageReceiver};
 use nym_statistics_common::clients::ClientStatsSender;
 use nym_statistics_common::generate_client_stats_id;
 use nym_task::connections::{ConnectionCommandReceiver, ConnectionCommandSender, LaneQueueLengths};
-use nym_task::{ShutdownManager, ShutdownTracker};
+use nym_task::ShutdownTracker;
 use nym_topology::provider_trait::TopologyProvider;
 use nym_topology::HardcodedTopologyProvider;
 use nym_validator_client::nym_api::NymApiClientExt;
@@ -881,8 +880,12 @@ where
         let shared_topology_accessor =
             TopologyAccessor::new(self.config.debug.topology.ignore_egress_epoch_role);
 
-        // Shutdown notifier for signalling tasks to stop
-        let shutdown = ShutdownHelper::new(self.shutdown)?;
+        // Create a shutdown tracker for this client - either as a child of provided tracker
+        // or get one from the registry
+        let shutdown_tracker = match self.shutdown {
+            Some(parent_tracker) => parent_tracker.child_tracker(),
+            None => nym_task::get_sdk_shutdown_tracker()?,
+        };
 
         // channels responsible for dealing with reply-related fun
         let (reply_controller_sender, reply_controller_receiver) =
@@ -914,7 +917,7 @@ where
             self.user_agent.clone(),
             generate_client_stats_id(*self_address.identity()),
             input_sender.clone(),
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         );
 
         // needs to be started as the first thing to block if required waiting for the gateway
@@ -924,14 +927,14 @@ where
             shared_topology_accessor.clone(),
             self_address.gateway(),
             self.wait_for_gateway,
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         )
         .await?;
 
         let gateway_packet_router = PacketRouter::new(
             ack_sender,
             mixnet_messages_sender,
-            shutdown.shutdown_token(),
+            shutdown_tracker.clone_shutdown_token(),
         );
 
         let gateway_transceiver = Self::setup_gateway_transceiver(
@@ -944,15 +947,15 @@ where
             stats_reporter.clone(),
             #[cfg(unix)]
             self.connection_fd_callback,
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         )
         .await?;
         let gateway_ws_fd = gateway_transceiver.ws_fd();
 
         let reply_storage = Self::setup_persistent_reply_storage(
             reply_storage_backend,
             key_rotation_config,
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         )
         .await?;
 
@@ -963,16 +966,18 @@ where
             reply_storage.key_storage(),
             reply_controller_sender.clone(),
             stats_reporter.clone(),
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         );
 
         // The message_sender is the transmitter for any component generating sphinx packets
         // that are to be sent to the mixnet. They are used by cover traffic stream and real
         // traffic stream.
         // The MixTrafficController then sends the actual traffic
 
-        let (message_sender, client_request_sender) =
-            Self::start_mix_traffic_controller(gateway_transceiver, shutdown.tracker());
+        let (message_sender, client_request_sender) = Self::start_mix_traffic_controller(
+            gateway_transceiver,
+            &shutdown_tracker.child_tracker(),
+        );
 
         // Channels that the websocket listener can use to signal downstream to the real traffic
         // controller that connections are closed.
@@ -1001,7 +1006,7 @@ where
             shared_lane_queue_lengths.clone(),
             client_connection_rx,
             stats_reporter.clone(),
-            shutdown.tracker(),
+            &shutdown_tracker.child_tracker(),
         );
 
         if !self
@@ -1017,7 +1022,7 @@ where
                 shared_topology_accessor.clone(),
                 message_sender,
                 stats_reporter.clone(),
-                shutdown.tracker(),
+                &shutdown_tracker.child_tracker(),
             );
         }
 
@@ -1045,7 +1050,7 @@ where
                 gateway_connection: GatewayConnection { gateway_ws_fd },
             },
             stats_reporter,
-            shutdown_handle: shutdown.into_internal(),
+            shutdown_handle: Some(shutdown_tracker), // The primary tracker for this client
             client_request_sender,
             forget_me: self.config.debug.forget_me,
             remember_me: self.config.debug.remember_me,
@@ -1061,7 +1066,7 @@ pub struct BaseClient {
     pub client_state: ClientState,
     pub stats_reporter: ClientStatsSender,
     pub client_request_sender: ClientRequestSender,
-    pub shutdown_handle: Option<ShutdownManager>,
+    pub shutdown_handle: Option<ShutdownTracker>,
     pub forget_me: ForgetMe,
     pub remember_me: RememberMe,
 }

common/client-core/src/error.rs

@@ -4,6 +4,7 @@
 use crate::client::mix_traffic::transceiver::ErasedGatewayError;
 use nym_crypto::asymmetric::ed25519::Ed25519RecoveryError;
 use nym_gateway_client::error::GatewayClientError;
+use nym_task::RegistryAccessError;
 use nym_topology::node::RoutingNodeError;
 use nym_topology::{NodeId, NymTopologyError};
 use nym_validator_client::nym_api::error::NymAPIError;
@@ -242,6 +243,9 @@ pub enum ClientCoreError {
 
     #[error("failed to select valid gateway due to incomputable latency")]
     GatewaySelectionFailure { source: WeightedError },
+
+    #[error("Could not access task registry, {0}")]
+    RegistryAccess(#[from] RegistryAccessError),
 }
 
 impl From<tungstenite::Error> for ClientCoreError {

common/task/src/cancellation/token.rs

@@ -17,6 +17,12 @@ pub struct ShutdownToken {
     inner: CancellationToken,
 }
 
+impl From<CancellationToken> for ShutdownToken {
+    fn from(inner: CancellationToken) -> Self {
+        ShutdownToken { inner }
+    }
+}
+
 impl ShutdownToken {
     /// A drop in no-op replacement for `send_status_msg` for easier migration from [TaskClient](crate::TaskClient).
     #[deprecated]
@@ -45,6 +51,13 @@ impl ShutdownToken {
         &self.inner
     }
 
+    /// Get an owned [CancellationToken](tokio_util::sync::CancellationToken) for public API use.
+    /// This is useful when you need to expose cancellation to SDK users without
+    /// exposing the internal ShutdownToken type.
+    pub fn to_cancellation_token(&self) -> CancellationToken {
+        self.inner.clone()
+    }
+
     /// Creates a `ShutdownToken` which will get cancelled whenever the
     /// current token gets cancelled. Unlike a cloned `ShutdownToken`,
     /// cancelling a child token does not cancel the parent token.

common/task/src/cancellation/tracker.rs

@@ -314,4 +314,40 @@ impl ShutdownTracker {
     pub fn clone_shutdown_token(&self) -> ShutdownToken {
         self.root_cancellation_token.clone()
     }
+
+    /// Create a child ShutdownTracker that inherits cancellation from this tracker
+    /// but has its own TaskTracker for managing sub-tasks.
+    ///
+    /// This enables hierarchical task management where:
+    /// - Parent cancellation flows to all children
+    /// - Each level tracks its own tasks independently
+    /// - Components can wait for their specific sub-tasks to complete
+    pub fn child_tracker(&self) -> ShutdownTracker {
+        // Child token inherits cancellation from parent
+        let child_token = self.root_cancellation_token.child_token();
+
+        // New TaskTracker for this level's tasks
+        let child_task_tracker = TaskTracker::new();
+
+        ShutdownTracker {
+            root_cancellation_token: child_token,
+            tracker: child_task_tracker,
+        }
+    }
+
+    /// Convenience method to perform a complete shutdown sequence.
+    /// This method:
+    /// 1. Signals cancellation to all tasks
+    /// 2. Closes the tracker to prevent new tasks
+    /// 3. Waits for all existing tasks to complete
+    pub async fn shutdown(self) {
+        // Signal cancellation to all tasks
+        self.root_cancellation_token.cancel();
+
+        // Close the tracker to prevent new tasks from being spawned
+        self.tracker.close();
+
+        // Wait for all existing tasks to complete
+        self.tracker.wait().await;
+    }
 }

common/task/src/lib.rs

@@ -5,6 +5,7 @@ pub mod cancellation;
 pub mod connections;
 pub mod event;
 pub mod manager;
+pub(crate) mod runtime_registry;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod signal;
 pub mod spawn;
@@ -18,3 +19,11 @@ pub use tokio_util::task::TaskTracker;
 
 #[cfg(not(target_arch = "wasm32"))]
 pub use signal::{wait_for_signal, wait_for_signal_and_error};
+
+pub use crate::runtime_registry::RegistryAccessError;
+
+/// Get or create a ShutdownTracker for SDK use.
+/// This provides automatic task management without requiring manual setup.
+pub fn get_sdk_shutdown_tracker() -> Result<ShutdownTracker, RegistryAccessError> {
+    Ok(runtime_registry::RuntimeRegistry::get_or_create_sdk()?.shutdown_tracker_owned())
+}

common/task/src/runtime_registry.rs

@@ -0,0 +1,96 @@
+// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
+// SPDX-License-Identifier: Apache-2.0
+
+use thiserror::Error;
+
+use crate::ShutdownManager;
+use std::sync::RwLock;
+use std::sync::{Arc, LazyLock};
+
+/// Global registry that manages ShutdownManagers transparently.
+/// This allows SDK components to get automatic task management without
+/// exposing the complexity to end users.
+pub(crate) struct RuntimeRegistry {
+    // For SDK clients: auto-created manager without signal handling
+    sdk_manager: RwLock<Option<Arc<ShutdownManager>>>,
+}
+
+#[derive(Debug, Error)]
+pub enum RegistryAccessError {
+    #[error("the runtime registry is poisoned")]
+    Poisoned,
+}
+
+impl RuntimeRegistry {
+    /// Get or create a ShutdownManager for SDK use.
+    /// This manager doesn't listen to OS signals, making it suitable for library use.
+    pub(crate) fn get_or_create_sdk() -> Result<Arc<ShutdownManager>, RegistryAccessError> {
+        let guard = REGISTRY
+            .sdk_manager
+            .read()
+            .map_err(|_| RegistryAccessError::Poisoned)?;
+        if let Some(manager) = guard.as_ref() {
+            return Ok(manager.clone());
+        }
+        drop(guard);
+
+        let mut guard = REGISTRY
+            .sdk_manager
+            .write()
+            .map_err(|_| RegistryAccessError::Poisoned)?;
+        Ok(guard
+            .get_or_insert_with(|| Arc::new(ShutdownManager::new_without_signals()))
+            .clone())
+    }
+
+    /// Check if an SDK manager has been created.
+    /// Useful for testing and debugging.
+    #[allow(dead_code)]
+    pub(crate) fn has_sdk_manager() -> Result<bool, RegistryAccessError> {
+        Ok(REGISTRY
+            .sdk_manager
+            .read()
+            .map_err(|_| RegistryAccessError::Poisoned)?
+            .is_some())
+    }
+
+    /// Clear the SDK manager.
+    /// This is primarily for testing to ensure isolation between tests.
+    #[cfg(test)]
+    pub(crate) async fn clear() -> Result<(), RegistryAccessError> {
+        *REGISTRY
+            .sdk_manager
+            .write()
+            .map_err(|_| RegistryAccessError::Poisoned)? = None;
+        Ok(())
+    }
+}
+
+/// Global instance of the runtime registry.
+/// Uses LazyLock for on-demand initialization.
+static REGISTRY: LazyLock<RuntimeRegistry> = LazyLock::new(|| RuntimeRegistry {
+    sdk_manager: RwLock::new(None),
+});
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_get_or_create_sdk() {
+        // Clear any existing manager
+        let _ = RuntimeRegistry::clear().await;
+
+        assert!(!RuntimeRegistry::has_sdk_manager().unwrap());
+
+        let manager1 = RuntimeRegistry::get_or_create_sdk().unwrap();
+        assert!(RuntimeRegistry::has_sdk_manager().unwrap());
+
+        let manager2 = RuntimeRegistry::get_or_create_sdk().unwrap();
+        // Should return the same instance
+        assert!(Arc::ptr_eq(&manager1, &manager2));
+
+        let _ = RuntimeRegistry::clear().await;
+        assert!(!RuntimeRegistry::has_sdk_manager().unwrap());
+    }
+}