Skip to main content

superstac_engine/health/
manager.rs

1use parking_lot::Mutex;
2use reqwest::Client;
3use std::{collections::HashMap, sync::Arc};
4use tokio::{task::JoinHandle, time::interval};
5
6use crate::types::SharedStorage;
7use superstac_core::{errors::SuperSTACError, models::catalog::Catalog};
8
9pub struct HealthManager {
10    tasks: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
11    client: Client,
12}
13
14impl HealthManager {
15    pub fn new(client: Client) -> Self {
16        Self {
17            tasks: Arc::new(Mutex::new(HashMap::new())),
18            client,
19        }
20    }
21
22    pub async fn start(
23        &self,
24        storage: SharedStorage,
25        catalogs: Vec<Catalog>,
26    ) -> Result<(), SuperSTACError> {
27        tracing::debug!(catalogs = catalogs.len(), "starting health monitor");
28
29        for catalog in catalogs {
30            let healthy = self.check_once(Arc::clone(&storage), &catalog).await;
31
32            if healthy {
33                // TODO, check settings to decide whether to spawn monitor or not
34                self.spawn_monitor(Arc::clone(&storage), catalog);
35            }
36        }
37
38        Ok(())
39    }
40
41    pub async fn stop_all(&self) {
42        let mut tasks = self.tasks.lock();
43
44        for (_, handle) in tasks.drain() {
45            handle.abort();
46        }
47    }
48
49    pub async fn stop(&self, id: &str) {
50        if let Some(handle) = self.tasks.lock().remove(id) {
51            handle.abort();
52        }
53    }
54
55    pub async fn check_once(&self, storage: SharedStorage, catalog: &Catalog) -> bool {
56        let endpoint = &catalog.health_status.endpoint;
57
58        let range = catalog.settings.healthy_status_code_range;
59
60        let healthy = match self.client.get(endpoint).send().await {
61            Ok(resp) => {
62                let status = resp.status().as_u16();
63
64                range.0 <= status && status <= range.1
65            }
66
67            Err(_) => false,
68        };
69
70        tracing::debug!(catalog = %catalog.id, healthy, "health check");
71
72        {
73            let mut storage = storage.lock();
74
75            if let Err(e) = storage.update_health(&catalog.id, healthy) {
76                tracing::warn!(catalog = %catalog.id, error = %e, "failed to persist health update");
77            }
78        }
79
80        healthy
81    }
82
83    pub fn spawn_monitor(&self, storage: SharedStorage, catalog: Catalog) {
84        let id = catalog.id.clone();
85
86        if self.tasks.lock().contains_key(&id) {
87            return;
88        }
89
90        let endpoint = catalog.health_status.endpoint.clone();
91
92        let range = catalog.settings.healthy_status_code_range;
93
94        let duration = catalog.settings.health_check_strategy.as_duration();
95
96        let client = self.client.clone();
97
98        let task_id = id.clone();
99
100        let handle = tokio::spawn(async move {
101            let mut ticker = interval(duration);
102
103            loop {
104                ticker.tick().await;
105
106                let healthy = match client.get(&endpoint).send().await {
107                    Ok(resp) => {
108                        let status = resp.status().as_u16();
109
110                        range.0 <= status && status <= range.1
111                    }
112
113                    Err(_) => false,
114                };
115
116                tracing::debug!(catalog = %task_id, healthy, "health tick");
117
118                {
119                    let mut storage = storage.lock();
120
121                    if let Err(e) = storage.update_health(&task_id, healthy) {
122                        tracing::warn!(
123                            catalog = %task_id,
124                            error = %e,
125                            "failed to persist health update"
126                        );
127                    }
128                }
129            }
130        });
131
132        tracing::debug!(
133            catalog = %id,
134            interval_secs = duration.as_secs(),
135            "spawned health monitor"
136        );
137
138        self.tasks.lock().insert(id, handle);
139    }
140}