use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use parking_lot::RwLock;
use crate::healthcheck::{HealthcheckReport, HealthcheckStatus, RestHealthcheckRegistry};
#[derive(Debug, Default)]
pub struct DependencyChecker {
draining: AtomicBool,
deps: RwLock<BTreeMap<String, bool>>,
}
impl DependencyChecker {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn register_dep(&self, dep_gear: impl Into<String>) {
self.deps.write().entry(dep_gear.into()).or_insert(false);
}
pub fn mark_resolved(&self, dep_gear: &str) -> bool {
if let Some(resolved) = self.deps.write().get_mut(dep_gear) {
let was = *resolved;
*resolved = true;
!was
} else {
false
}
}
#[must_use]
pub fn unresolved_deps(&self) -> Vec<String> {
self.deps
.read()
.iter()
.filter(|(_, resolved)| !**resolved)
.map(|(name, _)| name.clone())
.collect()
}
#[must_use]
pub fn all_resolved(&self) -> bool {
self.deps.read().values().all(|resolved| *resolved)
}
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)
}
#[must_use]
pub fn is_ready(&self) -> bool {
!self.is_draining() && self.all_resolved()
}
}
pub struct ReadinessHealthcheck {
deps: Arc<DependencyChecker>,
}
impl ReadinessHealthcheck {
#[must_use]
pub fn new(deps: Arc<DependencyChecker>) -> Self {
Self { deps }
}
}
#[async_trait::async_trait]
impl crate::healthcheck::Healthcheck for ReadinessHealthcheck {
fn name(&self) -> &'static str {
"readiness"
}
async fn check(&self) -> crate::healthcheck::HealthcheckResult {
use crate::healthcheck::HealthcheckResult;
if self.deps.is_draining() {
return HealthcheckResult::unhealthy("draining").with_code("draining");
}
let unresolved = self.deps.unresolved_deps();
if unresolved.is_empty() {
HealthcheckResult::healthy()
} else {
HealthcheckResult::unhealthy(format!(
"unresolved dependencies: {}",
unresolved.join(", ")
))
.with_code("deps_unresolved")
}
}
}
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 {
deps: Arc<DependencyChecker>,
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.deps.unresolved_deps())
.field("draining", &self.deps.is_draining())
.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>,
{
let deps = Arc::new(DependencyChecker::new());
for dep in critical_deps {
deps.register_dep(dep);
}
Self::from_checker(deps, healthchecks, check_timeout)
}
#[must_use]
pub fn from_checker(
deps: Arc<DependencyChecker>,
healthchecks: Arc<RestHealthcheckRegistry>,
check_timeout: Duration,
) -> Arc<Self> {
Arc::new(Self {
deps,
startup_complete: AtomicBool::new(false),
healthchecks,
check_timeout,
})
}
#[must_use]
pub fn dependency_checker(&self) -> Arc<DependencyChecker> {
Arc::clone(&self.deps)
}
pub fn register_dep(&self, dep_gear: impl Into<String>) {
self.deps.register_dep(dep_gear);
}
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) {
if self.deps.mark_resolved(name) {
tracing::info!(dep = %name, "critical dependency resolved");
}
}
#[must_use]
pub fn all_deps_resolved(&self) -> bool {
self.deps.all_resolved()
}
pub fn set_draining(&self, draining: bool) {
self.deps.set_draining(draining);
}
#[must_use]
pub fn is_draining(&self) -> bool {
self.deps.is_draining()
}
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.deps.is_draining();
let startup_complete = self.is_startup_complete();
let unresolved_deps: Vec<String> = self.deps.unresolved_deps();
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;
#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod dep_checker_tests {
use super::DependencyChecker;
#[test]
fn no_deps_is_ready() {
let c = DependencyChecker::new();
assert!(c.all_resolved());
assert!(c.is_ready());
assert!(c.unresolved_deps().is_empty());
}
#[test]
fn unresolved_dep_lists_and_blocks() {
let c = DependencyChecker::new();
c.register_dep("billing");
c.register_dep("inventory");
assert!(!c.is_ready());
assert_eq!(
c.unresolved_deps(),
vec!["billing".to_owned(), "inventory".to_owned()]
);
}
#[test]
fn ready_once_all_deps_resolved() {
let c = DependencyChecker::new();
c.register_dep("billing");
c.register_dep("inventory");
assert!(c.mark_resolved("billing")); assert!(!c.is_ready());
assert!(c.mark_resolved("inventory"));
assert!(c.is_ready());
}
#[test]
fn register_dep_is_idempotent_and_preserves_resolved() {
let c = DependencyChecker::new();
c.register_dep("billing");
assert!(c.mark_resolved("billing"));
c.register_dep("billing"); assert!(c.is_ready());
assert!(!c.mark_resolved("billing")); }
#[test]
fn mark_resolved_unknown_dep_is_noop() {
let c = DependencyChecker::new();
assert!(!c.mark_resolved("nope")); assert!(c.is_ready());
}
#[test]
fn draining_overrides_ready() {
let c = DependencyChecker::new();
assert!(c.is_ready());
c.set_draining(true);
assert!(c.is_draining());
assert!(!c.is_ready());
c.set_draining(false);
assert!(c.is_ready());
}
#[tokio::test]
async fn readiness_healthcheck_maps_state_to_status() {
use super::ReadinessHealthcheck;
use crate::healthcheck::{Healthcheck, HealthcheckStatus};
use std::sync::Arc;
let checker = Arc::new(DependencyChecker::new());
checker.register_dep("billing");
let hc = ReadinessHealthcheck::new(checker.clone());
assert_eq!(hc.name(), "readiness");
let starting = hc.check().await;
assert_eq!(starting.status, HealthcheckStatus::Unhealthy);
assert_eq!(starting.code.as_deref(), Some("deps_unresolved"));
assert!(starting.message.unwrap().contains("billing"));
checker.mark_resolved("billing");
assert_eq!(hc.check().await.status, HealthcheckStatus::Healthy);
checker.set_draining(true);
let draining = hc.check().await;
assert_eq!(draining.status, HealthcheckStatus::Unhealthy);
assert_eq!(draining.code.as_deref(), Some("draining"));
}
}