allow overwriting existing sdk shutdown manager

Author: simonwicky

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

@@ -1005,7 +1005,7 @@ where
         // or get one from the registry
         let shutdown_tracker = match self.shutdown {
             Some(parent_tracker) => parent_tracker.clone(),
-            None => nym_task::get_sdk_shutdown_tracker()?,
+            None => nym_task::create_sdk_shutdown_tracker()?,
         };
 
         Self::start_event_control(self.event_tx, event_receiver, &shutdown_tracker);

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

@@ -205,7 +205,7 @@ impl LoopCoverTrafficStream<OsRng> {
                 TrySendError::Full(_) => {
                     // This isn't a problem, if the channel is full means we're already sending the
                     // max amount of messages downstream can handle.
-                    tracing::debug!("Failed to send cover message - channel full");
+                    tracing::trace!("Failed to send cover message - channel full");
                 }
                 TrySendError::Closed(_) => {
                     tracing::warn!("Failed to send cover message - channel closed");

common/task/src/lib.rs

@@ -24,6 +24,6 @@ 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())
+pub fn create_sdk_shutdown_tracker() -> Result<ShutdownTracker, RegistryAccessError> {
+    Ok(runtime_registry::RuntimeRegistry::create_sdk()?.shutdown_tracker_owned())
 }

common/task/src/runtime_registry.rs

@@ -19,32 +19,47 @@ pub(crate) struct RuntimeRegistry {
 pub enum RegistryAccessError {
     #[error("the runtime registry is poisoned")]
     Poisoned,
+
+    #[error("The SDK ShutdownManager already exists")]
+    ExistingShutdownManager,
+
+    #[error("No existing SDK ShutdownManager")]
+    MissingShutdownManager,
 }
 
 impl RuntimeRegistry {
-    /// Get or create a ShutdownManager for SDK use.
+    /// 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);
-
+    /// This function overwrite any existing manager!
+    pub(crate) fn create_sdk() -> Result<Arc<ShutdownManager>, RegistryAccessError> {
         let mut guard = REGISTRY
             .sdk_manager
             .write()
             .map_err(|_| RegistryAccessError::Poisoned)?;
+
         Ok(guard
-            .get_or_insert_with(|| {
-                Arc::new(ShutdownManager::new_without_signals().with_cancel_on_panic())
-            })
+            .insert(Arc::new(
+                ShutdownManager::new_without_signals().with_cancel_on_panic(),
+            ))
             .clone())
     }
 
+    /// Get the  ShutdownManager for SDK use.
+    /// This manager doesn't listen to OS signals, making it suitable for library use.
+    /// Not yet used, but maybe in the future
+    #[allow(dead_code)]
+    pub(crate) fn get_sdk() -> Result<Arc<ShutdownManager>, RegistryAccessError> {
+        let guard = REGISTRY
+            .sdk_manager
+            .read()
+            .map_err(|_| RegistryAccessError::Poisoned)?;
+        if let Some(manager) = guard.as_ref() {
+            Ok(manager.clone())
+        } else {
+            Err(RegistryAccessError::MissingShutdownManager)
+        }
+    }
+
     /// Check if an SDK manager has been created.
     /// Useful for testing and debugging.
     #[allow(dead_code)]
@@ -85,10 +100,13 @@ mod tests {
 
         assert!(!RuntimeRegistry::has_sdk_manager().unwrap());
 
-        let manager1 = RuntimeRegistry::get_or_create_sdk().unwrap();
+        // Error if nothing was created
+        assert!(RuntimeRegistry::get_sdk().is_err());
+
+        let manager1 = RuntimeRegistry::create_sdk().unwrap();
         assert!(RuntimeRegistry::has_sdk_manager().unwrap());
 
-        let manager2 = RuntimeRegistry::get_or_create_sdk().unwrap();
+        let manager2 = RuntimeRegistry::get_sdk().unwrap();
         // Should return the same instance
         assert!(Arc::ptr_eq(&manager1, &manager2));
 

sdk/rust/nym-sdk/src/mixnet/client.rs

@@ -736,15 +736,11 @@ where
             base_builder = base_builder.with_topology_provider(topology_provider);
         }
 
-        // Use custom shutdown if provided, otherwise get from registry
-        let shutdown_tracker = match self.custom_shutdown {
-            Some(custom) => custom,
-            None => {
-                // Auto-create from registry for SDK use
-                nym_task::get_sdk_shutdown_tracker()?
-            }
-        };
-        base_builder = base_builder.with_shutdown(shutdown_tracker);
+        // Use custom shutdown if provided, otherwise the sdk one will be used later down the line
+        if let Some(shutdown_tracker) = self.custom_shutdown {
+            base_builder = base_builder.with_shutdown(shutdown_tracker);
+        }
+
         if let Some(event_tx) = self.event_tx {
             base_builder = base_builder.with_event_tx(event_tx);
         }