mecha10-diagnostics 0.6.3

//! Simulation backend subprocess health collector
//!
//! Unlike the poll-based collectors in this module (`docker`, `redis`, `system`), simulation
//! backends (e.g. `simulator`) are process supervisors around a subprocess - there's nothing
//! to poll. Instead, the supervisor calls [`SimulationCollector::record_start`] /
//! [`SimulationCollector::record_exit`] as lifecycle events happen (spawn, crash, restart),
//! and each call publishes the resulting snapshot to the diagnostic topic.

use crate::topics::*;
use crate::types::*;
use anyhow::Result;
use mecha10_core::prelude::*;
use mecha10_core::topics::Topic;

/// Tracks and publishes health metrics for a supervised simulation backend subprocess.
pub struct SimulationCollector {
    source: String,
    backend: String,
    running: bool,
    restart_count: u64,
    last_start_us: u64,
    last_exit_reason: Option<String>,
}

impl SimulationCollector {
    /// Create a new collector for the given simulation backend (e.g. "mujoco").
    ///
    /// `source` identifies the publishing node (e.g. "simulator"); `backend` identifies which
    /// simulation backend is being supervised.
    pub fn new(source: impl Into<String>, backend: impl Into<String>) -> Self {
        Self {
            source: source.into(),
            backend: backend.into(),
            running: false,
            restart_count: 0,
            last_start_us: 0,
            last_exit_reason: None,
        }
    }

    /// Record that the subprocess has (re)started, and publish the updated snapshot.
    pub async fn record_start(&mut self, ctx: &Context) -> Result<()> {
        let metrics = self.note_start();
        self.publish_metrics(ctx, metrics).await
    }

    /// Record that the subprocess has exited (cleanly, crashed, or failed to spawn), and
    /// publish the updated snapshot.
    pub async fn record_exit(&mut self, ctx: &Context, reason: impl Into<String>) -> Result<()> {
        let metrics = self.note_exit(reason);
        self.publish_metrics(ctx, metrics).await
    }

    /// Publish the current snapshot without changing state.
    pub async fn publish(&self, ctx: &Context) -> Result<()> {
        let metrics = self.snapshot();
        self.publish_metrics(ctx, metrics).await
    }

    /// Pure state transition for a (re)start. Split out from `record_start` so the transition
    /// logic is unit testable without a live `Context`/Redis connection.
    fn note_start(&mut self) -> SimulationConnectionMetrics {
        self.running = true;
        self.restart_count += 1;
        self.last_start_us = crate::types::now_micros();
        self.snapshot()
    }

    /// Pure state transition for an exit. Split out from `record_exit` for the same reason as
    /// `note_start`.
    fn note_exit(&mut self, reason: impl Into<String>) -> SimulationConnectionMetrics {
        self.running = false;
        self.last_exit_reason = Some(reason.into());
        self.snapshot()
    }

    /// Snapshot the current state into a publishable metrics payload.
    fn snapshot(&self) -> SimulationConnectionMetrics {
        SimulationConnectionMetrics {
            backend: self.backend.clone(),
            running: self.running,
            restart_count: self.restart_count,
            last_start_us: self.last_start_us,
            last_exit_reason: self.last_exit_reason.clone(),
        }
    }

    async fn publish_metrics(&self, ctx: &Context, metrics: SimulationConnectionMetrics) -> Result<()> {
        let msg = DiagnosticMessage::new(&self.source, metrics);
        ctx.publish_to(
            Topic::<DiagnosticMessage<SimulationConnectionMetrics>>::new(TOPIC_DIAGNOSTICS_SIMULATION_CONNECTION),
            &msg,
        )
        .await?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new_collector_starts_not_running_with_no_restarts() {
        let collector = SimulationCollector::new("simulator", "mujoco");
        assert!(!collector.running);
        assert_eq!(collector.backend, "mujoco");
        assert_eq!(collector.restart_count, 0);
        assert_eq!(collector.last_start_us, 0);
        assert!(collector.last_exit_reason.is_none());
    }

    #[test]
    fn note_start_marks_running_and_increments_restart_count() {
        let mut collector = SimulationCollector::new("simulator", "mujoco");

        let metrics = collector.note_start();
        assert!(metrics.running);
        assert_eq!(metrics.restart_count, 1);
        assert!(metrics.last_start_us > 0);
        assert!(collector.running);
        assert_eq!(collector.restart_count, 1);

        let metrics = collector.note_start();
        assert_eq!(metrics.restart_count, 2);
    }

    #[test]
    fn note_exit_marks_not_running_and_sets_reason_without_touching_restart_count() {
        let mut collector = SimulationCollector::new("simulator", "mujoco");

        collector.note_start();
        let metrics = collector.note_exit("exited with code 1");

        assert!(!metrics.running);
        assert_eq!(metrics.last_exit_reason.as_deref(), Some("exited with code 1"));
        // Restart count and last_start_us are untouched by an exit.
        assert_eq!(metrics.restart_count, 1);
        assert!(metrics.last_start_us > 0);
    }

    #[test]
    fn snapshot_reflects_backend_identifier() {
        let collector = SimulationCollector::new("simulator", "mujoco");
        assert_eq!(collector.snapshot().backend, "mujoco");
    }
}