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

init4_bin_base/utils/
otlp.rs
