mecha10_runtime/

runtime.rs

//! Main runtime for Mecha10 nodes

use crate::config::{LogFormat, RestartPolicy, RuntimeConfig};
use crate::health::HealthChecker;
use crate::launcher::Launcher;
use crate::lifecycle::{LifecycleManager, ModeConfig};
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::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

/// Main runtime for supervising and running nodes
pub struct Runtime {
    config: RuntimeConfig,
    health_checker: Arc<HealthChecker>,
    supervisor: Arc<Supervisor>,
    launcher: Option<Arc<Launcher>>,
    lifecycle: Option<LifecycleManager>,
    shutdown: ShutdownHandle,
}

impl Runtime {
    /// Create a new runtime builder
    pub fn builder() -> RuntimeBuilder {
        RuntimeBuilder::default()
    }

    /// Run a single node until completion or shutdown
    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);

        // Wait for either node completion or shutdown signal
        tokio::select! {
            result = handle.wait() => {
                result?;
                tracing::info!("Node {} completed", name);
            }
            _ = wait_for_signal() => {
                tracing::info!("Shutdown signal received");
                self.shutdown.shutdown();

                // Wait for graceful shutdown with timeout
                tokio::time::timeout(
                    self.config.shutdown_timeout,
                    self.wait_shutdown()
                ).await??;
            }
        }

        Ok(())
    }

    /// Run multiple nodes until all complete or shutdown
    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);

        // Wait for either all nodes to complete or shutdown signal
        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();

                // Wait for graceful shutdown with timeout
                tokio::time::timeout(
                    self.config.shutdown_timeout,
                    self.wait_shutdown()
                ).await??;
            }
        }

        Ok(())
    }

    /// Wait for all nodes to shut down
    async fn wait_shutdown(&self) -> Result<()> {
        // Health checker will track when nodes complete shutdown
        loop {
            let statuses = self.health_checker.check_all().await;
            if statuses.is_empty() {
                break;
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
        Ok(())
    }

    /// Get the health checker
    pub fn health_checker(&self) -> &Arc<HealthChecker> {
        &self.health_checker
    }

    /// Get the supervisor
    pub fn supervisor(&self) -> &Arc<Supervisor> {
        &self.supervisor
    }

    /// Get the launcher (if enabled)
    pub fn launcher(&self) -> Option<&Arc<Launcher>> {
        self.launcher.as_ref()
    }

    /// Get the lifecycle manager (if enabled)
    pub fn lifecycle(&self) -> Option<&LifecycleManager> {
        self.lifecycle.as_ref()
    }

    /// Get the shutdown handle
    pub fn shutdown(&self) -> &ShutdownHandle {
        &self.shutdown
    }
}

/// Builder for creating a runtime with custom configuration
#[derive(Default)]
pub struct RuntimeBuilder {
    config: RuntimeConfig,
    enable_launcher: bool,
    lifecycle_config: Option<(HashMap<String, ModeConfig>, String)>,
}

impl RuntimeBuilder {
    builder_setter!(
        /// Set log level (trace, debug, info, warn, error)
        into, log_level, config.log_level, impl Into<String>
    );

    builder_setter!(
        /// Set log format
        assign, log_format, config.log_format, LogFormat
    );

    builder_setter!(
        /// Set restart policy
        assign, restart_policy, config.restart_policy, RestartPolicy
    );

    builder_setter!(
        /// Set health check interval
        assign, health_check_interval, config.health_check_interval, Duration
    );

    builder_setter!(
        /// Set graceful shutdown timeout
        assign, shutdown_timeout, config.shutdown_timeout, Duration
    );

    builder_setter!(
        /// Enable launcher service for dynamic node management
        assign, with_launcher, enable_launcher, bool
    );

    /// Enable lifecycle management with mode configuration
    ///
    /// # Arguments
    ///
    /// * `modes` - Map of mode name to mode configuration
    /// * `default_mode` - Initial mode to start in
    pub fn with_lifecycle(mut self, modes: HashMap<String, ModeConfig>, default_mode: String) -> Self {
        self.lifecycle_config = Some((modes, default_mode));
        self
    }

    /// Build the runtime
    pub fn build(self) -> Runtime {
        // Initialize logging
        init_logging(&self.config.log_level, self.config.log_format.clone());

        // Create shutdown handle
        let shutdown = ShutdownHandle::new();

        // Create health checker
        let health_checker = Arc::new(HealthChecker::new(self.config.health_check_interval));

        // Create supervisor
        let supervisor = Arc::new(Supervisor::new(
            health_checker.clone(),
            shutdown.clone(),
            self.config.restart_policy.clone(),
        ));

        // Create launcher if enabled
        let launcher = if self.enable_launcher {
            Some(Arc::new(Launcher::new(supervisor.clone())))
        } else {
            None
        };

        // Create lifecycle manager if enabled
        let lifecycle = if let Some((modes, default_mode)) = self.lifecycle_config {
            Some(LifecycleManager::new(
                supervisor.clone() as Arc<dyn crate::lifecycle::SupervisorTrait>,
                modes,
                default_mode,
            ))
        } else {
            None
        };

        Runtime {
            config: self.config,
            health_checker,
            supervisor,
            launcher,
            lifecycle,
            shutdown,
        }
    }
}