opentelemetry_lambda_tower/
cold_start.rs

1//! Cold start detection for Lambda functions.
2//!
3//! Tracks whether the current invocation is a cold start (first invocation
4//! in this container) following the pattern from the OpenTelemetry Node.js
5//! Lambda instrumentation.
6
7use std::sync::atomic::{AtomicBool, Ordering};
8
9/// Global flag tracking whether we've seen an invocation yet.
10/// Starts as `true` and is set to `false` after the first invocation.
11static IS_COLD_START: AtomicBool = AtomicBool::new(true);
12
13/// Checks if this is a cold start and clears the flag.
14///
15/// Returns `true` on the first call (cold start), and `false` on all
16/// subsequent calls (warm starts).
17///
18/// This function also checks the `AWS_LAMBDA_INITIALIZATION_TYPE` environment
19/// variable. If set to `"provisioned-concurrency"`, the function returns
20/// `false` since provisioned concurrency pre-warms the Lambda container.
21///
22/// # Thread Safety
23///
24/// This function is safe to call from multiple threads. The atomic swap
25/// ensures that exactly one invocation will see `true`.
26///
27/// # Example
28///
29/// ```
30/// use opentelemetry_lambda_tower::check_cold_start;
31///
32/// // First call returns true (cold start)
33/// // Note: In tests, this may return false if other tests ran first
34/// let _ = check_cold_start();
35///
36/// // Subsequent calls return false (warm start)
37/// assert!(!check_cold_start());
38/// ```
39pub fn check_cold_start() -> bool {
40    // Check for provisioned concurrency - never a cold start
41    if std::env::var("AWS_LAMBDA_INITIALIZATION_TYPE")
42        .map(|v| v == "provisioned-concurrency")
43        .unwrap_or(false)
44    {
45        // Still clear the flag but return false
46        IS_COLD_START.store(false, Ordering::SeqCst);
47        return false;
48    }
49
50    // Atomically swap to false and return the previous value
51    // First call: swaps true -> false, returns true
52    // Subsequent calls: swaps false -> false, returns false
53    IS_COLD_START.swap(false, Ordering::SeqCst)
54}
55
56/// Resets the cold start flag for testing purposes.
57///
58/// # Safety
59///
60/// This function should only be used in tests. Using it in production
61/// code will cause incorrect cold start reporting.
62#[cfg(test)]
63pub(crate) fn reset_cold_start_for_testing() {
64    IS_COLD_START.store(true, Ordering::SeqCst);
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use serial_test::serial;
71
72    #[test]
73    #[serial]
74    fn test_cold_start_first_invocation() {
75        reset_cold_start_for_testing();
76
77        // First call should return true
78        assert!(check_cold_start());
79
80        // Second call should return false
81        assert!(!check_cold_start());
82
83        // Third call should also return false
84        assert!(!check_cold_start());
85    }
86
87    #[test]
88    #[serial]
89    fn test_provisioned_concurrency_not_cold_start() {
90        reset_cold_start_for_testing();
91
92        // Set provisioned concurrency environment variable
93        temp_env::with_var(
94            "AWS_LAMBDA_INITIALIZATION_TYPE",
95            Some("provisioned-concurrency"),
96            || {
97                // Should return false even on first call
98                assert!(!check_cold_start());
99            },
100        );
101    }
102
103    #[test]
104    #[serial]
105    fn test_other_initialization_type_is_cold_start() {
106        reset_cold_start_for_testing();
107
108        // Set a different initialization type
109        temp_env::with_var("AWS_LAMBDA_INITIALIZATION_TYPE", Some("on-demand"), || {
110            // Should still be a cold start
111            assert!(check_cold_start());
112        });
113    }
114}
opentelemetry_lambda_tower/cold_start.rs

opentelemetry_lambda_tower/
cold_start.rs
