init4_bin_base/utils/
otlp.rs

1use crate::utils::from_env::{FromEnv, FromEnvErr, FromEnvVar};
2use opentelemetry::{trace::TracerProvider, KeyValue};
3use opentelemetry_sdk::trace::SdkTracerProvider;
4use opentelemetry_sdk::Resource;
5use opentelemetry_semantic_conventions::{
6    attribute::{DEPLOYMENT_ENVIRONMENT_NAME, SERVICE_NAME, SERVICE_VERSION},
7    SCHEMA_URL,
8};
9use std::time::Duration;
10use tracing::level_filters::LevelFilter;
11use tracing_subscriber::Layer;
12use url::Url;
13
14const OTEL_ENDPOINT: &str = "OTEL_EXPORTER_OTLP_ENDPOINT";
15const OTEL_LEVEL: &str = "OTEL_LEVEL";
16const OTEL_TIMEOUT: &str = "OTEL_TIMEOUT";
17const OTEL_ENVIRONMENT: &str = "OTEL_ENVIRONMENT_NAME";
18
19/// Drop guard for the Otel provider. This will shutdown the provider when
20/// dropped, and generally should be held for the lifetime of the `main`
21/// function.
22///
23/// ```
24/// # use init4_bin_base::utils::otlp::{OtelConfig, OtelGuard};
25/// # fn test() {
26/// use init4_bin_base::utils::from_env::FromEnv;
27/// fn main() {
28///     let cfg = OtelConfig::from_env().unwrap();
29///     let guard = cfg.provider();
30///     // do stuff
31///     // drop the guard when the program is done
32/// }
33/// # }
34/// ```
35#[derive(Debug)]
36pub struct OtelGuard(SdkTracerProvider, tracing::Level);
37
38impl OtelGuard {
39    /// Get a tracer from the provider.
40    fn tracer(&self, s: &'static str) -> opentelemetry_sdk::trace::Tracer {
41        self.0.tracer(s)
42    }
43
44    /// Create a filtered tracing layer.
45    pub fn layer<S>(&self) -> impl Layer<S>
46    where
47        S: tracing::Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
48    {
49        let tracer = self.tracer("tracing-otel-subscriber");
50        tracing_opentelemetry::layer()
51            .with_tracer(tracer)
52            .with_filter(LevelFilter::from_level(self.1))
53    }
54}
55
56impl Drop for OtelGuard {
57    fn drop(&mut self) {
58        if let Err(err) = self.0.shutdown() {
59            eprintln!("{err:?}");
60        }
61    }
62}
63
64/// Otlp parse error.
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct OtlpParseError(String);
67
68impl From<String> for OtlpParseError {
69    fn from(s: String) -> Self {
70        Self(s)
71    }
72}
73
74impl core::fmt::Display for OtlpParseError {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        f.write_str(&format!("invalid OTLP protocol: {}", self.0))
77    }
78}
79
80impl core::error::Error for OtlpParseError {}
81
82/// Otel configuration. This struct is intended to be loaded from the env vars
83///
84/// The env vars it checks are:
85/// - `OTEL_EXPORTER_OTLP_ENDPOINT` - optional. The endpoint to send traces to,
86///   should be some valid URL. If not specified, then [`OtelConfig::load`]
87///   will return [`None`].
88/// - OTEL_LEVEL - optional. Specifies the minimum [`tracing::Level`] to
89///   export. Defaults to [`tracing::Level::DEBUG`].
90/// - OTEL_TIMEOUT - optional. Specifies the timeout for the exporter in
91///   **milliseconds**. Defaults to 1000ms, which is equivalent to 1 second.
92/// - OTEL_ENVIRONMENT_NAME - optional. Value for the `deployment.environment.
93///   name` resource key according to the OTEL conventions.
94#[derive(Debug, Clone)]
95#[non_exhaustive]
96pub struct OtelConfig {
97    /// The endpoint to send traces to, should be some valid HTTP endpoint for
98    /// OTLP.
99    pub endpoint: Url,
100
101    /// Defaults to DEBUG.
102    pub level: tracing::Level,
103    /// Defaults to 1 second. Specified in Milliseconds.
104    pub timeout: Duration,
105
106    /// OTEL convenition `deployment.environment.name`
107    pub environment: String,
108}
109
110impl FromEnv for OtelConfig {
111    type Error = url::ParseError;
112
113    fn from_env() -> Result<Self, FromEnvErr<Self::Error>> {
114        // load endpoint from env. ignore empty values (shortcut return None), parse, and print the error if any using inspect_err
115        let endpoint = Url::from_env_var(OTEL_ENDPOINT).inspect_err(|e| eprintln!("{e}"))?;
116
117        let level = tracing::Level::from_env_var(OTEL_LEVEL).unwrap_or(tracing::Level::DEBUG);
118
119        let timeout = Duration::from_env_var(OTEL_TIMEOUT).unwrap_or(Duration::from_millis(1000));
120
121        let environment = String::from_env_var(OTEL_ENVIRONMENT).unwrap_or("unknown".into());
122
123        Ok(Self {
124            endpoint,
125            level,
126            timeout,
127            environment,
128        })
129    }
130}
131
132impl OtelConfig {
133    /// Load from env vars.
134    ///
135    /// The env vars it checks are:
136    /// - `OTEL_EXPORTER_OTLP_ENDPOINT` - optional. The endpoint to send traces
137    ///    to. If missing or unparsable, this function will return [`None`], and
138    ///    OTLP exporting will be disabled.
139    /// - `OTEL_LEVEL` - optional. Specifies the minimum [`tracing::Level`] to
140    ///   export. Defaults to [`tracing::Level::DEBUG`].
141    /// - `OTEL_TIMEOUT` - optional. Specifies the timeout for the exporter in
142    ///   **milliseconds**. Defaults to 1000ms, which is equivalent to 1 second.
143    /// - `OTEL_ENVIRONMENT_NAME` - optional. Value for the
144    ///   `deployment.environment.name` resource key according to the OTEL
145    ///   conventions. Defaults to `"unknown"`.
146    pub fn load() -> Option<Self> {
147        Self::from_env().ok()
148    }
149
150    fn resource(&self) -> Resource {
151        Resource::builder()
152            .with_schema_url(
153                [
154                    KeyValue::new(SERVICE_NAME, env!("CARGO_PKG_NAME")),
155                    KeyValue::new(SERVICE_VERSION, env!("CARGO_PKG_VERSION")),
156                    KeyValue::new(DEPLOYMENT_ENVIRONMENT_NAME, self.environment.clone()),
157                ],
158                SCHEMA_URL,
159            )
160            .build()
161    }
162
163    /// Instantiate a new Otel provider, and start relevant tasks. Return a
164    /// guard that will shut down the provider when dropped.
165    pub fn provider(&self) -> OtelGuard {
166        let exporter = opentelemetry_otlp::SpanExporter::builder()
167            .with_http()
168            .build()
169            .unwrap();
170
171        let provider = SdkTracerProvider::builder()
172            // Customize sampling strategy
173            // If export trace to AWS X-Ray, you can use XrayIdGenerator
174            .with_resource(self.resource())
175            .with_batch_exporter(exporter)
176            .build();
177
178        OtelGuard(provider, self.level)
179    }
180}
181
182#[cfg(test)]
183mod test {
184    use super::*;
185
186    const URL: &str = "http://localhost:4317";
187
188    fn clear_env() {
189        std::env::remove_var(OTEL_ENDPOINT);
190        std::env::remove_var(OTEL_LEVEL);
191        std::env::remove_var(OTEL_TIMEOUT);
192        std::env::remove_var(OTEL_ENVIRONMENT);
193    }
194
195    fn run_clear_env<F>(f: F)
196    where
197        F: FnOnce(),
198    {
199        f();
200        clear_env();
201    }
202
203    #[test]
204    #[serial_test::serial]
205
206    fn test_env_read() {
207        run_clear_env(|| {
208            std::env::set_var(OTEL_ENDPOINT, URL);
209
210            let cfg = OtelConfig::load().unwrap();
211            assert_eq!(cfg.endpoint, URL.parse().unwrap());
212            assert_eq!(cfg.level, tracing::Level::DEBUG);
213            assert_eq!(cfg.timeout, std::time::Duration::from_millis(1000));
214            assert_eq!(cfg.environment, "unknown");
215        })
216    }
217
218    #[test]
219    #[serial_test::serial]
220    fn test_env_read_level() {
221        run_clear_env(|| {
222            std::env::set_var(OTEL_ENDPOINT, URL);
223            std::env::set_var(OTEL_LEVEL, "WARN");
224
225            let cfg = OtelConfig::load().unwrap();
226            assert_eq!(cfg.level, tracing::Level::WARN);
227        })
228    }
229
230    #[test]
231    #[serial_test::serial]
232    fn test_env_read_timeout() {
233        run_clear_env(|| {
234            std::env::set_var(OTEL_ENDPOINT, URL);
235            std::env::set_var(OTEL_TIMEOUT, "500");
236
237            let cfg = OtelConfig::load().unwrap();
238            assert_eq!(cfg.timeout, std::time::Duration::from_millis(500));
239        })
240    }
241
242    #[test]
243    #[serial_test::serial]
244    fn invalid_url() {
245        run_clear_env(|| {
246            std::env::set_var(OTEL_ENDPOINT, "not a url");
247
248            let cfg = OtelConfig::load();
249            assert!(cfg.is_none());
250        })
251    }
252}