use async_trait::async_trait;
use crate::{NidusError, Result};
#[async_trait]
pub trait LifecycleHook: Send + Sync + 'static {
async fn on_startup(&self) -> Result<()> {
Ok(())
}
async fn on_shutdown(&self) -> Result<()> {
Ok(())
}
}
#[derive(Default)]
pub struct LifecycleRunner {
hooks: Vec<Box<dyn LifecycleHook>>,
}
impl LifecycleRunner {
pub fn new() -> Self {
Self::default()
}
pub fn hook<H>(mut self, hook: H) -> Self
where
H: LifecycleHook,
{
self.hooks.push(Box::new(hook));
self
}
pub async fn startup(&self) -> Result<()> {
let span = tracing::info_span!("lifecycle.startup", hook_count = self.hooks.len());
let _entered = span.enter();
let mut started: Vec<usize> = Vec::new();
tracing::debug!(hook_count = self.hooks.len(), "lifecycle startup begin");
for (index, hook) in self.hooks.iter().enumerate() {
tracing::debug!(hook_index = index, "lifecycle startup hook begin");
if let Err(source) = hook.on_startup().await {
tracing::error!(
hook_index = index,
error = %source,
"lifecycle startup hook failed"
);
let mut rollback_errors = Vec::new();
for started_index in started.into_iter().rev() {
tracing::debug!(
hook_index = started_index,
"lifecycle startup rollback hook begin"
);
if let Err(error) = self.hooks[started_index].on_shutdown().await {
tracing::error!(
hook_index = started_index,
error = %error,
"lifecycle startup rollback hook failed"
);
rollback_errors.push(error);
} else {
tracing::debug!(
hook_index = started_index,
"lifecycle startup rollback hook complete"
);
}
}
return Err(NidusError::LifecycleStartup {
source: Box::new(source),
rollback_errors,
});
}
tracing::debug!(hook_index = index, "lifecycle startup hook complete");
started.push(index);
}
tracing::debug!(hook_count = self.hooks.len(), "lifecycle startup complete");
Ok(())
}
pub async fn shutdown(&self) -> Result<()> {
let span = tracing::info_span!("lifecycle.shutdown", hook_count = self.hooks.len());
let _entered = span.enter();
tracing::debug!(hook_count = self.hooks.len(), "lifecycle shutdown begin");
for (index, hook) in self.hooks.iter().enumerate().rev() {
tracing::debug!(hook_index = index, "lifecycle shutdown hook begin");
hook.on_shutdown().await?;
tracing::debug!(hook_index = index, "lifecycle shutdown hook complete");
}
tracing::debug!(hook_count = self.hooks.len(), "lifecycle shutdown complete");
Ok(())
}
pub(crate) fn empty() -> Self {
Self::new()
}
}