Add support for tracing-flame

Cargo.toml

@@ -35,8 +35,9 @@ required-features = ["log", "unstable"]
 [features]
 default = ["log"]
 # TODO: use "dep:{tracing-subscriber,evn_logger}" once our MSRV is 1.60 or higher.
-trace = ["tracing-subscriber", "test-log-macros/trace"]
+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;
 
@@ -16,7 +16,6 @@ use syn::ItemFn;
 use syn::Lit;
 use syn::Meta;
 
-
 #[proc_macro_attribute]
 pub fn test(attr: TokenStream, item: TokenStream) -> TokenStream {
   let item = parse_macro_input!(item as ItemFn);
@@ -59,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]
@@ -76,21 +75,28 @@ 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
     }
   };
   Ok(result)
 }
 
-
 #[derive(Debug, Default)]
 struct AttributeArgs {
   default_log_filter: Option<String>,
@@ -99,7 +105,7 @@ struct AttributeArgs {
 impl AttributeArgs {
   fn try_parse_attr_single(&mut self, attr: &Attribute) -> syn::Result<bool> {
     if !attr.path().is_ident("test_log") {
-      return Ok(false)
+      return Ok(false);
     }
 
     let nested_meta = attr.parse_args_with(Meta::parse)?;
@@ -109,7 +115,7 @@ impl AttributeArgs {
       return Err(syn::Error::new_spanned(
         &nested_meta,
         "Expected NameValue syntax, e.g. 'default_log_filter = \"debug\"'.",
-      ))
+      ));
     };
 
     let ident = if let Some(ident) = name_value.path.get_ident() {
@@ -118,7 +124,7 @@ impl AttributeArgs {
       return Err(syn::Error::new_spanned(
         &name_value.path,
         "Expected NameValue syntax, e.g. 'default_log_filter = \"debug\"'.",
-      ))
+      ));
     };
 
     let arg_ref = if ident == "default_log_filter" {
@@ -127,7 +133,7 @@ impl AttributeArgs {
       return Err(syn::Error::new_spanned(
         &name_value.path,
         "Unrecognized attribute, see documentation for details.",
-      ))
+      ));
     };
 
     if let Expr::Lit(lit) = &name_value.value {
@@ -142,14 +148,13 @@ impl AttributeArgs {
       return Err(syn::Error::new_spanned(
         &name_value.value,
         "Failed to parse value, expected a string",
-      ))
+      ));
     }
 
     Ok(true)
   }
 }
 
-
 /// Expand the initialization code for the `log` crate.
 #[cfg(feature = "log")]
 fn expand_logging_init(attribute_args: &AttributeArgs) -> Tokens {
@@ -179,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()
@@ -194,45 +199,16 @@ fn expand_tracing_init(attribute_args: &AttributeArgs) -> Tokens {
     quote! { ::test_log::tracing_subscriber::EnvFilter::from_default_env() }
   };
 
-  quote! {
-    {
-      let __internal_event_filter = {
-        use ::test_log::tracing_subscriber::fmt::format::FmtSpan;
-
-        match ::std::env::var_os("RUST_LOG_SPAN_EVENTS") {
-          Some(mut value) => {
-            value.make_ascii_lowercase();
-            let value = value.to_str().expect("test-log: RUST_LOG_SPAN_EVENTS must be valid UTF-8");
-            value
-              .split(",")
-              .map(|filter| match filter.trim() {
-                "new" => FmtSpan::NEW,
-                "enter" => FmtSpan::ENTER,
-                "exit" => FmtSpan::EXIT,
-                "close" => FmtSpan::CLOSE,
-                "active" => FmtSpan::ACTIVE,
-                "full" => FmtSpan::FULL,
-                _ => panic!("test-log: RUST_LOG_SPAN_EVENTS must contain filters separated by `,`.\n\t\
-                  For example: `active` or `new,close`\n\t\
-                  Supported filters: new, enter, exit, close, active, full\n\t\
-                  Got: {}", value),
-              })
-              .fold(FmtSpan::NONE, |acc, filter| filter | acc)
-          },
-          None => FmtSpan::NONE,
-        }
-      };
+  let name = name.to_string();
 
-      let _ = ::test_log::tracing_subscriber::FmtSubscriber::builder()
-        .with_env_filter(#env_filter)
-        .with_span_events(__internal_event_filter)
-        .with_test_writer()
-        .try_init();
-    }
+  quote! {
+    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/lib.rs

@@ -68,3 +68,7 @@ pub use tracing_subscriber;
 #[cfg(feature = "log")]
 #[doc(hidden)]
 pub use env_logger;
+
+#[cfg(feature = "trace")]
+#[doc(hidden)]
+pub mod tracing;

src/tracing.rs

@@ -0,0 +1,94 @@
+//! 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(name: &str, env_filter: impl Into<tracing_subscriber::EnvFilter>) -> TracingGuard {
+  let env_filter = env_filter.into();
+  let event_filter = eval_event_filter();
+
+  let fmt = tracing_subscriber::fmt::layer()
+    .with_ansi(true)
+    .with_span_events(event_filter)
+    .with_level(true)
+    .with_test_writer()
+    .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 {
+  match var_os("RUST_LOG_SPAN_EVENTS") {
+    Some(mut value) => {
+      value.make_ascii_lowercase();
+      let value = value
+        .to_str()
+        .expect("test-log: RUST_LOG_SPAN_EVENTS must be valid UTF-8");
+      value
+        .split(",")
+        .map(|filter| match filter.trim() {
+          "new" => FmtSpan::NEW,
+          "enter" => FmtSpan::ENTER,
+          "exit" => FmtSpan::EXIT,
+          "close" => FmtSpan::CLOSE,
+          "active" => FmtSpan::ACTIVE,
+          "full" => FmtSpan::FULL,
+          _ => panic!(
+            "test-log: RUST_LOG_SPAN_EVENTS must contain filters separated by `,`.\n\t\
+                  For example: `active` or `new,close`\n\t\
+                  Supported filters: new, enter, exit, close, active, full\n\t\
+                  Got: {}",
+            value
+          ),
+        })
+        .fold(FmtSpan::NONE, |acc, filter| filter | acc)
+    },
+    None => FmtSpan::NONE,
+  }
+}