micromegas/servers/
readiness.rs

1use micromegas_ingestion::data_lake_connection::DataLakeConnection;
2use micromegas_tracing::prelude::*;
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, Instant};
5
6/// Shared readiness probe used by the FlightSQL HTTP sidecar.
7///
8/// Probes DB + blob storage under a 2 s timeout and caches success for 1 s
9/// to avoid amplifying load under rapid ALB polling.
10pub struct ReadinessProbe {
11    lake: Arc<DataLakeConnection>,
12    ready_ok_until: Mutex<Option<Instant>>,
13}
14
15impl ReadinessProbe {
16    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
17        Self {
18            lake,
19            ready_ok_until: Mutex::new(None),
20        }
21    }
22
23    pub async fn check_ready(&self) -> bool {
24        let now = Instant::now();
25        {
26            let guard = self.ready_ok_until.lock().expect("readiness cache lock");
27            if let Some(ok_until) = *guard
28                && ok_until > now
29            {
30                return true;
31            }
32        }
33
34        let probe_db = instrument_named!(
35            sqlx::query("SELECT 1").execute(&self.lake.db_pool),
36            "sql_readiness_probe"
37        );
38        let probe_blob = self.lake.blob_storage.probe();
39
40        let result = tokio::time::timeout(Duration::from_secs(2), async {
41            tokio::join!(probe_db, probe_blob)
42        })
43        .await;
44
45        match result {
46            Ok((Ok(_), Ok(()))) => {
47                let mut guard = self.ready_ok_until.lock().expect("readiness cache lock");
48                *guard = Some(Instant::now() + Duration::from_secs(1));
49                true
50            }
51            _ => {
52                let mut guard = self.ready_ok_until.lock().expect("readiness cache lock");
53                *guard = None;
54                false
55            }
56        }
57    }
58}
micromegas/servers/readiness.rs

micromegas/servers/
readiness.rs
