Skip to main content

gwr_track/
builder.rs

1// Copyright (c) 2025 Graphcore Ltd. All rights reserved.
2
3//! Library functions to build trackers as defined by the user.
4
5use std::io::BufWriter;
6use std::rc::Rc;
7use std::{fs, io};
8
9use clap::Args;
10
11use crate::tracker::multi_tracker::MultiTracker;
12#[cfg(feature = "perfetto")]
13use crate::tracker::perfetto::PerfettoTracker;
14use crate::tracker::{CapnProtoTracker, EntityManager, TextTracker, TrackConfigError};
15use crate::{Tracker, Writer};
16
17/// Standard command-line arguments for tracker configuration.
18#[derive(Clone, Debug, Args)]
19pub struct TrackerArgs {
20    /// Enable logging to the console.
21    #[arg(long, default_value = "false")]
22    pub stdout: bool,
23
24    /// Level of log message to display.
25    #[arg(long, default_value = "Info")]
26    pub stdout_level: log::Level,
27
28    /// Set a regular expression for which entities should have logging level
29    /// set to `--stdout-level`. Others will have level set to `Error`.
30    #[arg(long, default_value = "")]
31    pub stdout_filter_regex: String,
32
33    /// Enable logging to binary file used by `gwr-spotter`.
34    #[arg(long, default_value = "false")]
35    pub binary: bool,
36
37    /// Level of binary trace events to record.
38    #[arg(long, default_value = "Trace")]
39    pub binary_level: log::Level,
40
41    /// Set a regular expression for which entities should have binary output
42    /// level set to `--binary-level`. Others will have level set to `Error`.
43    #[arg(long, default_value = "")]
44    pub binary_filter_regex: String,
45
46    /// The filename binary trace output is written to.
47    #[arg(long, default_value = "trace.bin")]
48    pub binary_file: String,
49
50    /// Enable logging to Perfetto file used by `gwr-spotter`.
51    #[cfg(feature = "perfetto")]
52    #[arg(long, default_value = "false")]
53    pub perfetto: bool,
54
55    /// Level of Perfetto trace events to record.
56    #[cfg(feature = "perfetto")]
57    #[arg(long, default_value = "Trace")]
58    pub perfetto_level: log::Level,
59
60    /// Set a regular expression for which entities should have Perfetto output
61    /// level set to `--perfetto-level`. Others will have level set to `Error`.
62    #[cfg(feature = "perfetto")]
63    #[arg(long, default_value = "")]
64    pub perfetto_filter_regex: String,
65
66    /// The filename Perfetto trace output is written to.
67    #[cfg(feature = "perfetto")]
68    #[arg(long, default_value = "trace.pftrace")]
69    pub perfetto_file: String,
70
71    /// Enable monitoring at the specified number of clock ticks.
72    #[arg(long)]
73    pub monitor_window_ticks: Option<u64>,
74
75    /// Set a regular expression for which ports should have monitors enabled.
76    #[arg(long, default_value = "")]
77    pub monitor_filter_regex: String,
78}
79
80impl TrackerArgs {
81    /// Return whether any tracker output has been explicitly requested.
82    #[must_use]
83    pub fn tracking_requested(&self) -> bool {
84        let requested = self.stdout || self.binary;
85        #[cfg(feature = "perfetto")]
86        let requested = requested || self.perfetto;
87        requested
88    }
89
90    /// Return whether any configured tracker will emit messages at `level`.
91    #[must_use]
92    pub fn level_enabled(&self, level: log::Level) -> bool {
93        let shown = (self.stdout && self.stdout_level >= level)
94            || (self.binary && self.binary_level >= level);
95        #[cfg(feature = "perfetto")]
96        let shown = shown || (self.perfetto && self.perfetto_level >= level);
97        shown
98    }
99
100    /// Ensure that if the specified feature is enabled then a tracker will be
101    /// showing messages of that level
102    pub fn ensure_visiblity(&mut self, feature: bool, feature_name: &str, level: log::Level) {
103        if feature && !self.level_enabled(level) {
104            self.stdout = true;
105            if self.stdout_level < log::Level::Info {
106                self.stdout_level = log::Level::Info;
107            }
108            eprintln!(
109                "WARNING: `{feature_name}` emits {level} messages, but no tracker is configured to show that level. Enabling stdout at {}.",
110                self.stdout_level
111            );
112        }
113    }
114
115    /// Convert these command-line arguments into a [`TrackersConfig`].
116    #[must_use]
117    pub fn trackers_config(&self) -> TrackersConfig<'_> {
118        TrackersConfig {
119            stdout: TrackerConfig {
120                enable: self.stdout,
121                level: self.stdout_level,
122                filter_regex: &self.stdout_filter_regex,
123                file: None,
124            },
125            binary: TrackerConfig {
126                enable: self.binary,
127                level: self.binary_level,
128                filter_regex: &self.binary_filter_regex,
129                file: Some(&self.binary_file),
130            },
131            #[cfg(feature = "perfetto")]
132            perfetto: TrackerConfig {
133                enable: self.perfetto,
134                level: self.perfetto_level,
135                filter_regex: &self.perfetto_filter_regex,
136                file: Some(&self.perfetto_file),
137            },
138            monitors: MonitorsConfig {
139                enable: self.monitor_window_ticks.is_some(),
140                window_size_ticks: self.monitor_window_ticks.unwrap_or(0),
141                filter_regex: &self.monitor_filter_regex,
142            },
143        }
144    }
145}
146
147/// Configuration options for an individual tracker.
148pub struct TrackerConfig<'a> {
149    /// Enable this tracker.
150    pub enable: bool,
151
152    /// Set the level at which this tracker should be enabled.
153    pub level: log::Level,
154
155    /// A regular expression to match which entities should have this level
156    /// applied.
157    pub filter_regex: &'a str,
158
159    /// If required, the name of the file to which the tracker will write.
160    pub file: Option<&'a str>,
161}
162
163impl Default for TrackerConfig<'_> {
164    fn default() -> Self {
165        Self {
166            enable: true,
167            level: log::Level::Warn,
168            filter_regex: "",
169            file: None,
170        }
171    }
172}
173
174/// Configuration options for monitoring.
175#[derive(Default)]
176pub struct MonitorsConfig<'a> {
177    /// Enable monitoring.
178    pub enable: bool,
179
180    /// Window size in clock ticks to process monitoring.
181    pub window_size_ticks: u64,
182
183    /// Regular expression for which entities should have monitoring enabled.
184    pub filter_regex: &'a str,
185}
186
187/// Configuration options for all tracking/monitoring.
188pub struct TrackersConfig<'a> {
189    /// Configuration for stdout.
190    pub stdout: TrackerConfig<'a>,
191
192    /// Configuration for binary trace file.
193    pub binary: TrackerConfig<'a>,
194
195    #[cfg(feature = "perfetto")]
196    /// Configuration for perfetto trace file.
197    pub perfetto: TrackerConfig<'a>,
198
199    /// Configuration for monitoring.
200    pub monitors: MonitorsConfig<'a>,
201}
202
203/// Create a tracker that prints to stdout
204///
205/// The user can pass a filter regular expression which will set the level only
206/// for matching Entities and set all other Entities to only emit errors.
207fn build_stdout_tracker(
208    config: &TrackerConfig,
209    monitors: &MonitorsConfig,
210) -> Result<Tracker, TrackConfigError> {
211    let default_level = if config.filter_regex.is_empty() {
212        config.level
213    } else {
214        log::Level::Error
215    };
216
217    let mut entity_manager = EntityManager::new(default_level);
218    if !config.filter_regex.is_empty() {
219        entity_manager.add_entity_level_filter(config.filter_regex, config.level)?;
220    }
221
222    if monitors.enable {
223        entity_manager
224            .set_monitor_window_size_for(monitors.filter_regex, monitors.window_size_ticks)?;
225    }
226
227    let stdout_writer = Box::new(std::io::BufWriter::new(io::stdout()));
228    Ok(Rc::new(TextTracker::new(entity_manager, stdout_writer)))
229}
230
231/// Same as the text tracker (see build_stdout_tracker) except will generate a
232/// binary file.
233fn build_binary_tracker(
234    config: &TrackerConfig,
235    monitors: &MonitorsConfig,
236) -> Result<Tracker, TrackConfigError> {
237    let default_level = if config.filter_regex.is_empty() {
238        config.level
239    } else {
240        log::Level::Error
241    };
242    let mut entity_manager = EntityManager::new(default_level);
243    if !config.filter_regex.is_empty() {
244        entity_manager.add_entity_level_filter(config.filter_regex, config.level)?;
245    }
246
247    if monitors.enable {
248        entity_manager
249            .set_monitor_window_size_for(monitors.filter_regex, monitors.window_size_ticks)?;
250    }
251
252    let bin_writer: Writer = Box::new(BufWriter::new(
253        fs::File::create(config.file.unwrap()).unwrap(),
254    ));
255    Ok(Rc::new(CapnProtoTracker::new(entity_manager, bin_writer)))
256}
257
258/// This tracker will produce a Perfetto trace file, which unlike the other
259/// tracker options can be viewed using the Perfetto UI, rather than
260/// gwr-spotter.
261#[cfg(feature = "perfetto")]
262fn build_perfetto_tracker(
263    config: &TrackerConfig,
264    monitors: &MonitorsConfig,
265) -> Result<Tracker, TrackConfigError> {
266    let default_level = if config.filter_regex.is_empty() {
267        config.level
268    } else {
269        log::Level::Error
270    };
271    let mut entity_manager = EntityManager::new(default_level);
272    if !config.filter_regex.is_empty() {
273        entity_manager.add_entity_level_filter(config.filter_regex, config.level)?;
274    }
275
276    if monitors.enable {
277        entity_manager
278            .set_monitor_window_size_for(monitors.filter_regex, monitors.window_size_ticks)?;
279    }
280
281    let bin_writer: Writer = Box::new(BufWriter::new(
282        fs::File::create(config.file.unwrap()).unwrap(),
283    ));
284    Ok(Rc::new(PerfettoTracker::new(entity_manager, bin_writer)))
285}
286
287/// Set up stdout/binary/Perfetto trackers according the the command-line
288/// arguments
289#[cfg(not(feature = "perfetto"))]
290pub fn setup_trackers(config: &TrackersConfig) -> Result<Tracker, TrackConfigError> {
291    let multi_tracker_required = config.stdout.enable && config.binary.enable;
292
293    if multi_tracker_required {
294        let mut tracker = MultiTracker::default();
295
296        if config.stdout.enable {
297            let log_tracker: Tracker = build_stdout_tracker(&config.stdout, &config.monitors)?;
298            tracker.add_tracker(log_tracker);
299        }
300        if config.binary.enable {
301            let trace_tracker: Tracker = build_binary_tracker(&config.binary, &config.monitors)?;
302            tracker.add_tracker(trace_tracker);
303        }
304
305        Ok(Rc::new(tracker))
306    } else if config.stdout.enable {
307        build_stdout_tracker(&config.stdout, &config.monitors)
308    } else if config.binary.enable {
309        build_binary_tracker(&config.binary, &config.monitors)
310    } else {
311        build_stdout_tracker(&TrackerConfig::default(), &MonitorsConfig::default())
312    }
313}
314
315/// Set up stdout/binary/Perfetto trackers according the the command-line
316/// arguments
317#[cfg(feature = "perfetto")]
318pub fn setup_trackers(config: &TrackersConfig) -> Result<Tracker, TrackConfigError> {
319    let multi_tracker_required = [
320        config.stdout.enable,
321        config.binary.enable,
322        config.perfetto.enable,
323    ]
324    .into_iter()
325    .filter(|x| *x)
326    .count()
327        > 1;
328
329    if multi_tracker_required {
330        let mut tracker = MultiTracker::default();
331
332        if config.stdout.enable {
333            let log_tracker: Tracker = build_stdout_tracker(&config.stdout, &config.monitors)?;
334            tracker.add_tracker(log_tracker);
335        }
336        if config.binary.enable {
337            let trace_tracker: Tracker = build_binary_tracker(&config.binary, &config.monitors)?;
338            tracker.add_tracker(trace_tracker);
339        }
340        if config.perfetto.enable {
341            let perfetto_tracker: Tracker =
342                build_perfetto_tracker(&config.perfetto, &config.monitors)?;
343            tracker.add_tracker(perfetto_tracker);
344        }
345
346        Ok(Rc::new(tracker))
347    } else if config.stdout.enable {
348        build_stdout_tracker(&config.stdout, &config.monitors)
349    } else if config.binary.enable {
350        build_binary_tracker(&config.binary, &config.monitors)
351    } else if config.perfetto.enable {
352        build_perfetto_tracker(&config.perfetto, &config.monitors)
353    } else {
354        build_stdout_tracker(&TrackerConfig::default(), &MonitorsConfig::default())
355    }
356}