tracing_profile/layers/
init_tracing.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use cfg_if::cfg_if;
4use thiserror::Error;
5use tracing::{level_filters::LevelFilter, Subscriber};
6use tracing_subscriber::filter::EnvFilter;
7use tracing_subscriber::{
8    filter::Filtered,
9    layer::SubscriberExt,
10    util::{SubscriberInitExt, TryInitError},
11    Layer,
12};
13
14use crate::{PrintTreeConfig, PrintTreeLayer};
15
16trait WithEnvFilter<S: Subscriber>: Layer<S> + Sized {
17    fn with_env_filter(self) -> Filtered<Self, EnvFilter, S> {
18        let env_level_filter = EnvFilter::builder()
19            .with_default_directive(LevelFilter::DEBUG.into())
20            .from_env_lossy();
21
22        self.with_filter(env_level_filter)
23    }
24}
25
26impl<S: Subscriber, T: Layer<S>> WithEnvFilter<S> for T {}
27
28#[derive(Debug, Error)]
29pub enum Error {
30    #[error("failed to initialize tracing: {0}")]
31    TryInit(#[from] TryInitError),
32    #[cfg(feature = "perfetto")]
33    #[error("failed to initialize Perfetto: {0}")]
34    Perfetto(#[from] perfetto_sys::Error),
35    #[cfg(feature = "perf_counters")]
36    #[error("failed to initialize PerfCounters: {0}")]
37    PerfCounters(#[from] std::io::Error),
38    #[cfg(feature = "gen_filename")]
39    #[error("failed to initialize filename builder: {0}")]
40    FilenameBuilder(#[from] crate::filename_builder::FilenameBuilderError),
41}
42
43// Type aliases to handle conditional compilation cleanly
44#[cfg(feature = "gen_filename")]
45mod filename_support {
46    pub use crate::filename_builder::TraceFilenameBuilder;
47    pub type BuilderOption = Option<TraceFilenameBuilder>;
48}
49
50#[cfg(not(feature = "gen_filename"))]
51mod filename_support {
52    pub type BuilderOption = Option<()>;
53}
54
55use filename_support::BuilderOption;
56
57/// Internal function that handles common layer setup logic.
58///
59/// Sets up all available tracing layers based on enabled features:
60/// - PrintTreeLayer (always enabled)
61/// - PerfettoLayer (if perfetto feature enabled)
62/// - IttApiLayer (if ittapi feature enabled)
63/// - TracyLayer (if tracy feature enabled)
64/// - PrintPerfCountersLayer (if perf_counters feature enabled)
65///
66/// The builder parameter allows customization of perfetto trace filenames
67/// when the gen_filename feature is enabled.
68fn init_tracing_internal(_builder: BuilderOption) -> Result<impl Drop, Error> {
69    // Create print tree layer
70    let (layer, guard) = PrintTreeLayer::new(PrintTreeConfig::default());
71    let layer = tracing_subscriber::registry().with(layer.with_env_filter());
72
73    // Add perfetto layer if feature is enabled
74    let (layer, guard) = {
75        cfg_if! {
76            if #[cfg(feature = "perfetto")] {
77                let (new_layer, new_guard) = match _builder {
78                    None => {
79                        crate::PerfettoLayer::new_from_env()?
80                    }
81                    Some(builder) => {
82                        crate::PerfettoLayer::new_from_env_with_builder(builder)
83                            .map_err(Error::Perfetto)?
84                    }
85                };
86                (layer.with(new_layer.with_env_filter()), crate::data::GuardWrapper::wrap(guard, new_guard))
87            } else {
88                (layer, guard)
89            }
90        }
91    };
92
93    // Add ITT API layer if feature is enabled
94    let (layer, guard) = {
95        cfg_if! {
96            if #[cfg(feature = "ittapi")] {
97                (layer.with(crate::IttApiLayer::new().with_env_filter()), guard)
98            } else {
99                (layer, guard)
100            }
101        }
102    };
103
104    // Add tracy layer if feature is enabled
105    let (layer, guard) = {
106        cfg_if! {
107            if #[cfg(feature = "tracy")] {
108                (layer.with(crate::TracyLayer::default().with_env_filter()), guard)
109            } else {
110                (layer, guard)
111            }
112        }
113    };
114
115    // Add perf counters layer if feature is enabled
116    let (layer, guard) = {
117        cfg_if! {
118            if #[cfg(feature = "perf_counters")] {
119                (layer.with(
120                    crate::PrintPerfCountersLayer::new(vec![
121                        ("instructions".to_string(), crate::PerfHardwareEvent::INSTRUCTIONS.into()),
122                        ("cycles".to_string(), crate::PerfHardwareEvent::CPU_CYCLES.into()),
123                    ])?
124                    .with_env_filter(),
125                ), guard)
126            } else {
127                (layer, guard)
128            }
129        }
130    };
131
132    // Try to initialize subscriber - OK if already set
133    match layer.try_init() {
134        Ok(()) => {
135            // First initialization succeeded
136        }
137        Err(_) => {
138            // Subscriber already initialized - this is fine
139            // The perfetto guard will still be created and work correctly
140            // Tracing events will go to the existing subscriber
141        }
142    }
143
144    Ok(guard)
145}
146
147/// Initialize the tracing with the default values depending on the features enabled and environment variables set.
148///
149/// The following layers are added:
150/// - `PrintTreeLayer` (added always)
151/// - `IttApiLayer` (added if feature `ittapi` is enabled)
152/// - `TracyLayer` (added if feature `tracy` is enabled)
153/// - `PrintPerfCountersLayer` (added if feature `perf_counters` is enabled)
154///
155/// Returns the guard that should be kept alive for the duration of the program.
156pub fn init_tracing() -> Result<impl Drop, Error> {
157    init_tracing_internal(None)
158}
159
160/// Initialize tracing with a custom TraceFilenameBuilder for perfetto traces.
161///
162/// This function provides the same functionality as `init_tracing()` but allows
163/// you to customize the perfetto trace filename using TraceFilenameBuilder.
164///
165/// # Example
166/// ```rust,no_run
167/// use tracing_profile::{init_tracing_with_builder, TraceFilenameBuilder};
168///
169/// let builder = TraceFilenameBuilder::new()
170///     .name("my_app")
171///     .iteration(1)
172///     .timestamp()
173///     .git_info();
174///
175/// let _guard = init_tracing_with_builder(builder).unwrap();
176/// ```
177#[cfg(feature = "gen_filename")]
178pub fn init_tracing_with_builder(
179    builder: crate::filename_builder::TraceFilenameBuilder,
180) -> Result<impl Drop, Error> {
181    init_tracing_internal(Some(builder))
182}