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