Skip to main content

systemprompt_api/services/server/
readiness.rs

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