feat: add tracing-flame

Author: ctron

Cargo.toml

@@ -37,6 +37,7 @@ default = ["log"]
 # TODO: use "dep:{tracing-subscriber,evn_logger}" once our MSRV is 1.60 or higher.
 trace = ["tracing-subscriber", "tracing-subscriber/ansi", "tracing-subscriber/tracing-log", "test-log-macros/trace"]
 log = ["env_logger", "test-log-macros/log"]
+tracing-flame = ["trace", "dep:tracing-flame"]
 # Enable unstable features. These are generally exempt from any semantic
 # versioning guarantees.
 unstable = ["test-log-macros/unstable"]
@@ -49,6 +50,8 @@ test-log-macros = {version = "0.2.15", path = "macros"}
 tracing-subscriber = {version = "0.3.17", default-features = false, optional = true, features = ["env-filter", "fmt"]}
 env_logger = {version = "0.11", default-features = false, optional = true}
 
+tracing-flame = { version = "0.2.0", optional = true }
+
 [dev-dependencies]
 logging = {version = "0.4.8", package = "log"}
 test-case = {version = "3.1"}

macros/src/lib.rs

@@ -4,7 +4,7 @@
 extern crate proc_macro;
 
 use proc_macro::TokenStream;
-use proc_macro2::TokenStream as Tokens;
+use proc_macro2::{Ident, TokenStream as Tokens};
 
 use quote::quote;
 
@@ -58,7 +58,7 @@ fn try_test(attr: TokenStream, input: ItemFn) -> syn::Result<Tokens> {
 
   let (attribute_args, ignored_attrs) = parse_attrs(attrs)?;
   let logging_init = expand_logging_init(&attribute_args);
-  let tracing_init = expand_tracing_init(&attribute_args);
+  let tracing_init = expand_tracing_init(&sig.ident, &attribute_args);
 
   let result = quote! {
     #[#inner_test]
@@ -75,13 +75,21 @@ fn try_test(attr: TokenStream, input: ItemFn) -> syn::Result<Tokens> {
       // The alternative would be to use fully qualified call syntax in
       // all initialization code, but that's much harder to control.
       mod init {
-        pub fn init() {
+        pub struct Guard<T> {
+            tracing: T,
+        }
+
+        pub fn init() -> Guard<impl std::any::Any> {
           #logging_init
-          #tracing_init
+          Guard {
+            tracing: {
+              #tracing_init
+            },
+          }
         }
       }
 
-      init::init();
+      let _ = init::init();
 
       #block
     }
@@ -176,7 +184,7 @@ fn expand_logging_init(_attribute_args: &AttributeArgs) -> Tokens {
 
 /// Expand the initialization code for the `tracing` crate.
 #[cfg(feature = "trace")]
-fn expand_tracing_init(attribute_args: &AttributeArgs) -> Tokens {
+fn expand_tracing_init(name: &Ident, attribute_args: &AttributeArgs) -> Tokens {
   let env_filter = if let Some(default_log_filter) = &attribute_args.default_log_filter {
     quote! {
       ::test_log::tracing_subscriber::EnvFilter::builder()
@@ -191,12 +199,16 @@ fn expand_tracing_init(attribute_args: &AttributeArgs) -> Tokens {
     quote! { ::test_log::tracing_subscriber::EnvFilter::from_default_env() }
   };
 
+  let name = name.to_string();
+
   quote! {
-      ::test_log::tracing::init(#env_filter);
+    let base = module_path!().split("::").map(std::path::Path::new).collect::<std::path::PathBuf>();
+    let name = format!("{}/{}", base.display(), #name);
+    Some(::test_log::tracing::init(&name, #env_filter))
   }
 }
 
 #[cfg(not(feature = "trace"))]
-fn expand_tracing_init(_attribute_args: &AttributeArgs) -> Tokens {
+fn expand_tracing_init(_name: &Ident, _attribute_args: &AttributeArgs) -> Tokens {
   quote! {}
 }

src/tracing.rs

@@ -1,17 +1,66 @@
 //! Support for tracing
 
 use std::env::var_os;
+use std::fs::File;
+use std::io::BufWriter;
+use std::path::Path;
 use tracing_subscriber::fmt::format::FmtSpan;
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::util::SubscriberInitExt;
+
+#[derive(Default)]
+pub struct TracingGuard {
+  #[cfg(feature = "tracing-flame")]
+  _flame: Option<tracing_flame::FlushGuard<BufWriter<File>>>,
+}
 
 /// Initialize the tracing
-pub fn init(env_filter: impl Into<tracing_subscriber::EnvFilter>) {
+pub fn init(name: &str, env_filter: impl Into<tracing_subscriber::EnvFilter>) -> TracingGuard {
+  let env_filter = env_filter.into();
   let event_filter = eval_event_filter();
 
-  let _ = tracing_subscriber::FmtSubscriber::builder()
-    .with_env_filter(env_filter)
+  let fmt = tracing_subscriber::fmt::layer()
+    .with_ansi(true)
     .with_span_events(event_filter)
+    .with_level(true)
     .with_test_writer()
-    .try_init();
+    .compact();
+
+  let layered = tracing_subscriber::registry().with(env_filter).with(fmt);
+
+  #[cfg(feature = "tracing-flame")]
+  {
+    return match std::env::var("TEST_LOG_FLAMES").ok() {
+      Some(base) => {
+        let path = format!("{base}/{name}.folded");
+        let path = Path::new(&path);
+
+        // ensure we have the parent dir
+        if let Some(parent) = path.parent() {
+          let _ = std::fs::create_dir_all(parent);
+        }
+
+        let (flame, guard) =
+          tracing_flame::FlameLayer::with_file(path).expect("Unable to initialize tracing-flame");
+
+        let _ = layered.with(flame).try_init();
+
+        TracingGuard {
+          _flame: Some(guard),
+        }
+      },
+      None => {
+        let _ = layered.try_init();
+        TracingGuard::default()
+      },
+    };
+  }
+
+  #[cfg(not(feature = "tracing-flame"))]
+  {
+    let layered = layered.with(env_filter).with(fmt).try_init();
+    TracingGuard::default()
+  }
 }
 
 fn eval_event_filter() -> FmtSpan {