use crate::{Context, HealthStatus, NodeState, Result};
use async_trait::async_trait;
use schemars::JsonSchema;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tracing::{debug, error, info, warn};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum HealthPriority {
Critical,
#[default]
Normal,
Background,
}
impl HealthPriority {
pub const fn interval(&self) -> Duration {
match self {
Self::Critical => Duration::from_secs(1),
Self::Normal => Duration::from_secs(5),
Self::Background => Duration::from_secs(30),
}
}
pub const fn ttl_seconds(&self) -> u64 {
match self {
Self::Critical => 5,
Self::Normal => 60,
Self::Background => 120,
}
}
pub const fn timeout(&self) -> Duration {
match self {
Self::Critical => Duration::from_secs(3),
Self::Normal => Duration::from_secs(15),
Self::Background => Duration::from_secs(90),
}
}
pub const fn timeout_micros(&self) -> u64 {
match self {
Self::Critical => 3_000_000,
Self::Normal => 15_000_000,
Self::Background => 90_000_000,
}
}
}
impl std::fmt::Display for HealthPriority {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Critical => write!(f, "critical"),
Self::Normal => write!(f, "normal"),
Self::Background => write!(f, "background"),
}
}
}
#[derive(Debug, Clone)]
pub struct HealthReportingConfig {
pub enabled: bool,
pub priority: HealthPriority,
pub store_history: bool,
}
impl Default for HealthReportingConfig {
fn default() -> Self {
Self {
enabled: true,
priority: HealthPriority::Normal,
store_history: true,
}
}
}
impl HealthReportingConfig {
pub fn critical() -> Self {
Self {
enabled: true,
priority: HealthPriority::Critical,
store_history: false, }
}
pub fn normal() -> Self {
Self::default()
}
pub fn background() -> Self {
Self {
enabled: true,
priority: HealthPriority::Background,
store_history: true,
}
}
pub fn interval(&self) -> Duration {
self.priority.interval()
}
pub fn ttl_seconds(&self) -> u64 {
self.priority.ttl_seconds()
}
}
pub trait RateConfig {
fn rate_hz(&self) -> f32;
fn create_interval(&self) -> tokio::time::Interval {
let duration = Duration::from_secs_f32(1.0 / self.rate_hz());
tokio::time::interval(duration)
}
fn period(&self) -> Duration {
Duration::from_secs_f32(1.0 / self.rate_hz())
}
}
pub trait HealthStatusExt {
fn healthy(message: impl Into<String>) -> Self;
fn unhealthy(message: impl Into<String>) -> Self;
}
impl HealthStatusExt for HealthStatus {
fn healthy(message: impl Into<String>) -> Self {
Self {
healthy: true,
last_check: crate::prelude::now_micros(),
failure_count: 0,
message: Some(message.into()),
}
}
fn unhealthy(message: impl Into<String>) -> Self {
Self {
healthy: false,
last_check: crate::prelude::now_micros(),
failure_count: 1,
message: Some(message.into()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "kebab-case")]
pub enum NodeType {
Imu,
Camera,
ImageClassifier,
ObjectDetector,
BehaviorExecutor,
LlmCommand,
Teleop,
Motor,
WebsocketBridge,
VideoStreamer,
Simulator,
Speaker,
Listener,
}
impl NodeType {
pub const ALL: &'static [NodeType] = &[
NodeType::Imu,
NodeType::Camera,
NodeType::ImageClassifier,
NodeType::ObjectDetector,
NodeType::BehaviorExecutor,
NodeType::LlmCommand,
NodeType::Teleop,
NodeType::Motor,
NodeType::WebsocketBridge,
NodeType::VideoStreamer,
NodeType::Simulator,
NodeType::Speaker,
NodeType::Listener,
];
pub fn as_str(&self) -> &'static str {
match self {
NodeType::Imu => "imu",
NodeType::Camera => "camera",
NodeType::ImageClassifier => "image-classifier",
NodeType::ObjectDetector => "object-detector",
NodeType::BehaviorExecutor => "behavior-executor",
NodeType::LlmCommand => "llm-command",
NodeType::Teleop => "teleop",
NodeType::Motor => "motor",
NodeType::WebsocketBridge => "websocket-bridge",
NodeType::VideoStreamer => "video-streamer",
NodeType::Simulator => "simulator",
NodeType::Speaker => "speaker",
NodeType::Listener => "listener",
}
}
}
impl std::fmt::Display for NodeType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "kebab-case")]
pub enum NodeTag {
Sensor,
Inertial,
Capture,
Ai,
Classifier,
Detector,
Vision,
Planner,
Control,
Nlp,
Input,
Actuator,
Motion,
Relay,
Websocket,
Streamer,
Video,
Simulator,
Logger,
}
#[async_trait]
pub trait NodeImpl: Send + Sync + Debug + 'static {
type Config: DeserializeOwned + Send + Sync + Debug;
const NODE_TYPE: Option<NodeType> = None;
const NODE_TAGS: &'static [NodeTag] = &[];
async fn init(config: Self::Config) -> Result<Self>
where
Self: Sized;
async fn run(&mut self, ctx: &Context) -> Result<()>;
async fn shutdown(&mut self) -> Result<()> {
Ok(())
}
async fn health_check(&self) -> HealthStatus {
HealthStatus {
healthy: true,
last_check: crate::prelude::now_micros(),
failure_count: 0,
message: Some("Node is running".to_string()),
}
}
fn preflight(&self) -> Result<()> {
Ok(())
}
}
#[derive(Clone)]
pub struct Node<T: NodeImpl> {
node: Arc<Mutex<T>>,
state: Arc<Mutex<NodeState>>,
}
impl<T: NodeImpl> Node<T> {
pub async fn from_config(config: T::Config) -> Result<Self> {
info!("Initializing node");
let node = T::init(config).await?;
Ok(Self {
node: Arc::new(Mutex::new(node)),
state: Arc::new(Mutex::new(NodeState::Starting)),
})
}
pub async fn start(&self, ctx: Context) -> Result<()> {
{
let node = self.node.lock().await;
node.preflight()?;
}
{
let mut state = self.state.lock().await;
*state = NodeState::Running;
info!("Node transitioned to Running state");
}
let result = {
let mut node = self.node.lock().await;
debug!("Starting node execution");
node.run(&ctx).await
};
match result {
Ok(()) => {
info!("Node completed successfully");
self.shutdown().await?;
}
Err(e) => {
error!("Node failed with error: {}", e);
{
let mut state = self.state.lock().await;
*state = NodeState::Crashed;
}
self.shutdown().await?;
return Err(e);
}
}
Ok(())
}
pub async fn shutdown(&self) -> Result<()> {
info!("Shutting down node");
{
let mut state = self.state.lock().await;
*state = NodeState::Stopping;
}
let result = {
let mut node = self.node.lock().await;
node.shutdown().await
};
{
let mut state = self.state.lock().await;
*state = NodeState::Stopped;
info!("Node stopped");
}
if let Err(e) = result {
warn!("Error during shutdown: {}", e);
return Err(e);
}
Ok(())
}
pub async fn state(&self) -> NodeState {
self.state.lock().await.clone()
}
pub async fn health_check(&self) -> HealthStatus {
let node = self.node.lock().await;
node.health_check().await
}
pub async fn is_healthy(&self) -> bool {
matches!(self.state().await, NodeState::Running | NodeState::Starting)
}
}
pub async fn run_node<T: NodeImpl>(
config: T::Config,
ctx: Context,
health_config: HealthReportingConfig,
) -> Result<()> {
crate::prelude::init_logging();
#[cfg(feature = "messaging")]
crate::log_publisher::attach(ctx.node_id(), ctx.clone());
let node = Node::<T>::from_config(config).await?;
#[cfg(feature = "messaging")]
{
if health_config.enabled {
let node_for_health = node.node.clone();
let ctx_for_health = ctx.clone();
let priority = health_config.priority;
let report_interval = health_config.interval();
let ttl_seconds = health_config.ttl_seconds();
let store_history = health_config.store_history;
tokio::spawn(async move {
use crate::health::HealthReporter;
let reporter = match HealthReporter::with_config_and_type_tags(
&ctx_for_health,
priority,
ttl_seconds,
store_history,
T::NODE_TYPE,
T::NODE_TAGS,
)
.await
{
Ok(r) => r,
Err(e) => {
warn!("Health reporting disabled: {}", e);
return;
}
};
info!(
"Health reporting enabled: priority={}, interval={:?}, ttl={}s, history={}",
priority, report_interval, ttl_seconds, store_history
);
let mut interval = tokio::time::interval(report_interval);
info!("Health reporting loop starting");
loop {
interval.tick().await;
debug!("Health tick - checking node health");
let msg_status = match node_for_health.try_lock() {
Ok(node) => {
let status = node.health_check().await;
crate::messages::system::HealthStatus {
timestamp: status.last_check,
level: if status.healthy {
crate::messages::system::HealthLevel::Ok
} else {
crate::messages::system::HealthLevel::Error
},
message: status.message.unwrap_or_else(|| "No message".to_string()),
diagnostics: None,
}
}
Err(_) => {
debug!("Node busy, reporting as healthy");
crate::messages::system::HealthStatus {
timestamp: crate::prelude::now_micros(),
level: crate::messages::system::HealthLevel::Ok,
message: "Node running (busy)".to_string(),
diagnostics: None,
}
}
};
if let Err(e) = reporter.report(msg_status).await {
warn!("Failed to report health: {}", e);
} else {
tracing::trace!("Health report sent successfully");
}
}
});
} else {
info!("Health reporting disabled");
}
}
#[cfg(not(feature = "messaging"))]
let _ = &health_config;
node.start(ctx).await
}
#[cfg(all(test, not(feature = "messaging")))]
mod tests {
use super::*;
use crate::error::Mecha10Error;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc as StdArc;
#[derive(Debug, Default, serde::Deserialize)]
struct FakeConfig;
#[derive(Debug)]
struct OkNode {
ran: StdArc<AtomicBool>,
}
#[async_trait]
impl NodeImpl for OkNode {
type Config = FakeConfig;
async fn init(_config: Self::Config) -> Result<Self> {
Ok(Self {
ran: StdArc::new(AtomicBool::new(false)),
})
}
async fn run(&mut self, _ctx: &Context) -> Result<()> {
self.ran.store(true, Ordering::SeqCst);
Ok(())
}
}
#[derive(Debug)]
struct RejectingPreflightNode {
ran: StdArc<AtomicBool>,
}
#[async_trait]
impl NodeImpl for RejectingPreflightNode {
type Config = FakeConfig;
async fn init(_config: Self::Config) -> Result<Self> {
Ok(Self {
ran: StdArc::new(AtomicBool::new(false)),
})
}
async fn run(&mut self, _ctx: &Context) -> Result<()> {
self.ran.store(true, Ordering::SeqCst);
Ok(())
}
fn preflight(&self) -> Result<()> {
Err(Mecha10Error::Configuration("boom: bad node content".to_string()))
}
}
#[tokio::test]
async fn default_preflight_is_a_noop_and_run_still_executes() {
let node = Node::<OkNode>::from_config(FakeConfig).await.unwrap();
let ctx = Context::new("test-node-ok").await.unwrap();
let ran_handle = { node.node.lock().await.ran.clone() };
node.start(ctx).await.unwrap();
assert!(ran_handle.load(Ordering::SeqCst), "run() should have executed");
}
#[tokio::test]
async fn start_fails_fast_when_preflight_rejects_config_and_never_calls_run() {
let node = Node::<RejectingPreflightNode>::from_config(FakeConfig).await.unwrap();
let ctx = Context::new("test-node-rejecting").await.unwrap();
let ran_handle = { node.node.lock().await.ran.clone() };
let result = node.start(ctx).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("boom: bad node content"));
assert!(
!ran_handle.load(Ordering::SeqCst),
"run() must not execute when preflight() rejects the node"
);
}
#[test]
fn node_type_kebab_case_matches_existing_builtin_node_types_strings() {
let expected: &[(NodeType, &str)] = &[
(NodeType::Imu, "imu"),
(NodeType::Camera, "camera"),
(NodeType::ImageClassifier, "image-classifier"),
(NodeType::ObjectDetector, "object-detector"),
(NodeType::BehaviorExecutor, "behavior-executor"),
(NodeType::LlmCommand, "llm-command"),
(NodeType::Teleop, "teleop"),
(NodeType::Motor, "motor"),
(NodeType::WebsocketBridge, "websocket-bridge"),
(NodeType::VideoStreamer, "video-streamer"),
(NodeType::Simulator, "simulator"),
(NodeType::Speaker, "speaker"),
(NodeType::Listener, "listener"),
];
assert_eq!(expected.len(), NodeType::ALL.len());
for (variant, kebab) in expected {
let json = serde_json::to_string(variant).unwrap();
assert_eq!(json, format!("\"{}\"", kebab));
assert_eq!(variant.as_str(), *kebab);
let round_tripped: NodeType = serde_json::from_str(&json).unwrap();
assert_eq!(round_tripped, *variant);
}
}
#[test]
fn node_tag_kebab_case_matches_expected_strings() {
let expected: &[(NodeTag, &str)] = &[
(NodeTag::Sensor, "sensor"),
(NodeTag::Inertial, "inertial"),
(NodeTag::Capture, "capture"),
(NodeTag::Ai, "ai"),
(NodeTag::Classifier, "classifier"),
(NodeTag::Detector, "detector"),
(NodeTag::Vision, "vision"),
(NodeTag::Planner, "planner"),
(NodeTag::Control, "control"),
(NodeTag::Nlp, "nlp"),
(NodeTag::Input, "input"),
(NodeTag::Actuator, "actuator"),
(NodeTag::Motion, "motion"),
(NodeTag::Relay, "relay"),
(NodeTag::Websocket, "websocket"),
(NodeTag::Streamer, "streamer"),
(NodeTag::Video, "video"),
(NodeTag::Simulator, "simulator"),
(NodeTag::Logger, "logger"),
];
for (variant, kebab) in expected {
let json = serde_json::to_string(variant).unwrap();
assert_eq!(json, format!("\"{}\"", kebab));
let round_tripped: NodeTag = serde_json::from_str(&json).unwrap();
assert_eq!(round_tripped, *variant);
}
}
#[test]
fn node_impl_defaults_node_type_to_none_and_tags_to_empty() {
assert_eq!(OkNode::NODE_TYPE, None);
assert!(OkNode::NODE_TAGS.is_empty());
}
}