vls_proxy/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 crate::util::otlp::new_tracer;
15#[cfg(feature = "otlp")]
16use opentelemetry::global;
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 stdout_layer = fmt::layer().with_writer(std::io::stdout);
45    let file_layer = fmt::layer().with_writer(file_writer);
46    let env_filter = env_filter();
47
48    let default_subscriber =
49        tracing_subscriber::registry().with(stdout_layer).with(file_layer).with(env_filter);
50
51    #[cfg(feature = "otlp")]
52    let default_subscriber = {
53        let tracer = new_tracer()?;
54        let otlp_layer = Some(OpenTelemetryLayer::new(tracer));
55
56        default_subscriber.with(Some(otlp_layer))
57    };
58
59    match default_subscriber.try_init() {
60        Ok(_) => Ok(OtelGuard::new(file_guard)),
61        Err(err) => Err(Box::new(err)),
62    }
63}
64
65pub struct OtelGuard {
66    _file_appender_guard: WorkerGuard,
67}
68
69impl OtelGuard {
70    pub fn new(file_appender_guard: WorkerGuard) -> Self {
71        Self { _file_appender_guard: file_appender_guard }
72    }
73}
74
75impl Drop for OtelGuard {
76    /** Shut down the current tracer provider. This will invoke the shutdown method on all span processors. Span processors should export remaining spans before return. */
77    fn drop(&mut self) {
78        #[cfg(feature = "otlp")]
79        global::shutdown_tracer_provider();
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use crate::util::observability::env_filter;
86    use std::time::Duration;
87    use tokio::time::sleep;
88    use tracing_subscriber::{fmt, layer::SubscriberExt};
89
90    #[tokio::test]
91    async fn test_setup_file_appender() {
92        std::env::set_var("RUST_LOG", "info");
93
94        let temp_dir = std::env::temp_dir();
95        let file_path = temp_dir.join(format!("test.log"));
96
97        let handle = tokio::spawn(async move {
98            let (file_writer, _file_guard) = super::setup_file_appender(temp_dir, "test");
99
100            let subscriber = tracing_subscriber::registry()
101                .with(fmt::layer().with_writer(file_writer))
102                .with(env_filter());
103
104            tracing::subscriber::set_global_default(subscriber)
105                .expect("setting default subscriber failed");
106
107            tracing::info!("test random date: 11/08/2001");
108
109            sleep(Duration::from_millis(10000)).await;
110        });
111
112        let _ = handle.await;
113
114        assert_eq!(file_path.exists(), true);
115        let contents = std::fs::read_to_string(&file_path).expect("failed to read file");
116        assert_eq!(contents.contains("test random date: 11/08/2001"), true);
117    }
118}