aion_server/observability/
health.rs

1//! Kubernetes-style liveness and readiness probes.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use aion_core::WorkflowId;
7use aion_store::ReadableEventStore;
8use axum::extract::State;
9use axum::http::StatusCode;
10use tokio::time::timeout;
11
12const LIVENESS_TIMEOUT: Duration = Duration::from_millis(50);
13const READINESS_TIMEOUT: Duration = Duration::from_millis(100);
14
15/// Cloneable state used by health probe handlers.
16#[derive(Clone)]
17pub struct HealthState {
18    store: Arc<dyn ReadableEventStore>,
19    runtime_initialized: bool,
20}
21
22impl HealthState {
23    /// Build health state from the store and runtime initialization flag.
24    #[must_use]
25    pub fn new(store: Arc<dyn ReadableEventStore>, runtime_initialized: bool) -> Self {
26        Self {
27            store,
28            runtime_initialized,
29        }
30    }
31}
32
33/// Liveness probe: validates that the async scheduler can run a trivial task promptly.
34pub async fn live() -> StatusCode {
35    let check = async {
36        let handle = tokio::spawn(async {});
37        handle.await.is_ok()
38    };
39
40    match timeout(LIVENESS_TIMEOUT, check).await {
41        Ok(true) => StatusCode::OK,
42        Ok(false) | Err(_) => StatusCode::SERVICE_UNAVAILABLE,
43    }
44}
45
46/// Readiness probe: validates runtime initialization and store reachability with a no-op read.
47pub async fn ready(State(state): State<HealthState>) -> StatusCode {
48    if !state.runtime_initialized {
49        return StatusCode::SERVICE_UNAVAILABLE;
50    }
51
52    let workflow_id = WorkflowId::new_v4();
53    match timeout(READINESS_TIMEOUT, state.store.read_history(&workflow_id)).await {
54        Ok(Ok(_)) => StatusCode::OK,
55        Ok(Err(_)) | Err(_) => StatusCode::SERVICE_UNAVAILABLE,
56    }
57}
aion_server/observability/health.rs

aion_server/observability/
health.rs
