use std::collections::BTreeSet;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use parking_lot::Mutex;
use crate::healthcheck::{HealthcheckReport, HealthcheckStatus, RestHealthcheckRegistry};
pub const DEFAULT_HEALTHCHECK_TIMEOUT: Duration = Duration::from_millis(500);
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum ReadinessLifecycle {
Starting,
Ready,
Degraded,
Draining,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ReadinessReport {
pub state: ReadinessLifecycle,
pub ready: bool,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub unresolved_deps: Vec<String>,
}
pub struct ReadinessState {
unresolved_deps: Mutex<BTreeSet<String>>,
draining: AtomicBool,
startup_complete: AtomicBool,
healthchecks: Arc<RestHealthcheckRegistry>,
check_timeout: Duration,
}
impl std::fmt::Debug for ReadinessState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ReadinessState")
.field("unresolved_deps", &self.unresolved_deps.lock())
.field("draining", &self.draining.load(Ordering::Relaxed))
.field(
"startup_complete",
&self.startup_complete.load(Ordering::Relaxed),
)
.field("check_timeout", &self.check_timeout)
.finish_non_exhaustive()
}
}
impl ReadinessState {
#[must_use]
pub fn new<I, S>(critical_deps: I, healthchecks: Arc<RestHealthcheckRegistry>) -> Arc<Self>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Self::with_check_timeout(critical_deps, healthchecks, DEFAULT_HEALTHCHECK_TIMEOUT)
}
#[must_use]
pub fn with_check_timeout<I, S>(
critical_deps: I,
healthchecks: Arc<RestHealthcheckRegistry>,
check_timeout: Duration,
) -> Arc<Self>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Arc::new(Self {
unresolved_deps: Mutex::new(critical_deps.into_iter().map(Into::into).collect()),
draining: AtomicBool::new(false),
startup_complete: AtomicBool::new(false),
healthchecks,
check_timeout,
})
}
pub fn mark_startup_complete(&self) {
self.startup_complete.store(true, Ordering::SeqCst);
}
#[must_use]
pub fn is_startup_complete(&self) -> bool {
self.startup_complete.load(Ordering::SeqCst)
}
pub fn mark_dep_resolved(&self, name: &str) {
let removed = self.unresolved_deps.lock().remove(name);
if removed {
tracing::info!(dep = %name, "critical dependency resolved");
}
}
#[must_use]
pub fn all_deps_resolved(&self) -> bool {
self.unresolved_deps.lock().is_empty()
}
pub fn set_draining(&self, draining: bool) {
self.draining.store(draining, Ordering::SeqCst);
}
#[must_use]
pub fn is_draining(&self) -> bool {
self.draining.load(Ordering::SeqCst)
}
pub async fn health_report(&self) -> HealthcheckReport {
self.healthchecks.report(self.check_timeout).await
}
pub async fn evaluate(&self) -> ReadinessReport {
let health = self.health_report().await;
let draining = self.is_draining();
let startup_complete = self.is_startup_complete();
let unresolved_deps: Vec<String> = self.unresolved_deps.lock().iter().cloned().collect();
let state = if draining {
ReadinessLifecycle::Draining
} else if !startup_complete
|| !unresolved_deps.is_empty()
|| health.status == HealthcheckStatus::Unhealthy
{
ReadinessLifecycle::Starting
} else if health.status == HealthcheckStatus::Degraded {
ReadinessLifecycle::Degraded
} else {
ReadinessLifecycle::Ready
};
let ready = matches!(
state,
ReadinessLifecycle::Ready | ReadinessLifecycle::Degraded
);
ReadinessReport {
state,
ready,
unresolved_deps,
}
}
}
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
#[path = "readiness_tests.rs"]
mod tests;