saluki_health/
api.rs

1use std::sync::{Arc, Mutex};
2
3use saluki_api::{
4    extract::State,
5    response::IntoResponse,
6    routing::{get, Router},
7    APIHandler, StatusCode,
8};
9use saluki_common::collections::FastHashMap;
10use serde::Serialize;
11
12use crate::RegistryState;
13
14#[derive(Serialize)]
15struct SimpleComponentState {
16    live: bool,
17    ready: bool,
18}
19
20/// State used for the healthy registry API handler.
21#[derive(Clone)]
22pub struct HealthRegistryState {
23    inner: Arc<Mutex<RegistryState>>,
24}
25
26impl HealthRegistryState {
27    fn get_response(&self, check_ready: bool) -> (StatusCode, String) {
28        // We specifically do this all here because we want to ensure the state is locked for both determining if the
29        // ready/live state is passing/failing, as well as serializing that same state data to JSON, to avoid
30        // inconsistencies between the two.
31
32        let health_state = {
33            let inner = self.inner.lock().unwrap();
34            let mut health_state = FastHashMap::default();
35
36            for component_state in &inner.component_state {
37                let simple_state = SimpleComponentState {
38                    live: component_state.is_live(),
39                    ready: component_state.is_ready(),
40                };
41                health_state.insert(component_state.name.clone(), simple_state);
42            }
43
44            health_state
45        };
46
47        // Run through the collected component health states to determine our overall passing/failing status
48        // depending on what endpoint this is being called for.
49        let passing = health_state
50            .values()
51            .all(|health| if check_ready { health.ready } else { health.live });
52        let status = if passing {
53            StatusCode::OK
54        } else {
55            StatusCode::SERVICE_UNAVAILABLE
56        };
57
58        let rendered = serde_json::to_string(&health_state).unwrap();
59
60        (status, rendered)
61    }
62
63    fn get_ready_response(&self) -> (StatusCode, String) {
64        self.get_response(true)
65    }
66
67    fn get_live_response(&self) -> (StatusCode, String) {
68        self.get_response(false)
69    }
70}
71
72/// An API handler for reporting the health of all components.
73///
74/// This handler exposes two main routes -- `/health/ready` and `/health/live` -- which return the overall readiness and
75/// liveness of all registered components, respectively. Each route will return a successful response (200 OK) if all
76/// components are ready/live, or a failure response (503 Service Unavailable) if any (or all) of the components are not
77/// ready/live, respectively.
78///
79/// In both cases, the response body will be a JSON object with all registered components, each with their individual
80/// readiness and liveness status, as well as the response latency (in seconds) for the component to respond to the
81/// latest liveness probe.
82pub struct HealthAPIHandler {
83    state: HealthRegistryState,
84}
85
86impl HealthAPIHandler {
87    pub(crate) fn from_state(inner: Arc<Mutex<RegistryState>>) -> Self {
88        Self {
89            state: HealthRegistryState { inner },
90        }
91    }
92
93    async fn ready_handler(State(state): State<HealthRegistryState>) -> impl IntoResponse {
94        state.get_ready_response()
95    }
96
97    async fn live_handler(State(state): State<HealthRegistryState>) -> impl IntoResponse {
98        state.get_live_response()
99    }
100}
101
102impl APIHandler for HealthAPIHandler {
103    type State = HealthRegistryState;
104
105    fn generate_initial_state(&self) -> Self::State {
106        self.state.clone()
107    }
108
109    fn generate_routes(&self) -> Router<Self::State> {
110        Router::new()
111            .route("/ready", get(Self::ready_handler))
112            .route("/live", get(Self::live_handler))
113    }
114}