runtara_sdk/
registry.rs

1// Copyright (C) 2025 SyncMyOrders Sp. z o.o.
2// SPDX-License-Identifier: AGPL-3.0-or-later
3//! Global SDK registry for the #[durable] macro.
4//!
5//! This module provides global SDK registration so the #[durable] macro
6//! can access the SDK without explicit passing. It also spawns a background
7//! heartbeat task to keep instances alive during long-running operations.
8
9use std::sync::Arc;
10use std::time::Duration;
11
12use once_cell::sync::OnceCell;
13use tokio::sync::Mutex;
14use tokio_util::sync::CancellationToken;
15use tracing::{debug, warn};
16
17use crate::RuntaraSdk;
18
19/// Global storage for the SDK instance.
20static SDK_INSTANCE: OnceCell<Arc<Mutex<RuntaraSdk>>> = OnceCell::new();
21
22/// Cancellation token for the background heartbeat task.
23static HEARTBEAT_CANCEL: OnceCell<CancellationToken> = OnceCell::new();
24
25/// Register an SDK instance globally for use by #[durable] functions.
26///
27/// This should be called once at application startup after creating and
28/// connecting the SDK. If the SDK is configured with a non-zero heartbeat
29/// interval, a background task will be spawned to send periodic heartbeats.
30///
31/// # Panics
32///
33/// Panics if called more than once.
34///
35/// # Example
36///
37/// ```ignore
38/// use runtara_sdk::{RuntaraSdk, register_sdk};
39///
40/// #[tokio::main]
41/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
42///     let sdk = RuntaraSdk::localhost("my-instance", "my-tenant")?;
43///     sdk.connect().await?;
44///
45///     // Register globally for #[durable] functions
46///     // This also starts a background heartbeat task (default: every 30s)
47///     register_sdk(sdk);
48///
49///     // Now #[durable] functions will use this SDK
50///     Ok(())
51/// }
52/// ```
53pub fn register_sdk(sdk: RuntaraSdk) {
54    let heartbeat_interval_ms = sdk.heartbeat_interval_ms();
55
56    let sdk_arc = Arc::new(Mutex::new(sdk));
57
58    if SDK_INSTANCE.set(sdk_arc.clone()).is_err() {
59        panic!("SDK already registered. register_sdk() should only be called once.");
60    }
61
62    // Spawn background heartbeat task if enabled
63    if heartbeat_interval_ms > 0 {
64        let cancel_token = CancellationToken::new();
65        let _ = HEARTBEAT_CANCEL.set(cancel_token.clone());
66
67        let interval = Duration::from_millis(heartbeat_interval_ms);
68
69        tokio::spawn(async move {
70            debug!(
71                interval_ms = heartbeat_interval_ms,
72                "Background heartbeat task started"
73            );
74
75            loop {
76                tokio::select! {
77                    biased;
78
79                    _ = cancel_token.cancelled() => {
80                        debug!("Background heartbeat task cancelled");
81                        break;
82                    }
83
84                    _ = tokio::time::sleep(interval) => {
85                        let sdk_guard = sdk_arc.lock().await;
86                        if let Err(e) = sdk_guard.heartbeat().await {
87                            warn!(error = %e, "Failed to send background heartbeat");
88                        } else {
89                            debug!("Background heartbeat sent");
90                        }
91                    }
92                }
93            }
94        });
95    }
96}
97
98/// Get a reference to the registered SDK.
99///
100/// # Panics
101///
102/// Panics if no SDK has been registered.
103pub fn sdk() -> &'static Arc<Mutex<RuntaraSdk>> {
104    SDK_INSTANCE
105        .get()
106        .expect("No SDK registered. Call register_sdk() at application startup.")
107}
108
109/// Try to get a reference to the registered SDK.
110///
111/// Returns `None` if no SDK has been registered.
112pub fn try_sdk() -> Option<&'static Arc<Mutex<RuntaraSdk>>> {
113    SDK_INSTANCE.get()
114}
115
116/// Stop the background heartbeat task.
117///
118/// This should be called when shutting down the SDK to cleanly stop
119/// the background heartbeat task. It's safe to call this multiple times
120/// or even if no heartbeat task was started.
121///
122/// # Example
123///
124/// ```ignore
125/// use runtara_sdk::{register_sdk, stop_heartbeat, sdk};
126///
127/// // ... at shutdown
128/// stop_heartbeat();
129/// let sdk_guard = sdk().lock().await;
130/// sdk_guard.close().await?;
131/// ```
132pub fn stop_heartbeat() {
133    if let Some(cancel_token) = HEARTBEAT_CANCEL.get() {
134        cancel_token.cancel();
135        debug!("Heartbeat cancellation requested");
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    // Note: These tests can't run in parallel due to global state.
142    // In a real test suite, you'd use a thread-local or per-test registry.
143
144    #[test]
145    fn test_try_sdk_returns_none_initially() {
146        // Can't test this reliably due to global state from other tests
147    }
148}
runtara_sdk/registry.rs

runtara_sdk/
registry.rs
