use crate::config::{LogFormat, RestartPolicy, RuntimeConfig};
use crate::health::HealthChecker;
use crate::launcher::Launcher;
use crate::logging::init_logging;
use crate::node::NodeRunner;
use crate::shutdown::{wait_for_signal, ShutdownHandle};
use crate::supervisor::Supervisor;
use anyhow::Result;
use mecha10_core::builder_setter;
use std::sync::Arc;
use std::time::Duration;
pub struct Runtime {
config: RuntimeConfig,
health_checker: Arc<HealthChecker>,
supervisor: Arc<Supervisor>,
launcher: Option<Arc<Launcher>>,
shutdown: ShutdownHandle,
}
impl Runtime {
pub fn builder() -> RuntimeBuilder {
RuntimeBuilder::default()
}
pub async fn run_node(&mut self, name: &str, node: Box<dyn NodeRunner>) -> Result<()> {
tracing::info!("Starting runtime for node: {}", name);
let handle = self.supervisor.launch_node(node);
tokio::select! {
result = handle.wait() => {
result?;
tracing::info!("Node {} completed", name);
}
_ = wait_for_signal() => {
tracing::info!("Shutdown signal received");
self.shutdown.shutdown();
tokio::time::timeout(
self.config.shutdown_timeout,
self.wait_shutdown()
).await??;
}
}
Ok(())
}
pub async fn run_nodes(&mut self, nodes: Vec<Box<dyn NodeRunner>>) -> Result<()> {
tracing::info!("Starting runtime for {} nodes", nodes.len());
let handles = self.supervisor.launch_nodes(nodes);
tokio::select! {
_ = async {
for handle in handles {
let _ = handle.wait().await;
}
} => {
tracing::info!("All nodes completed");
}
_ = wait_for_signal() => {
tracing::info!("Shutdown signal received");
self.shutdown.shutdown();
tokio::time::timeout(
self.config.shutdown_timeout,
self.wait_shutdown()
).await??;
}
}
Ok(())
}
async fn wait_shutdown(&self) -> Result<()> {
loop {
let statuses = self.health_checker.check_all().await;
if statuses.is_empty() {
break;
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
Ok(())
}
pub fn health_checker(&self) -> &Arc<HealthChecker> {
&self.health_checker
}
pub fn supervisor(&self) -> &Arc<Supervisor> {
&self.supervisor
}
pub fn launcher(&self) -> Option<&Arc<Launcher>> {
self.launcher.as_ref()
}
pub fn shutdown(&self) -> &ShutdownHandle {
&self.shutdown
}
}
#[derive(Default)]
pub struct RuntimeBuilder {
config: RuntimeConfig,
enable_launcher: bool,
}
impl RuntimeBuilder {
builder_setter!(
into, log_level, config.log_level, impl Into<String>
);
builder_setter!(
assign, log_format, config.log_format, LogFormat
);
builder_setter!(
assign, restart_policy, config.restart_policy, RestartPolicy
);
builder_setter!(
assign, health_check_interval, config.health_check_interval, Duration
);
builder_setter!(
assign, shutdown_timeout, config.shutdown_timeout, Duration
);
builder_setter!(
assign, with_launcher, enable_launcher, bool
);
pub fn build(self) -> Runtime {
init_logging(&self.config.log_level, self.config.log_format.clone());
let shutdown = ShutdownHandle::new();
let health_checker = Arc::new(HealthChecker::new(self.config.health_check_interval));
let supervisor = Arc::new(Supervisor::new(
health_checker.clone(),
shutdown.clone(),
self.config.restart_policy.clone(),
));
let launcher = if self.enable_launcher {
Some(Arc::new(Launcher::new(supervisor.clone())))
} else {
None
};
Runtime {
config: self.config,
health_checker,
supervisor,
launcher,
shutdown,
}
}
}