shadi_telemetry/
lib.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3
4use std::env;
5use std::path::{Path, PathBuf};
6use std::sync::{Once, OnceLock};
7
8use opentelemetry::KeyValue;
9use opentelemetry::trace::TracerProvider;
10use opentelemetry_otlp::WithExportConfig;
11use opentelemetry_sdk::{trace, Resource};
12use tracing_subscriber::layer::SubscriberExt;
13
14static INIT: Once = Once::new();
15static PROVIDER: OnceLock<trace::TracerProvider> = OnceLock::new();
16static FILE_GUARD: OnceLock<tracing_appender::non_blocking::WorkerGuard> = OnceLock::new();
17
18pub fn init(service_name: &str) {
19    INIT.call_once(|| {
20        let config = load_config(service_name);
21
22        if !telemetry_enabled(
23            &config.otlp_endpoint,
24            config.console_enabled,
25            config.file_path.as_deref(),
26        ) {
27            return;
28        }
29
30        let resource = Resource::new(vec![
31            KeyValue::new("service.name", config.service_name),
32            KeyValue::new("service.namespace", "shadi"),
33            KeyValue::new("service.version", env!("CARGO_PKG_VERSION")),
34            KeyValue::new("telemetry.sdk.language", "rust"),
35        ]);
36
37        let otel_layer = if !config.otlp_endpoint.is_empty() {
38            let exporter = opentelemetry_otlp::new_exporter()
39                .http()
40            .with_endpoint(config.otlp_endpoint);
41            let provider = opentelemetry_otlp::new_pipeline()
42                .tracing()
43                .with_exporter(exporter)
44                .with_trace_config(trace::Config::default().with_resource(resource))
45                .install_simple();
46
47            provider.ok().map(|provider| {
48                let _ = PROVIDER.set(provider);
49                let tracer = PROVIDER
50                    .get()
51                    .expect("telemetry provider")
52                    .tracer("shadi.telemetry");
53                tracing_opentelemetry::layer().with_tracer(tracer)
54            })
55        } else {
56            None
57        };
58
59        let file_layer = config.file_path.as_ref().and_then(|path| {
60            let (dir, file_name) = resolve_trace_path(path)?;
61            if std::fs::create_dir_all(&dir).is_err() {
62                return None;
63            }
64
65            let appender = tracing_appender::rolling::never(dir, file_name);
66            let (non_blocking, guard) = tracing_appender::non_blocking(appender);
67            let _ = FILE_GUARD.set(guard);
68
69            Some(
70                tracing_subscriber::fmt::layer()
71                    .json()
72                    .with_current_span(true)
73                    .with_span_list(true)
74                    .with_ansi(false)
75                    .with_writer(non_blocking),
76            )
77        });
78
79        let fmt_layer = config.console_enabled.then(|| tracing_subscriber::fmt::layer());
80
81        let subscriber = tracing_subscriber::registry()
82            .with(otel_layer)
83            .with(fmt_layer)
84            .with(file_layer);
85
86        let _ = tracing::subscriber::set_global_default(subscriber);
87    });
88}
89
90fn parse_bool_env(key: &str) -> bool {
91    let value = env::var(key).unwrap_or_default().trim().to_ascii_lowercase();
92    matches!(value.as_str(), "1" | "true" | "yes")
93}
94
95fn normalize_file_path(value: &str) -> Option<String> {
96    let trimmed = value.trim();
97    if trimmed.is_empty() {
98        None
99    } else {
100        Some(trimmed.to_string())
101    }
102}
103
104fn resolve_trace_path(path: &str) -> Option<(PathBuf, String)> {
105    let trace_path = Path::new(path);
106    let dir = trace_path.parent().unwrap_or_else(|| Path::new(".")).to_path_buf();
107    let file_name = trace_path
108        .file_name()
109        .and_then(|name| name.to_str())
110        .unwrap_or("traces.jsonl")
111        .to_string();
112    Some((dir, file_name))
113}
114
115#[derive(Debug, Clone)]
116struct TelemetryConfig {
117    otlp_endpoint: String,
118    console_enabled: bool,
119    file_path: Option<String>,
120    service_name: String,
121}
122
123fn load_config(default_service_name: &str) -> TelemetryConfig {
124    let otlp_endpoint = env::var("OTEL_EXPORTER_OTLP_ENDPOINT").unwrap_or_default();
125    let console_enabled = parse_bool_env("SHADI_OTEL_CONSOLE");
126    let file_path = env::var("SHADI_OTEL_FILE")
127        .ok()
128        .and_then(|value| normalize_file_path(&value));
129    let service_name = resolve_service_name(default_service_name);
130
131    TelemetryConfig {
132        otlp_endpoint,
133        console_enabled,
134        file_path,
135        service_name,
136    }
137}
138
139fn resolve_service_name(default_service_name: &str) -> String {
140    env::var("OTEL_SERVICE_NAME")
141        .ok()
142        .filter(|value| !value.trim().is_empty())
143        .unwrap_or_else(|| default_service_name.to_string())
144}
145
146fn telemetry_enabled(otlp_endpoint: &str, console_enabled: bool, file_path: Option<&str>) -> bool {
147    !otlp_endpoint.is_empty() || console_enabled || file_path.is_some()
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use std::sync::Mutex;
154    use std::time::{SystemTime, UNIX_EPOCH};
155
156    static ENV_LOCK: Mutex<()> = Mutex::new(());
157
158    #[test]
159    fn parse_bool_env_accepts_truthy_values() {
160        let _guard = ENV_LOCK.lock().unwrap();
161        std::env::set_var("SHADI_OTEL_CONSOLE", "1");
162        assert!(parse_bool_env("SHADI_OTEL_CONSOLE"));
163        std::env::set_var("SHADI_OTEL_CONSOLE", "true");
164        assert!(parse_bool_env("SHADI_OTEL_CONSOLE"));
165        std::env::set_var("SHADI_OTEL_CONSOLE", "yes");
166        assert!(parse_bool_env("SHADI_OTEL_CONSOLE"));
167        std::env::set_var("SHADI_OTEL_CONSOLE", "no");
168        assert!(!parse_bool_env("SHADI_OTEL_CONSOLE"));
169        std::env::remove_var("SHADI_OTEL_CONSOLE");
170    }
171
172    #[test]
173    fn normalize_file_path_trims_and_rejects_empty() {
174        assert_eq!(normalize_file_path(""), None);
175        assert_eq!(normalize_file_path("   "), None);
176        assert_eq!(normalize_file_path("/tmp/trace.jsonl"), Some("/tmp/trace.jsonl".to_string()));
177        assert_eq!(normalize_file_path("  ./traces.jsonl "), Some("./traces.jsonl".to_string()));
178    }
179
180    #[test]
181    fn resolve_trace_path_builds_dir_and_name() {
182        let (dir, file) = resolve_trace_path("/tmp/trace.jsonl").expect("path");
183        assert_eq!(dir, PathBuf::from("/tmp"));
184        assert_eq!(file, "trace.jsonl");
185    }
186
187    #[test]
188    fn resolve_trace_path_defaults_parent_for_bare_filename() {
189        let (dir, file) = resolve_trace_path("trace.jsonl").expect("path");
190        assert_eq!(dir, PathBuf::new());
191        assert_eq!(file, "trace.jsonl");
192    }
193
194    #[test]
195    fn load_config_reads_env_vars() {
196        let _guard = ENV_LOCK.lock().unwrap();
197        std::env::set_var("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318");
198        std::env::set_var("SHADI_OTEL_CONSOLE", "true");
199        std::env::set_var("SHADI_OTEL_FILE", "./traces.jsonl");
200        std::env::set_var("OTEL_SERVICE_NAME", "shadi-test");
201
202        let config = load_config("default");
203        assert_eq!(config.otlp_endpoint, "http://localhost:4318");
204        assert!(config.console_enabled);
205        assert_eq!(config.file_path, Some("./traces.jsonl".to_string()));
206        assert_eq!(config.service_name, "shadi-test");
207
208        std::env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
209        std::env::remove_var("SHADI_OTEL_CONSOLE");
210        std::env::remove_var("SHADI_OTEL_FILE");
211        std::env::remove_var("OTEL_SERVICE_NAME");
212    }
213
214    #[test]
215    fn resolve_service_name_defaults() {
216        let _guard = ENV_LOCK.lock().unwrap();
217        std::env::remove_var("OTEL_SERVICE_NAME");
218        let name = resolve_service_name("default-service");
219        assert_eq!(name, "default-service");
220    }
221
222    #[test]
223    fn telemetry_enabled_requires_any_sink() {
224        assert!(!telemetry_enabled("", false, None));
225        assert!(telemetry_enabled("http://localhost:4318", false, None));
226        assert!(telemetry_enabled("", true, None));
227        assert!(telemetry_enabled("", false, Some("/tmp/trace.jsonl")));
228    }
229
230    #[test]
231    fn init_configures_console_and_file_layers() {
232        let _guard = ENV_LOCK.lock().expect("env lock");
233        let nanos = SystemTime::now()
234            .duration_since(UNIX_EPOCH)
235            .expect("clock")
236            .as_nanos();
237        let file_path = std::env::temp_dir().join(format!("shadi-traces-{nanos}.jsonl"));
238
239        std::env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
240        std::env::set_var("SHADI_OTEL_CONSOLE", "true");
241        std::env::set_var("SHADI_OTEL_FILE", file_path.to_string_lossy().to_string());
242        std::env::set_var("OTEL_SERVICE_NAME", "shadi-telemetry-test");
243
244        init("default-service");
245
246        std::env::remove_var("OTEL_EXPORTER_OTLP_ENDPOINT");
247        std::env::remove_var("SHADI_OTEL_CONSOLE");
248        std::env::remove_var("SHADI_OTEL_FILE");
249        std::env::remove_var("OTEL_SERVICE_NAME");
250    }
251}
shadi_telemetry/lib.rs

shadi_telemetry/
lib.rs
