aion-server 0.4.0

Deployable HTTP, gRPC, WebSocket, and worker endpoint for Aion workflows.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

//! Kubernetes-style liveness and readiness probes.

use std::sync::Arc;
use std::time::Duration;

use aion_core::WorkflowId;
use aion_store::ReadableEventStore;
use axum::extract::State;
use axum::http::StatusCode;
use tokio::time::timeout;

const LIVENESS_TIMEOUT: Duration = Duration::from_millis(50);
const READINESS_TIMEOUT: Duration = Duration::from_millis(100);

/// Cloneable state used by health probe handlers.
#[derive(Clone)]
pub struct HealthState {
    store: Arc<dyn ReadableEventStore>,
    runtime_initialized: bool,
}

impl HealthState {
    /// Build health state from the store and runtime initialization flag.
    #[must_use]
    pub fn new(store: Arc<dyn ReadableEventStore>, runtime_initialized: bool) -> Self {
        Self {
            store,
            runtime_initialized,
        }
    }
}

/// Liveness probe: validates that the async scheduler can run a trivial task promptly.
pub async fn live() -> StatusCode {
    let check = async {
        let handle = tokio::spawn(async {});
        handle.await.is_ok()
    };

    match timeout(LIVENESS_TIMEOUT, check).await {
        Ok(true) => StatusCode::OK,
        Ok(false) | Err(_) => StatusCode::SERVICE_UNAVAILABLE,
    }
}

/// Readiness probe: validates runtime initialization and store reachability with a no-op read.
pub async fn ready(State(state): State<HealthState>) -> StatusCode {
    if !state.runtime_initialized {
        return StatusCode::SERVICE_UNAVAILABLE;
    }

    let workflow_id = WorkflowId::new_v4();
    match timeout(READINESS_TIMEOUT, state.store.read_history(&workflow_id)).await {
        Ok(Ok(_)) => StatusCode::OK,
        Ok(Err(_)) | Err(_) => StatusCode::SERVICE_UNAVAILABLE,
    }
}