use crate::ShutdownToken;
use async_trait::async_trait;
use nagisa_types::error::{Error, Result};
use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tokio::task::JoinSet;
#[derive(Clone, Default)]
pub struct ServiceBus {
inner: Arc<RwLock<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>>,
}
impl ServiceBus {
pub fn new() -> Self {
Self::default()
}
pub fn insert<T: Send + Sync + 'static>(&self, value: T) {
let mut map = self.inner.write().unwrap_or_else(|e| e.into_inner());
map.insert(TypeId::of::<T>(), Arc::new(value));
}
pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
let map = self.inner.read().unwrap_or_else(|e| e.into_inner());
map.get(&TypeId::of::<T>()).and_then(|v| Arc::clone(v).downcast::<T>().ok())
}
}
#[async_trait]
pub trait Service: Send + Sync + 'static {
fn id(&self) -> &'static str;
fn deps(&self) -> &'static [&'static str] {
&[]
}
async fn prepare(&self, bus: &ServiceBus) -> Result<()> {
let _ = bus;
Ok(())
}
async fn run(self: Arc<Self>, bus: ServiceBus, shutdown: ShutdownToken) -> Result<()> {
let _ = (bus, shutdown);
Ok(())
}
async fn cleanup(&self, bus: &ServiceBus) -> Result<()> {
let _ = bus;
Ok(())
}
}
#[derive(Default)]
pub struct Supervisor {
services: Vec<Arc<dyn Service>>,
optional: Vec<bool>,
bus: ServiceBus,
}
impl Supervisor {
pub fn new() -> Self {
Self::default()
}
#[allow(clippy::should_implement_trait)]
pub fn add(mut self, service: Arc<dyn Service>) -> Self {
self.services.push(service);
self.optional.push(false);
self
}
pub fn add_optional(mut self, service: Arc<dyn Service>) -> Self {
self.services.push(service);
self.optional.push(true);
self
}
pub fn bus(&self) -> ServiceBus {
self.bus.clone()
}
pub async fn run(self, shutdown: ShutdownToken) -> Result<()> {
let Supervisor { services, optional, bus } = self;
let order = toposort(&services)?;
let mut prepared: Vec<usize> = Vec::with_capacity(order.len());
for &idx in &order {
match services[idx].prepare(&bus).await {
Ok(()) => prepared.push(idx),
Err(e) if optional[idx] => {
tracing::warn!(service = services[idx].id(), error = %e, "可选服务 prepare 失败;已跳过");
}
Err(e) => {
tracing::error!(service = services[idx].id(), error = %e, "必选服务 prepare 失败,清理已就绪者");
cleanup_reverse(&services, &prepared, &bus).await;
return Err(e);
}
}
}
let local = shutdown.child_token();
let mut set: JoinSet<(usize, Result<()>)> = JoinSet::new();
for &idx in &prepared {
let svc = Arc::clone(&services[idx]);
let bus = bus.clone();
let child = local.child_token();
set.spawn(async move { (idx, svc.run(bus, child).await) });
}
let mut run_result: Result<()> = Ok(());
loop {
tokio::select! {
biased;
_ = shutdown.cancelled() => break,
joined = set.join_next() => {
match joined {
None => break,
Some(Ok((_idx, Ok(())))) => continue,
Some(Ok((idx, Err(e)))) if optional[idx] => {
tracing::warn!(service = services[idx].id(), error = %e, "可选服务 run 出错;已隔离,不影响其余");
continue;
}
Some(Ok((idx, Err(e)))) => {
tracing::error!(service = services[idx].id(), error = %e, "必选服务 run 出错,收束");
run_result = Err(e);
break;
}
Some(Err(join_err)) => {
tracing::error!(error = %join_err, "服务 run 任务 panic,收束");
run_result = Err(internal_err(format!("服务 run 任务 panic:{join_err}")));
break;
}
}
}
}
}
local.cancel();
set.shutdown().await;
cleanup_reverse(&services, &prepared, &bus).await;
run_result
}
}
async fn cleanup_reverse(services: &[Arc<dyn Service>], indices: &[usize], bus: &ServiceBus) {
for &idx in indices.iter().rev() {
if let Err(e) = services[idx].cleanup(bus).await {
tracing::error!(service = services[idx].id(), error = %e, "服务清理失败");
}
}
}
fn toposort(services: &[Arc<dyn Service>]) -> Result<Vec<usize>> {
let mut index: HashMap<&'static str, usize> = HashMap::with_capacity(services.len());
for (i, svc) in services.iter().enumerate() {
if index.insert(svc.id(), i).is_some() {
return Err(internal_err(format!("duplicate service id `{}`", svc.id())));
}
}
let n = services.len();
let mut dependents: Vec<Vec<usize>> = vec![Vec::new(); n];
let mut indegree: Vec<usize> = vec![0usize; n];
for (i, svc) in services.iter().enumerate() {
for dep in svc.deps() {
let Some(&d) = index.get(dep) else {
return Err(internal_err(format!("service `{}` depends on unknown service `{dep}`", svc.id())));
};
dependents[d].push(i);
indegree[i] += 1;
}
}
let mut queue: std::collections::VecDeque<usize> = (0..n).filter(|&i| indegree[i] == 0).collect();
let mut order: Vec<usize> = Vec::with_capacity(n);
while let Some(node) = queue.pop_front() {
order.push(node);
for &m in &dependents[node] {
indegree[m] -= 1;
if indegree[m] == 0 {
queue.push_back(m);
}
}
}
if order.len() != n {
let cyclic: Vec<&str> = (0..n).filter(|&i| indegree[i] > 0).map(|i| services[i].id()).collect();
return Err(internal_err(format!("dependency cycle among services: {cyclic:?}")));
}
Ok(order)
}
fn internal_err(msg: String) -> Error {
Error::action(msg)
}