aws_smithy_observability/
global.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Utilities for interacting with the currently set `GlobalTelemetryProvider`
7
8use std::{
9    mem,
10    sync::{Arc, LazyLock, RwLock},
11};
12
13use crate::{
14    error::{ErrorKind, GlobalTelemetryProviderError},
15    provider::{GlobalTelemetryProvider, TelemetryProvider},
16    ObservabilityError,
17};
18
19// Statically store the global provider
20static GLOBAL_TELEMETRY_PROVIDER: LazyLock<RwLock<GlobalTelemetryProvider>> =
21    LazyLock::new(|| RwLock::new(GlobalTelemetryProvider::new(TelemetryProvider::default())));
22
23/// Set the current global [TelemetryProvider].
24///
25/// This is meant to be run once at the beginning of an application. Will return an [Err] if the
26/// [RwLock] holding the global [TelemetryProvider] is locked or poisoned.
27pub fn set_telemetry_provider(new_provider: TelemetryProvider) -> Result<(), ObservabilityError> {
28    if let Ok(mut old_provider) = GLOBAL_TELEMETRY_PROVIDER.try_write() {
29        let new_global_provider = GlobalTelemetryProvider::new(new_provider);
30
31        let _ = mem::replace(&mut *old_provider, new_global_provider);
32
33        Ok(())
34    } else {
35        Err(ObservabilityError::new(
36            ErrorKind::Other,
37            GlobalTelemetryProviderError::new("Failed to set global TelemetryProvider."),
38        ))
39    }
40}
41
42/// Get an [Arc] reference to the current global [TelemetryProvider]. Will return an [Err] if the
43/// [RwLock] holding the global [TelemetryProvider] is locked or poisoned.
44pub fn get_telemetry_provider() -> Result<Arc<TelemetryProvider>, ObservabilityError> {
45    if let Ok(tp) = GLOBAL_TELEMETRY_PROVIDER.try_read() {
46        Ok(tp.telemetry_provider().clone())
47    } else {
48        Err(ObservabilityError::new(
49            ErrorKind::Other,
50            GlobalTelemetryProviderError::new("Failed to get global TelemetryProvider"),
51        ))
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use crate::provider::TelemetryProvider;
59    use serial_test::serial;
60
61    // Note: the tests in this module are run serially to prevent them from stepping on each other and poisoning the
62    // RwLock holding the GlobalTelemetryProvider
63    #[test]
64    #[serial]
65    fn can_set_global_telemetry_provider() {
66        let my_provider = TelemetryProvider::default();
67
68        // Set the new counter and get a reference to the old one
69        set_telemetry_provider(my_provider).unwrap();
70    }
71
72    #[test]
73    #[serial]
74    fn can_get_global_telemetry_provider() {
75        let curr_provider = get_telemetry_provider().unwrap();
76
77        // Use the global provider to create an instrument and record a value with it
78        let curr_mp = curr_provider.meter_provider();
79        let curr_meter = curr_mp.get_meter("TestMeter", None);
80        let instrument = curr_meter
81            .create_monotonic_counter("TestMonoCounter")
82            .build();
83        instrument.add(4, None, None);
84    }
85}