systemprompt_api/services/server/
readiness.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::OnceLock;
3use tokio::sync::broadcast;
4
5static API_READY: AtomicBool = AtomicBool::new(false);
6static READINESS_SENDER: OnceLock<broadcast::Sender<ReadinessEvent>> = OnceLock::new();
7
8#[derive(Debug, Clone, Copy)]
9pub enum ReadinessEvent {
10    ApiReady,
11    ApiShuttingDown,
12}
13
14pub fn init_readiness() -> broadcast::Receiver<ReadinessEvent> {
15    let sender = READINESS_SENDER.get_or_init(|| {
16        let (tx, _) = broadcast::channel(16);
17        tx
18    });
19    sender.subscribe()
20}
21
22pub fn get_readiness_receiver() -> broadcast::Receiver<ReadinessEvent> {
23    READINESS_SENDER
24        .get_or_init(|| {
25            let (tx, _) = broadcast::channel(16);
26            tx
27        })
28        .subscribe()
29}
30
31pub fn signal_ready() {
32    API_READY.store(true, Ordering::SeqCst);
33    if let Some(sender) = READINESS_SENDER.get() {
34        if sender.send(ReadinessEvent::ApiReady).is_err() {
35            tracing::debug!("No readiness receivers subscribed");
36        }
37    }
38}
39
40pub fn signal_shutdown() {
41    API_READY.store(false, Ordering::SeqCst);
42    if let Some(sender) = READINESS_SENDER.get() {
43        if sender.send(ReadinessEvent::ApiShuttingDown).is_err() {
44            tracing::debug!("No readiness receivers subscribed");
45        }
46    }
47}
48
49pub fn is_ready() -> bool {
50    API_READY.load(Ordering::SeqCst)
51}
52
53pub async fn wait_for_ready(timeout_secs: u64) -> bool {
54    if is_ready() {
55        return true;
56    }
57
58    let mut receiver = get_readiness_receiver();
59
60    tokio::time::timeout(std::time::Duration::from_secs(timeout_secs), async {
61        while let Ok(event) = receiver.recv().await {
62            if matches!(event, ReadinessEvent::ApiReady) {
63                return true;
64            }
65        }
66        false
67    })
68    .await
69    .map_err(|_| {
70        tracing::debug!(timeout_secs = timeout_secs, "Readiness wait timed out");
71    })
72    .unwrap_or(false)
73}
systemprompt_api/services/server/readiness.rs

systemprompt_api/services/server/
readiness.rs
