Skip to main content

mecha10_runtime/
runtime.rs

1//! Main runtime for Mecha10 nodes
2
3use 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
15/// Main runtime for supervising and running nodes
16pub 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    /// Create a new runtime builder
26    pub fn builder() -> RuntimeBuilder {
27        RuntimeBuilder::default()
28    }
29
30    /// Run a single node until completion or shutdown
31    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    /// Run multiple nodes until all complete or shutdown
56    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    /// Wait for all nodes to shut down
84    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    /// Get the health checker
96    pub fn health_checker(&self) -> &Arc<HealthChecker> {
97        &self.health_checker
98    }
99
100    /// Get the supervisor
101    pub fn supervisor(&self) -> &Arc<Supervisor> {
102        &self.supervisor
103    }
104
105    /// Get the launcher (if enabled)
106    pub fn launcher(&self) -> Option<&Arc<Launcher>> {
107        self.launcher.as_ref()
108    }
109
110    /// Get the shutdown handle
111    pub fn shutdown(&self) -> &ShutdownHandle {
112        &self.shutdown
113    }
114}
115
116/// Builder for creating a runtime with custom configuration
117#[derive(Default)]
118pub struct RuntimeBuilder {
119    config: RuntimeConfig,
120    enable_launcher: bool,
121}
122
123impl RuntimeBuilder {
124    builder_setter!(
125        /// Set log level (trace, debug, info, warn, error)
126        into, log_level, config.log_level, impl Into<String>
127    );
128
129    builder_setter!(
130        /// Set log format
131        assign, log_format, config.log_format, LogFormat
132    );
133
134    builder_setter!(
135        /// Set restart policy
136        assign, restart_policy, config.restart_policy, RestartPolicy
137    );
138
139    builder_setter!(
140        /// Set health check interval
141        assign, health_check_interval, config.health_check_interval, Duration
142    );
143
144    builder_setter!(
145        /// Set graceful shutdown timeout
146        assign, shutdown_timeout, config.shutdown_timeout, Duration
147    );
148
149    builder_setter!(
150        /// Enable launcher service for dynamic node management
151        assign, with_launcher, enable_launcher, bool
152    );
153
154    /// Build the runtime
155    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}