mecha10_runtime/
runtime.rs1use crate::config::{LogFormat, RestartPolicy, RuntimeConfig};
4use crate::health::HealthChecker;
5use crate::launcher::Launcher;
6use crate::logging::init_logging;
7use crate::node::NodeRunner;
8use crate::shutdown::{wait_for_signal, ShutdownHandle};
9use crate::supervisor::Supervisor;
10use anyhow::Result;
11use mecha10_core::builder_setter;
12use std::sync::Arc;
13use std::time::Duration;
14
15pub struct Runtime {
16 config: RuntimeConfig,
17 health_checker: Arc<HealthChecker>,
18 supervisor: Arc<Supervisor>,
19 launcher: Option<Arc<Launcher>>,
20 shutdown: ShutdownHandle,
21}
22
23impl Runtime {
24 pub fn builder() -> RuntimeBuilder {
26 RuntimeBuilder::default()
27 }
28
29 pub async fn run_node(&mut self, name: &str, node: Box<dyn NodeRunner>) -> Result<()> {
31 tracing::info!("Starting runtime for node: {}", name);
32
33 let handle = self.supervisor.launch_node(node);
34
35 tokio::select! {
36 result = handle.wait() => {
37 result?;
38 tracing::info!("Node {} completed", name);
39 }
40 _ = wait_for_signal() => {
41 tracing::info!("Shutdown signal received");
42 self.shutdown.shutdown();
43
44 tokio::time::timeout(
45 self.config.shutdown_timeout,
46 self.wait_shutdown()
47 ).await??;
48 }
49 }
50
51 Ok(())
52 }
53
54 pub async fn run_nodes(&mut self, nodes: Vec<Box<dyn NodeRunner>>) -> Result<()> {
56 tracing::info!("Starting runtime for {} nodes", nodes.len());
57
58 let handles = self.supervisor.launch_nodes(nodes);
59
60 tokio::select! {
61 _ = async {
62 for handle in handles {
63 let _ = handle.wait().await;
64 }
65 } => {
66 tracing::info!("All nodes completed");
67 }
68 _ = wait_for_signal() => {
69 tracing::info!("Shutdown signal received");
70 self.shutdown.shutdown();
71
72 tokio::time::timeout(
73 self.config.shutdown_timeout,
74 self.wait_shutdown()
75 ).await??;
76 }
77 }
78
79 Ok(())
80 }
81
82 async fn wait_shutdown(&self) -> Result<()> {
84 loop {
85 let statuses = self.health_checker.check_all().await;
86 if statuses.is_empty() {
87 break;
88 }
89 tokio::time::sleep(Duration::from_millis(100)).await;
90 }
91 Ok(())
92 }
93
94 pub fn health_checker(&self) -> &Arc<HealthChecker> {
96 &self.health_checker
97 }
98
99 pub fn supervisor(&self) -> &Arc<Supervisor> {
101 &self.supervisor
102 }
103
104 pub fn launcher(&self) -> Option<&Arc<Launcher>> {
106 self.launcher.as_ref()
107 }
108
109 pub fn shutdown(&self) -> &ShutdownHandle {
111 &self.shutdown
112 }
113}
114
115#[derive(Default)]
116pub struct RuntimeBuilder {
117 config: RuntimeConfig,
118 enable_launcher: bool,
119}
120
121impl RuntimeBuilder {
122 builder_setter!(
123 into, log_level, config.log_level, impl Into<String>
125 );
126
127 builder_setter!(
128 assign, log_format, config.log_format, LogFormat
130 );
131
132 builder_setter!(
133 assign, restart_policy, config.restart_policy, RestartPolicy
135 );
136
137 builder_setter!(
138 assign, health_check_interval, config.health_check_interval, Duration
140 );
141
142 builder_setter!(
143 assign, shutdown_timeout, config.shutdown_timeout, Duration
145 );
146
147 builder_setter!(
148 assign, with_launcher, enable_launcher, bool
150 );
151
152 pub fn build(self) -> Runtime {
154 init_logging(&self.config.log_level, self.config.log_format.clone());
155
156 let shutdown = ShutdownHandle::new();
157
158 let health_checker = Arc::new(HealthChecker::new(self.config.health_check_interval));
159
160 let supervisor = Arc::new(Supervisor::new(
161 health_checker.clone(),
162 shutdown.clone(),
163 self.config.restart_policy.clone(),
164 ));
165
166 let launcher = if self.enable_launcher {
167 Some(Arc::new(Launcher::new(supervisor.clone())))
168 } else {
169 None
170 };
171
172 Runtime {
173 config: self.config,
174 health_checker,
175 supervisor,
176 launcher,
177 shutdown,
178 }
179 }
180}