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