vls_util/
observability.rs

1use std::error::Error;
2use std::path::Path;
3
4use tracing::Level;
5use tracing_appender::non_blocking::{NonBlocking, WorkerGuard};
6use tracing_appender::rolling;
7use tracing_subscriber::fmt;
8use tracing_subscriber::layer::SubscriberExt;
9
10use tracing_subscriber::util::SubscriberInitExt;
11use tracing_subscriber::EnvFilter;
12
13#[cfg(feature = "otlp")]
14use super::otlp::new_tracer_provider;
15#[cfg(feature = "otlp")]
16use opentelemetry::{global, trace::TracerProvider as _};
17#[cfg(feature = "otlp")]
18use tracing_opentelemetry::OpenTelemetryLayer;
19
20/** create a non blocking tracing file appender with daily rolling */
21pub fn setup_file_appender<P: AsRef<Path>>(datadir: P, who: &str) -> (NonBlocking, WorkerGuard) {
22    let file_appender = rolling::never(datadir.as_ref(), format!("{}.log", who));
23
24    tracing_appender::non_blocking(file_appender)
25}
26
27/** create a RUST_LOG env based log filter with default level set to info */
28pub fn env_filter() -> EnvFilter {
29    EnvFilter::builder().with_default_directive(Level::INFO.into()).from_env_lossy()
30}
31
32/**
33 * Initialize tracing-subscriber with env filter based on RUST_LOG env variable.
34 * fmt layer is used to print logs to stdout.
35 * OpenTelemetryLayer is used to export logs to the configured OTLP endpoint.
36 * fmt layer with custom writer is used to write logs to log file in datadir.
37*/
38pub fn init_tracing_subscriber<P: AsRef<Path>>(
39    datadir: P,
40    who: &str,
41) -> Result<OtelGuard, Box<dyn Error>> {
42    let (file_writer, file_guard) = setup_file_appender(datadir, who);
43
44    let format = fmt::format()
45        .with_level(true)
46        .with_ansi(true)
47        .with_target(false)
48        .with_source_location(true)
49        .compact();
50
51    let stdout_layer = fmt::layer().event_format(format.clone()).with_writer(std::io::stdout);
52    let file_layer = fmt::layer().event_format(format).with_writer(file_writer);
53    let env_filter = env_filter();
54
55    let default_subscriber =
56        tracing_subscriber::registry().with(stdout_layer).with(file_layer).with(env_filter);
57
58    #[cfg(feature = "otlp")]
59    let default_subscriber = {
60        let tracer = new_tracer_provider()?.tracer(who.to_string());
61        let otlp_trace_layer = OpenTelemetryLayer::new(tracer);
62        default_subscriber.with(otlp_trace_layer)
63    };
64
65    match default_subscriber.try_init() {
66        Ok(_) => Ok(OtelGuard::new(file_guard)),
67        Err(err) => Err(Box::new(err)),
68    }
69}
70
71pub struct OtelGuard {
72    _file_appender_guard: WorkerGuard,
73}
74
75impl OtelGuard {
76    pub fn new(file_appender_guard: WorkerGuard) -> Self {
77        Self { _file_appender_guard: file_appender_guard }
78    }
79}
80
81impl Drop for OtelGuard {
82    /** Shut down the current tracer provider. This will invoke the shutdown method on all span processors. Span processors should export remaining spans before return. */
83    fn drop(&mut self) {
84        #[cfg(feature = "otlp")]
85        global::shutdown_tracer_provider();
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use crate::observability::env_filter;
92    use std::time::Duration;
93    use tokio::time::sleep;
94    use tracing_subscriber::{fmt, layer::SubscriberExt};
95
96    #[tokio::test]
97    async fn test_setup_file_appender() {
98        std::env::set_var("RUST_LOG", "info");
99
100        let temp_dir = std::env::temp_dir();
101        let file_path = temp_dir.join(format!("test.log"));
102
103        let handle = tokio::spawn(async move {
104            let (file_writer, _file_guard) = super::setup_file_appender(temp_dir, "test");
105
106            let subscriber = tracing_subscriber::registry()
107                .with(fmt::layer().with_writer(file_writer))
108                .with(env_filter());
109
110            tracing::subscriber::set_global_default(subscriber)
111                .expect("setting default subscriber failed");
112
113            tracing::info!("test random date: 11/08/2001");
114
115            sleep(Duration::from_millis(10000)).await;
116        });
117
118        let _ = handle.await;
119
120        assert_eq!(file_path.exists(), true);
121        let contents = std::fs::read_to_string(&file_path).expect("failed to read file");
122        assert_eq!(contents.contains("test random date: 11/08/2001"), true);
123    }
124}