Skip to main content

opentelemetry_xray/telemetry/
provider.rs

1/* Copyright © 2025, CosmicMind, Inc. */
2
3use opentelemetry::trace::TracerProvider as _;
4use opentelemetry::{global, KeyValue};
5use opentelemetry_aws::trace::{XrayIdGenerator, XrayPropagator};
6use opentelemetry_sdk::trace::{self, SdkTracerProvider};
7use opentelemetry_sdk::Resource;
8use tracing_opentelemetry::OpenTelemetryLayer;
9use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
10
11use crate::error::TelemetryError;
12use crate::exporter::JsonExporter;
13
14/// TelemetryConfig holds configuration information for initializing tracing.
15pub struct TelemetryConfig {
16    /// The name of your service.
17    pub service_name: String,
18    /// The version of your service.
19    pub service_version: String,
20    /// The deployment environment (e.g., "development", "production").
21    pub deployment_env: String,
22    /// The log level (e.g., "info", "debug").
23    pub log_level: String,
24}
25
26#[derive(Debug)]
27pub struct Telemetry {
28    tracer_provider: SdkTracerProvider,
29}
30
31impl Telemetry {
32    /// Creates a new Telemetry instance with a tracer provider.
33    pub fn new(tracer_provider: SdkTracerProvider) -> Self {
34        Self { tracer_provider }
35    }
36
37    /// Initialize with default configuration.
38    pub fn init(config: TelemetryConfig) -> Result<Self, TelemetryError> {
39        // Set the global X-Ray propagator.
40        Self::init_propagator();
41        // Build the tracer provider with the provided config.
42        let tracer_provider = Self::init_provider(&config)?;
43        // Set up the global subscriber using the validated config.
44        Self::init_subscriber(&config, &tracer_provider)?;
45
46        Ok(Self { tracer_provider })
47    }
48
49    /// Initialize the global X-Ray propagator.
50    pub fn init_propagator() {
51        global::set_text_map_propagator(XrayPropagator::default());
52    }
53
54    /// Initialize a tracer provider with the given configuration.
55    /// All configuration checks are done here.
56    pub fn init_provider(config: &TelemetryConfig) -> Result<SdkTracerProvider, TelemetryError> {
57        if config.service_name.is_empty() {
58            return Err(TelemetryError::InvalidConfiguration(
59                "service_name cannot be empty".into(),
60            ));
61        }
62        if config.service_version.is_empty() {
63            return Err(TelemetryError::InvalidConfiguration(
64                "service_version cannot be empty".into(),
65            ));
66        }
67        if config.deployment_env.is_empty() {
68            return Err(TelemetryError::InvalidConfiguration(
69                "deployment_env cannot be empty".into(),
70            ));
71        }
72
73        let resource = Resource::builder().with_attributes(vec![
74            KeyValue::new("service.name", config.service_name.clone()),
75            KeyValue::new("service.version", config.service_version.clone()),
76            KeyValue::new("deployment.environment", config.deployment_env.clone()),
77        ]);
78
79        let provider = SdkTracerProvider::builder()
80            .with_id_generator(XrayIdGenerator::default())
81            .with_sampler(trace::Sampler::AlwaysOn)
82            .with_simple_exporter(JsonExporter::new(config.deployment_env.clone()))
83            .with_resource(resource.build())
84            .build();
85
86        Ok(provider)
87    }
88
89    /// Initialize the global subscriber with the given configuration and tracer provider.
90    pub fn init_subscriber(
91        config: &TelemetryConfig,
92        provider: &SdkTracerProvider,
93    ) -> Result<(), TelemetryError> {
94        let tracer = provider.tracer(config.service_name.clone());
95        let subscriber = Registry::default()
96            .with(OpenTelemetryLayer::new(tracer))
97            .with(EnvFilter::new(&config.log_level));
98
99        tracing::subscriber::set_global_default(subscriber)
100            .map_err(|e| TelemetryError::SubscriberError(e.to_string()))
101    }
102
103    /// Get a reference to the underlying tracer provider.
104    pub fn provider(&self) -> &SdkTracerProvider {
105        &self.tracer_provider
106    }
107
108    /// Shuts down the tracer provider, flushing any remaining spans.
109    pub fn shutdown(&self) {
110        let _ = self.tracer_provider.shutdown();
111    }
112}
113
114impl Drop for Telemetry {
115    fn drop(&mut self) {
116        let _ = self.tracer_provider.shutdown();
117    }
118}