Skip to content

Commit

Permalink
feat: add tracing-flame
Browse files Browse the repository at this point in the history
  • Loading branch information
ctron committed Apr 24, 2024
1 parent 8058ba8 commit 04d2684
Show file tree
Hide file tree
Showing 4 changed files with 76 additions and 15 deletions.
3 changes: 3 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand All @@ -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"}
Expand Down
28 changes: 20 additions & 8 deletions macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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]
Expand All @@ -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
}
Expand Down Expand Up @@ -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()
Expand All @@ -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! {}
}
57 changes: 53 additions & 4 deletions src/tracing.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down
3 changes: 0 additions & 3 deletions tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,12 @@ use tracing::error;
use tracing::info;
use tracing::instrument;


mod something {
pub type Error = String;
}

use something::Error;


#[test_log::test]
fn without_return_type() {
assert_eq!(2 + 2, 4);
Expand Down Expand Up @@ -131,7 +129,6 @@ impl<T> Foo for T {}
#[test_log::test]
fn unambiguous_map() {}


/// A module used for testing the `test` attribute after importing it
/// via `use` instead of using fuller qualified syntax.
mod local {
Expand Down

0 comments on commit 04d2684

Please sign in to comment.