mod lifecycle;
#[cfg(test)]
mod tests;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use dashmap::DashMap;
use serde::Serialize;
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;
use bamboo_domain::mcp_config::ReconnectConfig;
use bamboo_plugin::manifest::{GracefulShutdown, HealthCheckSpec};
#[derive(Debug, Clone)]
pub struct ServiceRuntimeConfig {
pub id: String,
pub plugin_id: String,
pub name: Option<String>,
pub command: PathBuf,
pub args: Vec<String>,
pub cwd: Option<PathBuf>,
pub env: HashMap<String, String>,
pub health_check: HealthCheckSpec,
pub restart_policy: ReconnectConfig,
pub graceful_shutdown: GracefulShutdown,
pub user_config_path: PathBuf,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ServiceState {
Starting,
Running,
Degraded,
Crashed,
Restarting,
Stopping,
Stopped,
}
#[derive(Debug, Clone, Serialize)]
pub struct ServiceStatusSnapshot {
pub id: String,
pub plugin_id: String,
pub state: ServiceState,
pub pid: Option<u32>,
pub restart_count: u32,
pub last_error: Option<String>,
}
#[derive(Debug, thiserror::Error)]
pub enum ServiceManagerError {
#[error("service '{0}' is already running")]
AlreadyRunning(String),
#[error("service '{0}' is not running")]
NotRunning(String),
}
pub(crate) struct ServiceRuntime {
config: ServiceRuntimeConfig,
state: RwLock<ServiceState>,
pid: AtomicU32,
restart_count: AtomicU32,
last_error: RwLock<Option<String>>,
shutdown: AtomicBool,
stop_token: CancellationToken,
supervisor: RwLock<Option<tokio::task::JoinHandle<()>>>,
}
#[derive(Default)]
pub struct ServiceManager {
runtimes: DashMap<String, Arc<ServiceRuntime>>,
}
impl ServiceManager {
pub fn new() -> Self {
Self::default()
}
pub async fn start_service(
&self,
config: ServiceRuntimeConfig,
) -> Result<(), ServiceManagerError> {
let id = config.id.clone();
let runtime = Arc::new(ServiceRuntime {
config,
state: RwLock::new(ServiceState::Starting),
pid: AtomicU32::new(0),
restart_count: AtomicU32::new(0),
last_error: RwLock::new(None),
shutdown: AtomicBool::new(false),
stop_token: CancellationToken::new(),
supervisor: RwLock::new(None),
});
match self.runtimes.entry(id) {
dashmap::mapref::entry::Entry::Occupied(occupied) => {
return Err(ServiceManagerError::AlreadyRunning(occupied.key().clone()));
}
dashmap::mapref::entry::Entry::Vacant(vacant) => {
vacant.insert(runtime.clone());
}
}
let handle = tokio::spawn(lifecycle::run_supervisor(runtime.clone()));
*runtime.supervisor.write().await = Some(handle);
Ok(())
}
pub async fn stop_service(&self, id: &str) -> Result<(), ServiceManagerError> {
let Some((_, runtime)) = self.runtimes.remove(id) else {
return Err(ServiceManagerError::NotRunning(id.to_string()));
};
runtime.shutdown.store(true, Ordering::SeqCst);
runtime.stop_token.cancel();
let handle = runtime.supervisor.write().await.take();
if let Some(handle) = handle {
let _ = handle.await;
}
Ok(())
}
pub fn is_running(&self, id: &str) -> bool {
self.runtimes.contains_key(id)
}
pub async fn status(&self, id: &str) -> Option<ServiceStatusSnapshot> {
let runtime = self.runtimes.get(id)?.clone();
Some(lifecycle::snapshot(&runtime).await)
}
pub async fn list_status(&self) -> Vec<ServiceStatusSnapshot> {
let runtimes: Vec<Arc<ServiceRuntime>> = self
.runtimes
.iter()
.map(|entry| entry.value().clone())
.collect();
let mut out = Vec::with_capacity(runtimes.len());
for runtime in &runtimes {
out.push(lifecycle::snapshot(runtime).await);
}
out
}
}