use crate::health::HealthCheckerTrait;
use crate::models::{HealthResult, Service};
use async_trait::async_trait;
use std::time::{Duration, Instant};
use tokio::process::Command;
use tokio::time::timeout;
use tracing::debug;
pub struct ScriptHealthChecker {
script: String,
timeout: Duration,
expected_exit_code: Option<i32>,
parse_output: bool,
}
impl ScriptHealthChecker {
pub fn new(script: String, timeout: Duration) -> Self {
Self {
script,
timeout,
expected_exit_code: Some(0),
parse_output: false,
}
}
pub fn with_exit_code(mut self, code: Option<i32>) -> Self {
self.expected_exit_code = code;
self
}
pub fn with_output_parsing(mut self, parse: bool) -> Self {
self.parse_output = parse;
self
}
async fn execute_script(&self, service: &Service) -> Result<(i32, String, String), String> {
let script = self
.script
.replace("{name}", &service.name)
.replace("{host}", &service.host)
.replace("{port}", &service.port.to_string())
.replace("{pid}", &service.pid.map(|p| p.to_string()).unwrap_or_default());
let parts: Vec<&str> = script.trim().split_whitespace().collect();
if parts.is_empty() {
return Err("Empty script".to_string());
}
let command = parts[0];
let args = &parts[1..];
debug!("Executing health check script: {} {:?}", command, args);
let start = Instant::now();
let output = timeout(
self.timeout,
Command::new(command).args(args).output(),
)
.await
.map_err(|_| format!("Script timeout after {:?}", self.timeout))?
.map_err(|e| format!("Failed to execute: {}", e))?;
let elapsed = start.elapsed();
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
let exit_code = output.status.code().unwrap_or(-1);
debug!(
"Script executed in {:?}, exit code: {}",
elapsed, exit_code
);
Ok((exit_code, stdout, stderr))
}
}
#[async_trait]
impl HealthCheckerTrait for ScriptHealthChecker {
async fn check(&self, service: &Service) -> HealthResult {
match self.execute_script(service).await {
Ok((exit_code, stdout, stderr)) => {
if let Some(expected) = self.expected_exit_code {
if exit_code != expected {
return HealthResult::unhealthy(
format!(
"Script exited with code {} (expected {})",
exit_code, expected
),
Some(if !stderr.is_empty() {
stderr
} else {
stdout
}),
);
}
}
if self.parse_output {
let output_lower = stdout.to_lowercase();
if output_lower.contains("healthy") || output_lower.contains("ok") {
return HealthResult::healthy(0);
} else if output_lower.contains("unhealthy")
|| output_lower.contains("error")
|| output_lower.contains("fail")
{
return HealthResult::unhealthy(
"Script reported unhealthy".to_string(),
Some(stdout),
);
}
}
HealthResult::healthy(0)
}
Err(e) => HealthResult::unhealthy(
format!("Script execution failed: {}", e),
None,
),
}
}
fn supports(&self, _service: &Service) -> bool {
true
}
}