init4_bin_base/utils/
otlp.rs

1use crate::utils::from_env::{EnvItemInfo, 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/// Otel configuration. This struct is intended to be loaded from the env vars
65///
66/// The env vars it checks are:
67/// - `OTEL_EXPORTER_OTLP_ENDPOINT` - optional. The endpoint to send traces to,
68///   should be some valid URL. If not specified, then [`OtelConfig::load`]
69///   will return [`None`].
70/// - OTEL_LEVEL - optional. Specifies the minimum [`tracing::Level`] to
71///   export. Defaults to [`tracing::Level::DEBUG`].
72/// - OTEL_TIMEOUT - optional. Specifies the timeout for the exporter in
73///   **milliseconds**. Defaults to 1000ms, which is equivalent to 1 second.
74/// - OTEL_ENVIRONMENT_NAME - optional. Value for the `deployment.environment.
75///   name` resource key according to the OTEL conventions.
76#[derive(Debug, Clone)]
77#[non_exhaustive]
78pub struct OtelConfig {
79    /// The endpoint to send traces to, should be some valid HTTP endpoint for
80    /// OTLP.
81    pub endpoint: Url,
82
83    /// Defaults to DEBUG.
84    pub level: tracing::Level,
85
86    /// Defaults to 1 second. Specified in Milliseconds.
87    pub timeout: Duration,
88
89    /// OTEL convenition `deployment.environment.name`
90    pub environment: String,
91}
92
93impl FromEnv for OtelConfig {
94    type Error = url::ParseError;
95
96    fn inventory() -> Vec<&'static EnvItemInfo> {
97        vec![
98            &EnvItemInfo {
99                var: OTEL_ENDPOINT,
100                description:
101                    "OTLP endpoint to send traces to, a url. If missing, disables OTLP exporting.",
102                optional: true,
103            },
104            &EnvItemInfo {
105                var: OTEL_LEVEL,
106                description: "OTLP level to export, defaults to DEBUG. Permissible values are: TRACE, DEBUG, INFO, WARN, ERROR, OFF",
107                optional: true,
108            },
109            &EnvItemInfo {
110                var: OTEL_TIMEOUT,
111                description: "OTLP timeout in milliseconds",
112                optional: true,
113            },
114            &EnvItemInfo {
115                var: OTEL_ENVIRONMENT,
116                description: "OTLP environment name, a string",
117                optional: true,
118            },
119        ]
120    }
121
122    fn from_env() -> Result<Self, FromEnvErr<Self::Error>> {
123        // load endpoint from env. ignore empty values (shortcut return None), parse, and print the error if any using inspect_err
124        let endpoint = Url::from_env_var(OTEL_ENDPOINT)?;
125
126        let level = tracing::Level::from_env_var(OTEL_LEVEL).unwrap_or(tracing::Level::DEBUG);
127
128        let timeout = Duration::from_env_var(OTEL_TIMEOUT).unwrap_or(Duration::from_millis(1000));
129
130        let environment = String::from_env_var(OTEL_ENVIRONMENT).unwrap_or("unknown".into());
131
132        Ok(Self {
133            endpoint,
134            level,
135            timeout,
136            environment,
137        })
138    }
139}
140
141impl OtelConfig {
142    /// Load from env vars.
143    ///
144    /// The env vars it checks are:
145    /// - `OTEL_EXPORTER_OTLP_ENDPOINT` - optional. The endpoint to send traces
146    ///   to. If missing or unparsable, this function will return [`None`], and
147    ///   OTLP exporting will be disabled.
148    /// - `OTEL_LEVEL` - optional. Specifies the minimum [`tracing::Level`] to
149    ///   export. Defaults to [`tracing::Level::DEBUG`].
150    /// - `OTEL_TIMEOUT` - optional. Specifies the timeout for the exporter in
151    ///   **milliseconds**. Defaults to 1000ms, which is equivalent to 1 second.
152    /// - `OTEL_ENVIRONMENT_NAME` - optional. Value for the
153    ///   `deployment.environment.name` resource key according to the OTEL
154    ///   conventions. Defaults to `"unknown"`.
155    pub fn load() -> Option<Self> {
156        Self::from_env().ok()
157    }
158
159    fn resource(&self) -> Resource {
160        Resource::builder()
161            .with_schema_url(
162                [
163                    KeyValue::new(SERVICE_NAME, env!("CARGO_PKG_NAME")),
164                    KeyValue::new(SERVICE_VERSION, env!("CARGO_PKG_VERSION")),
165                    KeyValue::new(DEPLOYMENT_ENVIRONMENT_NAME, self.environment.clone()),
166                ],
167                SCHEMA_URL,
168            )
169            .build()
170    }
171
172    /// Instantiate a new Otel provider, and start relevant tasks. Return a
173    /// guard that will shut down the provider when dropped.
174    pub fn provider(&self) -> OtelGuard {
175        let exporter = opentelemetry_otlp::SpanExporter::builder()
176            .with_http()
177            .build()
178            .unwrap();
179
180        let provider = SdkTracerProvider::builder()
181            // Customize sampling strategy
182            // If export trace to AWS X-Ray, you can use XrayIdGenerator
183            .with_resource(self.resource())
184            .with_batch_exporter(exporter)
185            .build();
186
187        OtelGuard(provider, self.level)
188    }
189}
190
191#[cfg(test)]
192mod test {
193    use super::*;
194
195    const URL: &str = "http://localhost:4317";
196
197    fn clear_env() {
198        std::env::remove_var(OTEL_ENDPOINT);
199        std::env::remove_var(OTEL_LEVEL);
200        std::env::remove_var(OTEL_TIMEOUT);
201        std::env::remove_var(OTEL_ENVIRONMENT);
202    }
203
204    fn run_clear_env<F>(f: F)
205    where
206        F: FnOnce(),
207    {
208        f();
209        clear_env();
210    }
211
212    #[test]
213    #[serial_test::serial]
214
215    fn test_env_read() {
216        run_clear_env(|| {
217            std::env::set_var(OTEL_ENDPOINT, URL);
218
219            let cfg = OtelConfig::load().unwrap();
220            assert_eq!(cfg.endpoint, URL.parse().unwrap());
221            assert_eq!(cfg.level, tracing::Level::DEBUG);
222            assert_eq!(cfg.timeout, std::time::Duration::from_millis(1000));
223            assert_eq!(cfg.environment, "unknown");
224        })
225    }
226
227    #[test]
228    #[serial_test::serial]
229    fn test_env_read_level() {
230        run_clear_env(|| {
231            std::env::set_var(OTEL_ENDPOINT, URL);
232            std::env::set_var(OTEL_LEVEL, "WARN");
233
234            let cfg = OtelConfig::load().unwrap();
235            assert_eq!(cfg.level, tracing::Level::WARN);
236        })
237    }
238
239    #[test]
240    #[serial_test::serial]
241    fn test_env_read_timeout() {
242        run_clear_env(|| {
243            std::env::set_var(OTEL_ENDPOINT, URL);
244            std::env::set_var(OTEL_TIMEOUT, "500");
245
246            let cfg = OtelConfig::load().unwrap();
247            assert_eq!(cfg.timeout, std::time::Duration::from_millis(500));
248        })
249    }
250
251    #[test]
252    #[serial_test::serial]
253    fn invalid_url() {
254        run_clear_env(|| {
255            std::env::set_var(OTEL_ENDPOINT, "not a url");
256
257            let cfg = OtelConfig::load();
258            assert!(cfg.is_none());
259        })
260    }
261}