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
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    /// Create a new runtime builder
25    pub fn builder() -> RuntimeBuilder {
26        RuntimeBuilder::default()
27    }
28
29    /// Run a single node until completion or shutdown
30    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    /// Run multiple nodes until all complete or shutdown
55    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    /// Wait for all nodes to shut down
83    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    /// Get the health checker
95    pub fn health_checker(&self) -> &Arc<HealthChecker> {
96        &self.health_checker
97    }
98
99    /// Get the supervisor
100    pub fn supervisor(&self) -> &Arc<Supervisor> {
101        &self.supervisor
102    }
103
104    /// Get the launcher (if enabled)
105    pub fn launcher(&self) -> Option<&Arc<Launcher>> {
106        self.launcher.as_ref()
107    }
108
109    /// Get the shutdown handle
110    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        /// Set log level (trace, debug, info, warn, error)
124        into, log_level, config.log_level, impl Into<String>
125    );
126
127    builder_setter!(
128        /// Set log format
129        assign, log_format, config.log_format, LogFormat
130    );
131
132    builder_setter!(
133        /// Set restart policy
134        assign, restart_policy, config.restart_policy, RestartPolicy
135    );
136
137    builder_setter!(
138        /// Set health check interval
139        assign, health_check_interval, config.health_check_interval, Duration
140    );
141
142    builder_setter!(
143        /// Set graceful shutdown timeout
144        assign, shutdown_timeout, config.shutdown_timeout, Duration
145    );
146
147    builder_setter!(
148        /// Enable launcher service for dynamic node management
149        assign, with_launcher, enable_launcher, bool
150    );
151
152    /// Build the runtime
153    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}