use async_trait::async_trait;
use std::sync::Arc;
use crate::errors::CoreError;
#[async_trait]
pub trait AsyncInitializable: Send + Sync {
async fn initialize(&self) -> Result<(), CoreError>;
fn is_initialized(&self) -> bool {
true }
}
#[async_trait]
pub trait Disposable: Send + Sync {
async fn dispose(&self) -> Result<(), CoreError>;
}
#[async_trait]
pub trait LifecycleManaged: AsyncInitializable + Disposable {}
impl<T> LifecycleManaged for T where T: AsyncInitializable + Disposable {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ServiceState {
Registered,
Created,
Initialized,
Disposing,
Disposed,
}
pub struct ServiceLifecycleManager {
initializable_services: Vec<Arc<dyn AsyncInitializable>>,
disposable_services: Vec<Arc<dyn Disposable>>,
initializable_service_types: Vec<String>,
state: ServiceState,
disposal_handle: Option<tokio::task::JoinHandle<()>>,
}
impl std::fmt::Debug for ServiceLifecycleManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ServiceLifecycleManager")
.field(
"initializable_services_count",
&self.initializable_services.len(),
)
.field("disposable_services_count", &self.disposable_services.len())
.field("state", &self.state)
.field("has_disposal_handle", &self.disposal_handle.is_some())
.finish()
}
}
impl ServiceLifecycleManager {
pub fn new() -> Self {
Self {
initializable_services: Vec::new(),
disposable_services: Vec::new(),
initializable_service_types: Vec::new(),
state: ServiceState::Registered,
disposal_handle: None,
}
}
pub fn add_initializable<T: AsyncInitializable + 'static>(&mut self, service: Arc<T>) {
self.initializable_services.push(service);
self.initializable_service_types
.push(std::any::type_name::<T>().to_string());
}
pub fn add_disposable<T: Disposable + 'static>(&mut self, service: Arc<T>) {
self.disposable_services.push(service);
}
pub fn add_lifecycle_managed<T: LifecycleManaged + 'static>(&mut self, service: Arc<T>) {
let service_clone = service.clone();
self.initializable_services.push(service_clone);
self.initializable_service_types
.push(std::any::type_name::<T>().to_string());
self.disposable_services.push(service);
}
pub async fn initialize_all(&mut self) -> Result<(), CoreError> {
if self.state != ServiceState::Registered {
return Err(CoreError::InvalidServiceDescriptor {
message: format!("Cannot initialize services in state: {:?}", self.state),
});
}
self.state = ServiceState::Created;
for (index, service) in self.initializable_services.iter().enumerate() {
let service_type = self
.initializable_service_types
.get(index)
.map(|s| s.as_str())
.unwrap_or("unknown");
service
.initialize()
.await
.map_err(|e| CoreError::ServiceInitializationFailed {
service_type: service_type.to_string(),
source: Box::new(e),
})?;
}
self.state = ServiceState::Initialized;
Ok(())
}
pub async fn initialize_all_with_timeout(
&mut self,
timeout: std::time::Duration,
) -> Result<(), CoreError> {
let init_future = self.initialize_all();
match tokio::time::timeout(timeout, init_future).await {
Ok(result) => result,
Err(_) => Err(CoreError::ServiceInitializationFailed {
service_type: "timeout".to_string(),
source: Box::new(CoreError::InvalidServiceDescriptor {
message: format!("Service initialization timed out after {:?}", timeout),
}),
}),
}
}
pub async fn dispose_all(&mut self) -> Result<(), CoreError> {
if self.state == ServiceState::Disposed || self.state == ServiceState::Disposing {
return Ok(()); }
self.state = ServiceState::Disposing;
for service in self.disposable_services.iter().rev() {
if let Err(e) = service.dispose().await {
eprintln!("Error disposing service: {:?}", e);
}
}
self.state = ServiceState::Disposed;
self.disposal_handle = None; Ok(())
}
pub fn schedule_disposal(&mut self) {
if self.is_disposed() || self.disposal_handle.is_some() {
return; }
let services = std::mem::take(&mut self.disposable_services);
self.state = ServiceState::Disposing;
let handle = tokio::spawn(async move {
for service in services.iter().rev() {
if let Err(e) = service.dispose().await {
eprintln!("Error disposing service in background: {:?}", e);
}
}
});
self.disposal_handle = Some(handle);
}
pub async fn wait_for_disposal(&mut self) -> Result<(), CoreError> {
if let Some(handle) = self.disposal_handle.take() {
handle.await.map_err(|e| CoreError::SystemError {
message: format!("Disposal task failed: {}", e),
source: None,
})?;
self.state = ServiceState::Disposed;
}
Ok(())
}
pub fn state(&self) -> ServiceState {
self.state
}
pub fn is_initialized(&self) -> bool {
self.state == ServiceState::Initialized
}
pub fn is_disposed(&self) -> bool {
self.state == ServiceState::Disposed
}
pub fn initializable_count(&self) -> usize {
self.initializable_services.len()
}
pub fn disposable_count(&self) -> usize {
self.disposable_services.len()
}
}
impl Default for ServiceLifecycleManager {
fn default() -> Self {
Self::new()
}
}
impl Drop for ServiceLifecycleManager {
fn drop(&mut self) {
if !self.is_disposed() && !self.disposable_services.is_empty() {
if let Ok(_handle) = tokio::runtime::Handle::try_current() {
self.schedule_disposal();
} else {
eprintln!(
"Warning: ServiceLifecycleManager dropped with {} undisposed services. \
No tokio runtime available for async disposal. \
Call dispose_all() explicitly before dropping.",
self.disposable_services.len()
);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
#[derive(Default)]
struct TestService {
initialized: AtomicBool,
disposed: AtomicBool,
}
#[async_trait]
impl AsyncInitializable for TestService {
async fn initialize(&self) -> Result<(), CoreError> {
self.initialized.store(true, Ordering::SeqCst);
Ok(())
}
fn is_initialized(&self) -> bool {
self.initialized.load(Ordering::SeqCst)
}
}
#[async_trait]
impl Disposable for TestService {
async fn dispose(&self) -> Result<(), CoreError> {
self.disposed.store(true, Ordering::SeqCst);
Ok(())
}
}
#[tokio::test]
async fn test_lifecycle_manager_initialization() {
let mut manager = ServiceLifecycleManager::new();
let service = Arc::new(TestService::default());
assert!(!service.is_initialized());
manager.add_lifecycle_managed(service.clone());
manager.initialize_all().await.unwrap();
assert!(service.is_initialized());
assert!(manager.is_initialized());
}
#[tokio::test]
async fn test_lifecycle_manager_disposal() {
let mut manager = ServiceLifecycleManager::new();
let service = Arc::new(TestService::default());
manager.add_lifecycle_managed(service.clone());
manager.initialize_all().await.unwrap();
assert!(!service.disposed.load(Ordering::SeqCst));
manager.dispose_all().await.unwrap();
assert!(service.disposed.load(Ordering::SeqCst));
assert!(manager.is_disposed());
}
#[tokio::test]
async fn test_initialization_timeout() {
#[derive(Default)]
struct SlowService;
#[async_trait]
impl AsyncInitializable for SlowService {
async fn initialize(&self) -> Result<(), CoreError> {
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
Ok(())
}
}
let mut manager = ServiceLifecycleManager::new();
let service = Arc::new(SlowService::default());
manager.add_initializable(service);
let result = manager
.initialize_all_with_timeout(std::time::Duration::from_millis(50))
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_background_disposal() {
let mut manager = ServiceLifecycleManager::new();
let service = Arc::new(TestService::default());
manager.add_lifecycle_managed(service.clone());
manager.initialize_all().await.unwrap();
assert!(!service.disposed.load(Ordering::SeqCst));
manager.schedule_disposal();
manager.wait_for_disposal().await.unwrap();
assert!(service.disposed.load(Ordering::SeqCst));
assert!(manager.is_disposed());
}
#[tokio::test]
async fn test_drop_with_runtime() {
let service = Arc::new(TestService::default());
{
let mut manager = ServiceLifecycleManager::new();
manager.add_lifecycle_managed(service.clone());
manager.initialize_all().await.unwrap();
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
assert!(service.disposed.load(Ordering::SeqCst));
}
#[tokio::test]
async fn test_initialization_error_with_service_type() {
#[derive(Default)]
struct FailingService;
#[async_trait]
impl AsyncInitializable for FailingService {
async fn initialize(&self) -> Result<(), CoreError> {
Err(CoreError::InvalidServiceDescriptor {
message: "Test initialization failure".to_string(),
})
}
}
let mut manager = ServiceLifecycleManager::new();
let service = Arc::new(FailingService::default());
manager.add_initializable(service);
let result = manager.initialize_all().await;
assert!(result.is_err());
let error = result.unwrap_err();
if let CoreError::ServiceInitializationFailed { service_type, .. } = error {
assert!(service_type.contains("FailingService"));
assert!(!service_type.eq("unknown"));
} else {
panic!(
"Expected ServiceInitializationFailed error, got: {:?}",
error
);
}
}
}