use core::time::Duration;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::RwLock;
use tokio::time::sleep;
#[ derive( Debug, Clone, Copy, PartialEq, Eq ) ]
pub enum HealthCheckStrategy
{
Ping,
LightweightApi,
FullEndpoint,
}
#[ derive( Debug, Clone ) ]
pub struct HealthCheckConfig
{
pub endpoint : String,
pub strategy : HealthCheckStrategy,
pub check_interval : Duration,
pub timeout : Duration,
pub unhealthy_threshold : u32,
}
impl Default for HealthCheckConfig
{
#[ inline ]
fn default() -> Self
{
Self {
endpoint : String::new( ),
strategy : HealthCheckStrategy::LightweightApi,
check_interval : Duration::from_secs( 30 ),
timeout : Duration::from_secs( 5 ),
unhealthy_threshold : 3,
}
}
}
#[ derive( Debug, Clone ) ]
pub struct HealthStatus
{
pub healthy : bool,
pub latency_ms : u64,
pub consecutive_failures : u32,
pub total_checks : u64,
pub last_check : Instant,
}
impl Default for HealthStatus
{
#[ inline ]
fn default() -> Self
{
Self {
healthy : true,
latency_ms : 0,
consecutive_failures : 0,
total_checks : 0,
last_check : Instant::now( ),
}
}
}
#[ derive( Debug ) ]
pub struct HealthCheckerState
{
pub status : HealthStatus,
pub monitoring : bool,
}
#[ derive( Debug, Clone ) ]
pub struct HealthChecker
{
config : HealthCheckConfig,
pub state : Arc< RwLock< HealthCheckerState > >,
}
impl HealthChecker
{
#[ inline ]
#[ must_use ]
pub fn new( config : HealthCheckConfig ) -> Self
{
Self {
config,
state : Arc::new( RwLock::new( HealthCheckerState {
status : HealthStatus::default( ),
monitoring : false,
} )),
}
}
#[ inline ]
pub async fn check_health( &self ) -> Result< HealthStatus, HealthCheckError >
{
let start = Instant::now( );
let result = match self.config.strategy
{
HealthCheckStrategy::Ping => self.ping_check( ).await,
HealthCheckStrategy::LightweightApi => self.lightweight_api_check( ).await,
HealthCheckStrategy::FullEndpoint => self.full_endpoint_check( ).await,
};
let latency = start.elapsed( );
#[ allow( clippy::cast_possible_truncation ) ]
let latency_ms = latency.as_millis( ) as u64;
let mut state = self.state.write( ).await;
state.status.total_checks += 1;
state.status.last_check = Instant::now( );
state.status.latency_ms = latency_ms;
match result
{
Ok( ( )) => {
state.status.consecutive_failures = 0;
state.status.healthy = true;
Ok( state.status.clone( ))
}
Err( e ) => {
state.status.consecutive_failures += 1;
if state.status.consecutive_failures >= self.config.unhealthy_threshold
{
state.status.healthy = false;
}
Err( e )
}
}
}
#[ inline ]
pub async fn start_monitoring( &self ) -> MonitorHandle
{
let mut state = self.state.write( ).await;
state.monitoring = true;
drop( state );
let checker = self.clone( );
let handle = tokio::spawn( async move {
checker.monitoring_loop( ).await;
} );
MonitorHandle { handle }
}
#[ inline ]
pub async fn stop_monitoring( &self )
{
let mut state = self.state.write( ).await;
state.monitoring = false;
}
#[ inline ]
pub async fn get_status( &self ) -> HealthStatus
{
let state = self.state.read( ).await;
state.status.clone( )
}
#[ inline ]
pub async fn is_monitoring( &self ) -> bool
{
let state = self.state.read( ).await;
state.monitoring
}
#[ inline ]
pub async fn reset( &self )
{
let mut state = self.state.write( ).await;
state.status = HealthStatus::default( );
}
async fn monitoring_loop( &self )
{
loop
{
{
let state = self.state.read( ).await;
if !state.monitoring
{
break;
}
}
let _ = self.check_health( ).await;
sleep( self.config.check_interval ).await;
}
}
async fn ping_check( &self ) -> Result< ( ), HealthCheckError >
{
let client = reqwest::Client::builder( )
.timeout( self.config.timeout )
.build( )
.map_err( |e| HealthCheckError::CheckFailed {
reason : format!( "Failed to create client : {e}" ),
} )?;
let response = client
.head( &self.config.endpoint )
.send( )
.await
.map_err( |e| {
if e.is_timeout( )
{
HealthCheckError::Timeout
} else {
HealthCheckError::CheckFailed {
reason : format!( "Ping failed : {e}" ),
}
}
} )?;
if response.status( ).is_success( ) || response.status( ).is_redirection( )
{
Ok( ( ))
} else {
Err( HealthCheckError::CheckFailed {
reason : format!( "Ping returned status : {}", response.status( )),
} )
}
}
async fn lightweight_api_check( &self ) -> Result< ( ), HealthCheckError >
{
let client = reqwest::Client::builder( )
.timeout( self.config.timeout )
.build( )
.map_err( |e| HealthCheckError::CheckFailed {
reason : format!( "Failed to create client : {e}" ),
} )?;
let response = client
.get( &self.config.endpoint )
.send( )
.await
.map_err( |e| {
if e.is_timeout( )
{
HealthCheckError::Timeout
} else {
HealthCheckError::CheckFailed {
reason : format!( "API check failed : {e}" ),
}
}
} )?;
if response.status( ).is_success( )
{
Ok( ( ))
} else {
Err( HealthCheckError::CheckFailed {
reason : format!( "API returned status : {}", response.status( )),
} )
}
}
async fn full_endpoint_check( &self ) -> Result< ( ), HealthCheckError >
{
let client = reqwest::Client::builder( )
.timeout( self.config.timeout )
.build( )
.map_err( |e| HealthCheckError::CheckFailed {
reason : format!( "Failed to create client : {e}" ),
} )?;
let test_payload = serde_json::json!( {
"inputs": "test"
} );
let response = client
.post( &self.config.endpoint )
.json( &test_payload )
.send( )
.await
.map_err( |e| {
if e.is_timeout( )
{
HealthCheckError::Timeout
} else {
HealthCheckError::CheckFailed {
reason : format!( "Full endpoint check failed : {e}" ),
}
}
} )?;
if response.status( ).is_success( ) || response.status( ).is_client_error( )
{
Ok( ( ))
} else {
Err( HealthCheckError::CheckFailed {
reason : format!( "Endpoint returned status : {}", response.status( )),
} )
}
}
}
#[ derive( Debug ) ]
pub struct MonitorHandle
{
handle : tokio::task::JoinHandle< ( ) >,
}
impl MonitorHandle
{
#[ inline ]
pub async fn stop( self )
{
self.handle.abort( );
let _ = self.handle.await;
}
}
#[ derive( Debug ) ]
pub enum HealthCheckError
{
Timeout,
CheckFailed {
reason : String,
},
}
impl core::fmt::Display for HealthCheckError
{
#[ inline ]
fn fmt( &self, f : &mut core::fmt::Formatter< '_ > ) -> core::fmt::Result
{
match self
{
Self::Timeout => write!( f, "Health check timed out" ),
Self::CheckFailed { reason } => write!( f, "Health check failed : {reason}" ),
}
}
}
impl std::error::Error for HealthCheckError {}