use crate::context::Context;
use crate::error::{Mecha10Error, Result};
use crate::messages::{HealthLevel, HealthStatus, Message};
use crate::node::{HealthPriority, NodeTag, NodeType};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Default)]
pub struct InvalidInputTracker {
count: Arc<AtomicU64>,
last_reason: Arc<Mutex<Option<String>>>,
}
impl InvalidInputTracker {
pub fn new() -> Self {
Self::default()
}
pub fn record(&self, reason: impl Into<String>) {
self.count.fetch_add(1, Ordering::Relaxed);
if let Ok(mut last) = self.last_reason.lock() {
*last = Some(reason.into());
}
}
pub fn total(&self) -> u64 {
self.count.load(Ordering::Relaxed)
}
pub fn last_reason(&self) -> Option<String> {
self.last_reason.lock().ok().and_then(|guard| guard.clone())
}
pub fn fold_into(&self, healthy: HealthStatus) -> HealthStatus {
let total = self.total();
if total == 0 {
return healthy;
}
let reason = self.last_reason().unwrap_or_else(|| "unknown".to_string());
HealthStatus::degraded(format!(
"{total} invalid runtime input(s) received; most recent: {reason}"
))
}
}
#[cfg(test)]
mod invalid_input_tracker_tests {
use super::*;
#[test]
fn fresh_tracker_passes_healthy_status_through_unchanged() {
let tracker = InvalidInputTracker::new();
let folded = tracker.fold_into(HealthStatus::healthy());
assert_eq!(folded.level, HealthLevel::Ok);
}
#[test]
fn recording_degrades_and_names_the_most_recent_reason_and_total() {
let tracker = InvalidInputTracker::new();
tracker.record("unknown format 'foo'");
tracker.record("unknown format 'bar'");
let folded = tracker.fold_into(HealthStatus::healthy());
assert_eq!(folded.level, HealthLevel::Degraded);
assert!(
folded.message.contains('2'),
"message should mention the total: {}",
folded.message
);
assert!(
folded.message.contains("bar"),
"message should mention the most recent reason: {}",
folded.message
);
}
#[test]
fn total_reflects_number_of_records() {
let tracker = InvalidInputTracker::new();
assert_eq!(tracker.total(), 0);
tracker.record("x");
tracker.record("y");
assert_eq!(tracker.total(), 2);
}
}
pub struct HealthReporter {
node_id: String,
redis_client: redis::Client,
ctx: Arc<Context>,
priority: HealthPriority,
ttl_seconds: u64,
store_history: bool,
node_type: Option<NodeType>,
tags: Vec<NodeTag>,
}
impl HealthReporter {
pub async fn new(ctx: &Context) -> Result<Self> {
Self::with_config(ctx, HealthPriority::Normal, 60, true).await
}
pub async fn with_config(
ctx: &Context,
priority: HealthPriority,
ttl_seconds: u64,
store_history: bool,
) -> Result<Self> {
Self::with_config_and_type_tags(ctx, priority, ttl_seconds, store_history, None, &[]).await
}
pub async fn with_config_and_type_tags(
ctx: &Context,
priority: HealthPriority,
ttl_seconds: u64,
store_history: bool,
node_type: Option<NodeType>,
tags: &[NodeTag],
) -> Result<Self> {
let redis_url = Self::get_redis_url()?;
let redis_client = redis::Client::open(redis_url.as_str())
.map_err(|e| Mecha10Error::Other(format!("Failed to connect to Redis for health reporting: {}", e)))?;
Ok(Self {
node_id: ctx.node_id().to_string(),
redis_client,
ctx: Arc::new(ctx.clone()),
priority,
ttl_seconds,
store_history,
node_type,
tags: tags.to_vec(),
})
}
pub async fn report(&self, status: HealthStatus) -> Result<()> {
use redis::AsyncCommands;
let mut conn = self
.redis_client
.get_multiplexed_async_connection()
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to get Redis connection: {}", e)))?;
let report = HealthReport {
node_id: self.node_id.clone(),
status,
priority: self.priority,
reported_at: crate::prelude::now_micros(),
node_type: self.node_type,
tags: self.tags.clone(),
};
let key = format!("mecha10:health:{}", self.node_id);
let value = serde_json::to_string(&report)
.map_err(|e| Mecha10Error::Other(format!("Failed to serialize health report: {}", e)))?;
conn.set_ex::<_, _, ()>(&key, value.clone(), self.ttl_seconds)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to store health status: {}", e)))?;
if self.store_history {
let history_key = format!("mecha10:health:history:{}", self.node_id);
let score = report.reported_at as f64;
conn.zadd::<_, _, _, ()>(&history_key, &value, score)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to store health history: {}", e)))?;
let cutoff_time = (report.reported_at - 86_400_000_000) as f64;
conn.zrembyscore::<_, _, _, ()>(&history_key, f64::MIN, cutoff_time)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to clean old health history: {}", e)))?;
}
self.ctx.publish_raw("/system/health", &report).await?;
Ok(())
}
pub async fn report_healthy(&self) -> Result<()> {
self.report(HealthStatus::healthy()).await
}
pub async fn report_degraded(&self, message: impl Into<String>) -> Result<()> {
self.report(HealthStatus::degraded(message)).await
}
pub async fn report_unhealthy(&self, message: impl Into<String>) -> Result<()> {
self.report(HealthStatus::unhealthy(message)).await
}
fn get_redis_url() -> Result<String> {
let url = Context::get_redis_url()?;
tracing::debug!("HealthReporter using Redis URL: {}", url);
Ok(url)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthReport {
pub node_id: String,
pub status: HealthStatus,
#[serde(default)]
pub priority: HealthPriority,
pub reported_at: u64,
#[serde(default)]
pub node_type: Option<NodeType>,
#[serde(default)]
pub tags: Vec<NodeTag>,
}
impl Message for HealthReport {}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemHealth {
pub level: HealthLevel,
pub total_nodes: usize,
pub healthy_nodes: usize,
pub degraded_nodes: usize,
pub unhealthy_nodes: usize,
pub unresponsive_nodes: usize,
pub timestamp: u64,
pub nodes: HashMap<String, NodeHealthInfo>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeHealthInfo {
pub node_id: String,
pub status: HealthStatus,
#[serde(default)]
pub priority: HealthPriority,
pub reported_at: u64,
pub age_micros: u64,
pub responsive: bool,
#[serde(default)]
pub node_type: Option<NodeType>,
#[serde(default)]
pub tags: Vec<NodeTag>,
}
pub struct HealthMonitor {
redis_client: redis::Client,
cache: Arc<RwLock<Option<(SystemHealth, std::time::Instant)>>>,
cache_ttl: Duration,
}
impl HealthMonitor {
pub async fn connect(redis_url: &str) -> Result<Self> {
let redis_client = redis::Client::open(redis_url)
.map_err(|e| Mecha10Error::Other(format!("Failed to connect to Redis: {}", e)))?;
Ok(Self {
redis_client,
cache: Arc::new(RwLock::new(None)),
cache_ttl: Duration::from_secs(5),
})
}
pub async fn system_health(&self) -> Result<SystemHealth> {
{
let cache = self.cache.read().await;
if let Some((health, cached_at)) = cache.as_ref() {
if cached_at.elapsed() < self.cache_ttl {
return Ok(health.clone());
}
}
}
let health = self.fetch_system_health().await?;
{
let mut cache = self.cache.write().await;
*cache = Some((health.clone(), std::time::Instant::now()));
}
Ok(health)
}
async fn fetch_system_health(&self) -> Result<SystemHealth> {
use redis::AsyncCommands;
let mut conn = self
.redis_client
.get_multiplexed_async_connection()
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to get Redis connection: {}", e)))?;
let pattern = "mecha10:health:*";
let keys: Vec<String> = conn
.keys(pattern)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to list health keys: {}", e)))?;
let keys: Vec<_> = keys.into_iter().filter(|k| !k.contains(":history:")).collect();
let mut nodes = HashMap::new();
let mut healthy_count = 0;
let mut degraded_count = 0;
let mut unhealthy_count = 0;
let mut unresponsive_count = 0;
let now = crate::prelude::now_micros();
for key in keys {
let value: Option<String> = conn
.get(&key)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to get health value: {}", e)))?;
if let Some(value) = value {
if let Ok(report) = serde_json::from_str::<HealthReport>(&value) {
let age_micros = now.saturating_sub(report.reported_at);
let timeout_micros = report.priority.timeout_micros();
let responsive = age_micros < timeout_micros;
if !responsive {
unresponsive_count += 1;
} else {
match report.status.level {
HealthLevel::Ok => healthy_count += 1,
HealthLevel::Degraded => degraded_count += 1,
HealthLevel::Error => unhealthy_count += 1,
HealthLevel::Unknown => unresponsive_count += 1,
}
}
nodes.insert(
report.node_id.clone(),
NodeHealthInfo {
node_id: report.node_id,
status: report.status,
priority: report.priority,
reported_at: report.reported_at,
age_micros,
responsive,
node_type: report.node_type,
tags: report.tags,
},
);
}
}
}
let level = if unhealthy_count > 0 || unresponsive_count > nodes.len() / 2 {
HealthLevel::Error
} else if degraded_count > 0 || unresponsive_count > 0 {
HealthLevel::Degraded
} else if healthy_count > 0 {
HealthLevel::Ok
} else {
HealthLevel::Unknown
};
Ok(SystemHealth {
level,
total_nodes: nodes.len(),
healthy_nodes: healthy_count,
degraded_nodes: degraded_count,
unhealthy_nodes: unhealthy_count,
unresponsive_nodes: unresponsive_count,
timestamp: now,
nodes,
})
}
pub async fn node_health(&self, node_id: &str) -> Result<Option<NodeHealthInfo>> {
let system_health = self.system_health().await?;
Ok(system_health.nodes.get(node_id).cloned())
}
pub async fn healthy_nodes(&self) -> Result<Vec<NodeHealthInfo>> {
self.nodes_matching(|n| n.responsive && n.status.level == HealthLevel::Ok)
.await
}
pub async fn degraded_nodes(&self) -> Result<Vec<NodeHealthInfo>> {
self.nodes_matching(|n| n.responsive && n.status.level == HealthLevel::Degraded)
.await
}
pub async fn unhealthy_nodes(&self) -> Result<Vec<NodeHealthInfo>> {
self.nodes_matching(|n| n.responsive && n.status.level == HealthLevel::Error)
.await
}
pub async fn unresponsive_nodes(&self) -> Result<Vec<NodeHealthInfo>> {
self.nodes_matching(|n| !n.responsive).await
}
async fn nodes_matching(&self, predicate: impl Fn(&NodeHealthInfo) -> bool) -> Result<Vec<NodeHealthInfo>> {
let system_health = self.system_health().await?;
Ok(system_health.nodes.values().filter(|n| predicate(n)).cloned().collect())
}
pub async fn clear_cache(&self) {
let mut cache = self.cache.write().await;
*cache = None;
}
pub async fn node_health_history(&self, node_id: &str, duration: Duration) -> Result<Vec<HealthReport>> {
use redis::AsyncCommands;
let mut conn = self
.redis_client
.get_multiplexed_async_connection()
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to connect to Redis: {}", e)))?;
let history_key = format!("mecha10:health:history:{}", node_id);
let now = crate::prelude::now_micros();
let duration_micros = duration.as_micros() as u64;
let start_time = now.saturating_sub(duration_micros);
let results: Vec<String> = conn
.zrangebyscore(&history_key, start_time as f64, now as f64)
.await
.map_err(|e| Mecha10Error::Other(format!("Failed to retrieve health history: {}", e)))?;
let mut reports = Vec::new();
for json_str in results {
match serde_json::from_str::<HealthReport>(&json_str) {
Ok(report) => reports.push(report),
Err(e) => {
eprintln!("Failed to deserialize health report: {}", e);
continue;
}
}
}
Ok(reports)
}
}
#[async_trait::async_trait]
pub trait HealthReportingExt {
async fn start_health_reporting<F, Fut>(&self, health_check_fn: F) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send;
async fn start_health_reporting_with_config<F, Fut>(
&self,
priority: HealthPriority,
interval: Duration,
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send;
async fn start_health_reporting_with_type_tags<F, Fut>(
&self,
node_type: Option<NodeType>,
tags: &'static [NodeTag],
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send;
async fn start_health_reporting_with_config_and_type_tags<F, Fut>(
&self,
priority: HealthPriority,
interval: Duration,
node_type: Option<NodeType>,
tags: &'static [NodeTag],
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send;
}
#[async_trait::async_trait]
impl HealthReportingExt for Context {
async fn start_health_reporting<F, Fut>(&self, health_check_fn: F) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send,
{
self.start_health_reporting_with_config(HealthPriority::Normal, Duration::from_secs(5), health_check_fn)
.await
}
async fn start_health_reporting_with_config<F, Fut>(
&self,
priority: HealthPriority,
interval: Duration,
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send,
{
self.start_health_reporting_with_config_and_type_tags(priority, interval, None, &[], health_check_fn)
.await
}
async fn start_health_reporting_with_type_tags<F, Fut>(
&self,
node_type: Option<NodeType>,
tags: &'static [NodeTag],
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send,
{
self.start_health_reporting_with_config_and_type_tags(
HealthPriority::Normal,
Duration::from_secs(5),
node_type,
tags,
health_check_fn,
)
.await
}
async fn start_health_reporting_with_config_and_type_tags<F, Fut>(
&self,
priority: HealthPriority,
interval: Duration,
node_type: Option<NodeType>,
tags: &'static [NodeTag],
health_check_fn: F,
) -> Result<()>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = HealthStatus> + Send,
{
let reporter = HealthReporter::with_config_and_type_tags(self, priority, 60, true, node_type, tags).await?;
tracing::info!(
"Health reporting enabled: priority={:?}, interval={:?}",
priority,
interval
);
tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
loop {
ticker.tick().await;
let status = health_check_fn().await;
if let Err(e) = reporter.report(status).await {
tracing::debug!("Failed to report health: {}", e);
}
}
});
Ok(())
}
}