use crate::{Context, HealthStatus, NodeState, Result};
use async_trait::async_trait;
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()),
}
}
}
#[async_trait]
pub trait NodeImpl: Send + Sync + Debug + 'static {
type Config: DeserializeOwned + Send + Sync + Debug;
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()),
}
}
}
#[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 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();
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(&ctx_for_health, priority, ttl_seconds, store_history).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
}