allow overwriting existing sdk shutdown manager

This commit is contained in:
Simon Wicky
2025-10-24 16:17:29 +02:00
parent c61df79182
commit fd5a95fa4d
5 changed files with 43 additions and 29 deletions
@@ -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);
@@ -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");
+2 -2
View File
@@ -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())
}
+34 -16
View File
@@ -19,30 +19,45 @@ 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> {
/// 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
.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() {
return Ok(manager.clone());
Ok(manager.clone())
} else {
Err(RegistryAccessError::MissingShutdownManager)
}
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().with_cancel_on_panic())
})
.clone())
}
/// Check if an SDK manager has been created.
@@ -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));
+5 -9
View File
@@ -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);
}