utils::logging: implement tracing_chrome & tracing_flame support

libs/utils/Cargo.toml

@@ -61,6 +61,8 @@ const_format.workspace = true
 # to use tokio channels as streams, this is faster to compile than async_stream
 # why is it only here? no other crate should use it, streams are rarely needed.
 tokio-stream = { version = "0.1.14" }
+tracing-chrome = "0.7.1"
+tracing-flame = "0.2.0"
 
 serde_path_to_error.workspace = true
 

libs/utils/src/lib.rs

@@ -27,6 +27,7 @@ pub mod auth;
 pub mod id;
 
 mod hex;
+
 pub use hex::Hex;
 
 // http endpoint utils

libs/utils/src/logging.rs

@@ -1,4 +1,8 @@
-use std::str::FromStr;
+use std::{
+    io::BufWriter,
+    str::FromStr,
+    sync::{Arc, Mutex},
+};
 
 use anyhow::Context;
 use metrics::{IntCounter, IntCounterVec};
@@ -98,11 +102,27 @@ pub enum Output {
     Stderr,
 }
 
+/// Keep alive and drop it before the program terminates.
+#[allow(dead_code)] // We need to store the `Arc<>` for drop semantics.
+#[must_use]
+pub struct FlushGuard(Arc<Mutex<FlushGuardInner>>);
+
+struct FlushGuardInner {
+    _tracing_chrome_layer: Option<tracing_chrome::FlushGuard>,
+    _tracing_flame_layer: Option<tracing_flame::FlushGuard<BufWriter<std::fs::File>>>,
+}
+
+impl From<FlushGuardInner> for FlushGuard {
+    fn from(value: FlushGuardInner) -> Self {
+        Self(Arc::new(Mutex::new(value)))
+    }
+}
+
 pub fn init(
     log_format: LogFormat,
     tracing_error_layer_enablement: TracingErrorLayerEnablement,
     output: Output,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<FlushGuard> {
     // We fall back to printing all spans at info-level or above if
     // the RUST_LOG environment variable is not set.
     let rust_log_env_filter = || {
@@ -113,8 +133,28 @@ pub fn init(
     // NB: the order of the with() calls does not matter.
     // See https://docs.rs/tracing-subscriber/0.3.16/tracing_subscriber/layer/index.html#per-layer-filtering
     use tracing_subscriber::prelude::*;
-    let r = tracing_subscriber::registry();
-    let r = r.with({
+
+    // https://users.rust-lang.org/t/how-can-i-init-tracing-registry-dynamically-with-multiple-outputs/94307/6
+    #[derive(Default)]
+    struct LayerStack {
+        layers:
+            Option<Box<dyn tracing_subscriber::Layer<tracing_subscriber::Registry> + Sync + Send>>,
+    }
+    impl LayerStack {
+        fn add_layer<L>(&mut self, new_layer: L)
+        where
+            L: tracing_subscriber::Layer<tracing_subscriber::Registry> + Send + Sync,
+        {
+            let new = match self.layers.take() {
+                Some(layers) => Some(layers.and_then(new_layer).boxed()),
+                None => Some(new_layer.boxed()),
+            };
+            self.layers = new;
+        }
+    }
+    let mut layers = LayerStack::default();
+
+    layers.add_layer({
         let log_layer = tracing_subscriber::fmt::layer()
             .with_target(false)
             .with_ansi(false)
@@ -131,17 +171,51 @@ pub fn init(
         };
         log_layer.with_filter(rust_log_env_filter())
     });
-    let r = r.with(
+
+    layers.add_layer(
         TracingEventCountLayer(&TRACING_EVENT_COUNT_METRIC).with_filter(rust_log_env_filter()),
     );
+
+    let tracing_chrome_layer_flush_guard =
+        if crate::env::var("NEON_UTILS_LOGGING_ENABLE_TRACING_CHROME").unwrap_or(false) {
+            let (layer, guard) = tracing_chrome::ChromeLayerBuilder::new()
+                .trace_style(tracing_chrome::TraceStyle::Async)
+                .build();
+            layers.add_layer(layer.with_filter(rust_log_env_filter()));
+            Some(guard)
+        } else {
+            None
+        };
+
+    let tracing_flame_flush_guard =
+        if crate::env::var("NEON_UTILS_LOGGING_ENABLE_TRACING_FLAME").unwrap_or(false) {
+            let (layer, guard) = tracing_flame::FlameLayer::with_file("./tracing.folded").unwrap();
+            let layer = layer
+                .with_empty_samples(false)
+                .with_module_path(false)
+                .with_file_and_line(false)
+                .with_threads_collapsed(true);
+            layers.add_layer(layer.with_filter(rust_log_env_filter()));
+            Some(guard)
+        } else {
+            None
+        };
+
     match tracing_error_layer_enablement {
-        TracingErrorLayerEnablement::EnableWithRustLogFilter => r
-            .with(tracing_error::ErrorLayer::default().with_filter(rust_log_env_filter()))
-            .init(),
-        TracingErrorLayerEnablement::Disabled => r.init(),
+        TracingErrorLayerEnablement::EnableWithRustLogFilter => layers
+            .add_layer(tracing_error::ErrorLayer::default().with_filter(rust_log_env_filter())),
+        TracingErrorLayerEnablement::Disabled => (),
     }
 
-    Ok(())
+    let r = tracing_subscriber::registry();
+    r.with(layers.layers.expect("we add at least one layer"))
+        .init();
+
+    Ok(FlushGuardInner {
+        _tracing_chrome_layer: tracing_chrome_layer_flush_guard,
+        _tracing_flame_layer: tracing_flame_flush_guard,
+    }
+    .into())
 }
 
 /// Disable the default rust panic hook by using `set_hook`.