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 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(())
    }

    /// Run with launcher service for dynamic node management
    pub async fn run_with_launcher(&mut self) -> Result<()> {
        let _launcher = self
            .launcher
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Launcher not enabled. Use builder.with_launcher(true)"))?;

        tracing::info!("Starting runtime with launcher service");

        // Wait for shutdown signal
        wait_for_signal().await;
        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 with lifecycle management
    ///
    /// This starts the runtime with lifecycle management enabled. The lifecycle manager
    /// will handle mode transitions and start/stop nodes accordingly.
    ///
    /// The runtime will:
    /// 1. Start nodes for the default mode
    /// 2. Wait for shutdown signal
    /// 3. Gracefully shut down all nodes
    ///
    /// Mode transitions should be triggered externally via the lifecycle manager's
    /// `change_mode()` method.
    pub async fn run_with_lifecycle(&mut self) -> Result<()> {
        let lifecycle = self
            .lifecycle
            .as_mut()
            .ok_or_else(|| anyhow::anyhow!("Lifecycle not enabled. Use builder.with_lifecycle()"))?;

        tracing::info!("Starting runtime with lifecycle management");

        // Start nodes for default mode
        let default_mode = lifecycle.current_mode().to_string();
        tracing::info!("Initializing default mode: {}", default_mode);

        match lifecycle.change_mode(&default_mode).await {
            Ok(result) => {
                tracing::info!("Default mode initialized: started {} nodes", result.nodes_started.len());
            }
            Err(e) => {
                tracing::error!("Failed to initialize default mode: {}", e);
                return Err(e);
            }
        }

        // Wait for shutdown signal
        wait_for_signal().await;
        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 mutable lifecycle manager (if enabled)
    pub fn lifecycle_mut(&mut self) -> Option<&mut LifecycleManager> {
        self.lifecycle.as_mut()
    }

    /// 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 {
    /// Set log level (trace, debug, info, warn, error)
    pub fn log_level(mut self, level: impl Into<String>) -> Self {
        self.config.log_level = level.into();
        self
    }

    /// Set log format
    pub fn log_format(mut self, format: LogFormat) -> Self {
        self.config.log_format = format;
        self
    }

    /// Set restart policy
    pub fn restart_policy(mut self, policy: RestartPolicy) -> Self {
        self.config.restart_policy = policy;
        self
    }

    /// Set health check interval
    pub fn health_check_interval(mut self, interval: Duration) -> Self {
        self.config.health_check_interval = interval;
        self
    }

    /// Set graceful shutdown timeout
    pub fn shutdown_timeout(mut self, timeout: Duration) -> Self {
        self.config.shutdown_timeout = timeout;
        self
    }

    /// Enable launcher service for dynamic node management
    pub fn with_launcher(mut self, enable: bool) -> Self {
        self.enable_launcher = enable;
        self
    }

    /// 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,
        }
    }
}