iai_callgrind_runner/runner/
common.rs

1//! This module contains elements which are common to library and binary benchmarks
2
3mod defaults {
4    pub const SANDBOX_ENABLED: bool = false;
5    pub const SANDBOX_FIXTURES_FOLLOW_SYMLINKS: bool = false;
6}
7
8use std::ffi::OsString;
9use std::fmt::Display;
10use std::path::PathBuf;
11use std::process::{Child, Command, Stdio as StdStdio};
12use std::time::{Duration, Instant};
13
14use anyhow::Result;
15use log::{debug, info, log_enabled, trace, Level};
16use tempfile::TempDir;
17
18use super::args::NoCapture;
19use super::format::{OutputFormatKind, SummaryFormatter};
20use super::meta::Metadata;
21use super::summary::BenchmarkSummary;
22use crate::api::{self, Pipe};
23use crate::error::Error;
24use crate::util::{copy_directory, make_absolute, write_all_to_stderr};
25
26/// The `Baselines` type
27pub type Baselines = (Option<String>, Option<String>);
28
29/// the [`Assistant`] kind
30#[derive(Debug, Clone)]
31pub enum AssistantKind {
32    /// The `setup` function
33    Setup,
34    /// The `teardown` function
35    Teardown,
36}
37
38/// An `Assistant` corresponds to the `setup` or `teardown` functions in the UI
39#[derive(Debug, Clone)]
40pub struct Assistant {
41    envs: Vec<(OsString, OsString)>,
42    group_name: Option<String>,
43    indices: Option<(usize, usize)>,
44    kind: AssistantKind,
45    pipe: Option<Pipe>,
46    run_parallel: bool,
47}
48/// Contains benchmark summaries of (binary, library) benchmark runs and their execution time
49///
50/// Used to print a final summary after all benchmarks.
51#[derive(Debug, Default)]
52pub struct BenchmarkSummaries {
53    /// The benchmark summaries
54    pub summaries: Vec<BenchmarkSummary>,
55    /// The execution time of all benchmarks.
56    pub total_time: Option<Duration>,
57}
58
59/// The `Config` contains all the information extracted from the UI invocation of the runner
60#[derive(Debug)]
61pub struct Config {
62    /// The path to the compiled binary with the benchmark harness
63    pub bench_bin: PathBuf,
64    /// The path to the benchmark file which contains the benchmark harness
65    pub bench_file: PathBuf,
66    /// The [`Metadata`]
67    pub meta: Metadata,
68    /// The module path of the benchmark file
69    pub module_path: ModulePath,
70    /// The package directory of the package in which `iai-callgrind` (not the runner) is used
71    pub package_dir: PathBuf,
72}
73
74/// A helper struct similar to a file path but for module paths with the `::` delimiter
75#[derive(Debug, PartialEq, Eq, Clone)]
76pub struct ModulePath(String);
77
78/// The `Sandbox` in which benchmarks should be runs
79///
80/// As soon as the `Sandbox` is dropped the temporary directory is deleted.
81#[derive(Debug)]
82pub struct Sandbox {
83    current_dir: PathBuf,
84    temp_dir: Option<TempDir>,
85}
86
87impl Assistant {
88    /// The setup or teardown of the `main` macro
89    pub fn new_main_assistant(
90        kind: AssistantKind,
91        envs: Vec<(OsString, OsString)>,
92        run_parallel: bool,
93    ) -> Self {
94        Self {
95            kind,
96            group_name: None,
97            indices: None,
98            pipe: None,
99            envs,
100            run_parallel,
101        }
102    }
103
104    /// The setup or teardown of a `binary_benchmark_group` or `library_benchmark_group`
105    pub fn new_group_assistant(
106        kind: AssistantKind,
107        group_name: &str,
108        envs: Vec<(OsString, OsString)>,
109        run_parallel: bool,
110    ) -> Self {
111        Self {
112            kind,
113            group_name: Some(group_name.to_owned()),
114            indices: None,
115            pipe: None,
116            envs,
117            run_parallel,
118        }
119    }
120
121    /// The setup or teardown function of a `Bench`
122    ///
123    /// This is currently only used by binary benchmarks. Library benchmarks use a completely
124    /// different logic for setup and teardown functions specified in a `#[bench]`, `#[benches]` and
125    /// `#[library_benchmark]` and don't need to be executed via the compiled benchmark.
126    pub fn new_bench_assistant(
127        kind: AssistantKind,
128        group_name: &str,
129        indices: (usize, usize),
130        pipe: Option<Pipe>,
131        envs: Vec<(OsString, OsString)>,
132        run_parallel: bool,
133    ) -> Self {
134        Self {
135            kind,
136            group_name: Some(group_name.to_owned()),
137            indices: Some(indices),
138            pipe,
139            envs,
140            run_parallel,
141        }
142    }
143
144    /// Run the `Assistant` by calling the benchmark binary with the needed arguments
145    ///
146    /// We don't run the assistant if `--load-baseline` was given on the command-line!
147    pub fn run(&self, config: &Config, module_path: &ModulePath) -> Result<Option<Child>> {
148        if config.meta.args.load_baseline.is_some() {
149            return Ok(None);
150        }
151
152        let id = self.kind.id();
153        let nocapture = config.meta.args.nocapture;
154
155        let mut command = Command::new(&config.bench_bin);
156        command.envs(self.envs.iter().cloned());
157        command.arg("--iai-run");
158
159        if let Some(group_name) = &self.group_name {
160            command.arg(group_name);
161        }
162
163        command.arg(&id);
164
165        if let Some((group_index, bench_index)) = &self.indices {
166            command.args([group_index.to_string(), bench_index.to_string()]);
167        }
168
169        nocapture.apply(&mut command);
170
171        match &self.pipe {
172            Some(Pipe::Stdout) => {
173                command.stdout(StdStdio::piped());
174            }
175            Some(Pipe::Stderr) => {
176                command.stderr(StdStdio::piped());
177            }
178            _ => {}
179        }
180
181        if self.pipe.is_some() || self.run_parallel {
182            let child = command
183                .spawn()
184                .map_err(|error| Error::LaunchError(config.bench_bin.clone(), error.to_string()))?;
185            return Ok(Some(child));
186        }
187
188        match nocapture {
189            NoCapture::False => {
190                let output = command
191                    .output()
192                    .map_err(|error| {
193                        Error::LaunchError(config.bench_bin.clone(), error.to_string())
194                    })
195                    .and_then(|output| {
196                        if output.status.success() {
197                            Ok(output)
198                        } else {
199                            let status = output.status;
200                            Err(Error::ProcessError(
201                                module_path.join(&id).to_string(),
202                                Some(output),
203                                status,
204                                None,
205                            ))
206                        }
207                    })?;
208
209                if log_enabled!(Level::Info) && !output.stdout.is_empty() {
210                    info!("{id} function in group '{module_path}': stdout:");
211                    write_all_to_stderr(&output.stdout);
212                }
213
214                if log_enabled!(Level::Info) && !output.stderr.is_empty() {
215                    info!("{id} function in group '{module_path}': stderr:");
216                    write_all_to_stderr(&output.stderr);
217                }
218            }
219            NoCapture::True | NoCapture::Stderr | NoCapture::Stdout => {
220                command
221                    .status()
222                    .map_err(|error| {
223                        Error::LaunchError(config.bench_bin.clone(), error.to_string())
224                    })
225                    .and_then(|status| {
226                        if status.success() {
227                            Ok(())
228                        } else {
229                            Err(Error::ProcessError(
230                                format!("{module_path}::{id}"),
231                                None,
232                                status,
233                                None,
234                            ))
235                        }
236                    })?;
237            }
238        }
239
240        Ok(None)
241    }
242}
243
244impl AssistantKind {
245    /// Return the assistant kind `id` as string
246    pub fn id(&self) -> String {
247        match self {
248            Self::Setup => "setup",
249            Self::Teardown => "teardown",
250        }
251        .to_owned()
252    }
253}
254
255impl BenchmarkSummaries {
256    /// Add a [`BenchmarkSummary`]
257    pub fn add_summary(&mut self, summary: BenchmarkSummary) {
258        self.summaries.push(summary);
259    }
260
261    /// Add another `BenchmarkSummary`
262    ///
263    /// Ignores the execution time.
264    pub fn add_other(&mut self, other: Self) {
265        other.summaries.into_iter().for_each(|s| {
266            self.add_summary(s);
267        });
268    }
269
270    /// Return true if any regressions were encountered
271    pub fn is_regressed(&self) -> bool {
272        self.summaries.iter().any(BenchmarkSummary::is_regressed)
273    }
274
275    /// Set the total execution from `start` to `now`
276    pub fn elapsed(&mut self, start: Instant) {
277        self.total_time = Some(start.elapsed());
278    }
279
280    /// Return the number of total benchmarks
281    pub fn num_benchmarks(&self) -> usize {
282        self.summaries.len()
283    }
284
285    /// Print the summary if not prevented by command-line arguments
286    ///
287    /// If `nosummary` is true or [`OutputFormatKind`] is any kind of `JSON` format the summary is
288    /// not printed.
289    pub fn print(&self, nosummary: bool, output_format_kind: OutputFormatKind) {
290        if !nosummary {
291            SummaryFormatter::new(output_format_kind).print(self);
292        }
293    }
294}
295
296impl ModulePath {
297    /// Create a new `ModulePath`
298    ///
299    /// There is no validity check if the path contains valid characters or not and the path is
300    /// created as is.
301    pub fn new(path: &str) -> Self {
302        Self(path.to_owned())
303    }
304
305    /// Join this module path with another string (unchecked)
306    #[must_use]
307    pub fn join(&self, path: &str) -> Self {
308        let new = format!("{}::{path}", self.0);
309        Self(new)
310    }
311
312    /// Return the module path as string
313    pub fn as_str(&self) -> &str {
314        &self.0
315    }
316
317    /// Return the first segment of the module path if any
318    pub fn first(&self) -> Option<Self> {
319        self.0
320            .split_once("::")
321            .map(|(first, _)| Self::new(first))
322            .or_else(|| (!self.0.is_empty()).then_some(self.clone()))
323    }
324
325    /// Return the last segment of the module path if any
326    pub fn last(&self) -> Option<Self> {
327        self.0.rsplit_once("::").map(|(_, last)| Self::new(last))
328    }
329
330    /// Return the parent module path if present
331    pub fn parent(&self) -> Option<Self> {
332        self.0
333            .rsplit_once("::")
334            .map(|(prefix, _)| Self::new(prefix))
335    }
336
337    /// Return a vector which contains all segments of the module path without the delimiter
338    pub fn components(&self) -> Vec<&str> {
339        self.0.split("::").collect()
340    }
341}
342
343impl Display for ModulePath {
344    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345        f.write_str(&self.0)
346    }
347}
348
349impl Sandbox {
350    /// Setup the `Sandbox` if enabled
351    ///
352    /// If enabled, create a temporary directory which has a standardized length. Then copy fixtures
353    /// into the temporary directory. Finally, set the current directory to this temporary
354    /// directory.
355    pub fn setup(inner: &api::Sandbox, meta: &Metadata) -> Result<Self> {
356        let enabled = inner.enabled.unwrap_or(defaults::SANDBOX_ENABLED);
357        let follow_symlinks = inner
358            .follow_symlinks
359            .unwrap_or(defaults::SANDBOX_FIXTURES_FOLLOW_SYMLINKS);
360        let current_dir = std::env::current_dir().map_err(|error| {
361            Error::SandboxError(format!("Failed to detect current directory: {error}"))
362        })?;
363
364        let temp_dir = if enabled {
365            debug!("Creating sandbox");
366
367            let temp_dir = tempfile::tempdir().map_err(|error| {
368                Error::SandboxError(format!("Failed creating temporary directory: {error}"))
369            })?;
370
371            for fixture in &inner.fixtures {
372                if fixture.is_relative() {
373                    let absolute_path = make_absolute(&meta.project_root, fixture);
374                    copy_directory(&absolute_path, temp_dir.path(), follow_symlinks)?;
375                } else {
376                    copy_directory(fixture, temp_dir.path(), follow_symlinks)?;
377                }
378            }
379
380            trace!(
381                "Changing current directory to sandbox directory: '{}'",
382                temp_dir.path().display()
383            );
384
385            let path = temp_dir.path();
386            std::env::set_current_dir(path).map_err(|error| {
387                Error::SandboxError(format!(
388                    "Failed setting current directory to sandbox directory: '{error}'"
389                ))
390            })?;
391            Some(temp_dir)
392        } else {
393            debug!(
394                "Sandbox disabled: Running benchmarks in current directory '{}'",
395                current_dir.display()
396            );
397            None
398        };
399
400        Ok(Self {
401            current_dir,
402            temp_dir,
403        })
404    }
405
406    /// Reset the current directory and delete the temporary directory if present
407    pub fn reset(self) -> Result<()> {
408        if let Some(temp_dir) = self.temp_dir {
409            std::env::set_current_dir(&self.current_dir).map_err(|error| {
410                Error::SandboxError(format!("Failed to reset current directory: {error}"))
411            })?;
412
413            if log_enabled!(Level::Debug) {
414                debug!("Removing temporary workspace");
415                if let Err(error) = temp_dir.close() {
416                    debug!("Error trying to delete temporary workspace: {error}");
417                }
418            } else {
419                _ = temp_dir.close();
420            }
421        }
422
423        Ok(())
424    }
425}
426
427impl From<ModulePath> for String {
428    fn from(value: ModulePath) -> Self {
429        value.to_string()
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use rstest::rstest;
436
437    use super::*;
438
439    #[rstest]
440    #[case::empty("", None)]
441    #[case::single("first", Some("first"))]
442    #[case::two("first::second", Some("first"))]
443    #[case::three("first::second::third", Some("first"))]
444    fn test_module_path_first(#[case] module_path: &str, #[case] expected: Option<&str>) {
445        let expected = expected.map(ModulePath::new);
446        let actual = ModulePath::new(module_path).first();
447
448        assert_eq!(actual, expected);
449    }
450}