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#[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 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 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
72pub 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}