iai_callgrind_runner/runner/
common.rs

1use std::ffi::OsString;
2use std::fmt::Display;
3use std::path::PathBuf;
4use std::process::{Child, Command, Stdio as StdStdio};
5
6use anyhow::Result;
7use log::{debug, info, log_enabled, trace, Level};
8use tempfile::TempDir;
9
10use super::args::NoCapture;
11use super::meta::Metadata;
12use crate::api::{self, Pipe};
13use crate::error::Error;
14use crate::util::{copy_directory, make_absolute, write_all_to_stderr};
15
16mod defaults {
17    pub const SANDBOX_FIXTURES_FOLLOW_SYMLINKS: bool = false;
18    pub const SANDBOX_ENABLED: bool = false;
19}
20
21#[derive(Debug, Clone)]
22pub struct Assistant {
23    kind: AssistantKind,
24    group_name: Option<String>,
25    indices: Option<(usize, usize)>,
26    pipe: Option<Pipe>,
27    envs: Vec<(OsString, OsString)>,
28    run_parallel: bool,
29}
30
31#[derive(Debug, Clone)]
32pub enum AssistantKind {
33    Setup,
34    Teardown,
35}
36
37#[derive(Debug)]
38pub struct Config {
39    pub package_dir: PathBuf,
40    pub bench_file: PathBuf,
41    pub module_path: ModulePath,
42    pub bench_bin: PathBuf,
43    pub meta: Metadata,
44}
45
46#[derive(Debug, PartialEq, Eq, Clone)]
47pub struct ModulePath(String);
48
49#[derive(Debug)]
50pub struct Sandbox {
51    current_dir: PathBuf,
52    temp_dir: Option<TempDir>,
53}
54
55impl Assistant {
56    /// The setup or teardown of the `main` macro
57    pub fn new_main_assistant(
58        kind: AssistantKind,
59        envs: Vec<(OsString, OsString)>,
60        run_parallel: bool,
61    ) -> Self {
62        Self {
63            kind,
64            group_name: None,
65            indices: None,
66            pipe: None,
67            envs,
68            run_parallel,
69        }
70    }
71
72    /// The setup or teardown of a `binary_benchmark_group` or `library_benchmark_group`
73    pub fn new_group_assistant(
74        kind: AssistantKind,
75        group_name: &str,
76        envs: Vec<(OsString, OsString)>,
77        run_parallel: bool,
78    ) -> Self {
79        Self {
80            kind,
81            group_name: Some(group_name.to_owned()),
82            indices: None,
83            pipe: None,
84            envs,
85            run_parallel,
86        }
87    }
88
89    /// The setup or teardown function of a `Bench`
90    ///
91    /// This is currently only used by binary benchmarks. Library benchmarks use a completely
92    /// different logic for setup and teardown functions specified in a `#[bench]`, `#[benches]` and
93    /// `#[library_benchmark]` and don't need to be executed via the compiled benchmark.
94    pub fn new_bench_assistant(
95        kind: AssistantKind,
96        group_name: &str,
97        indices: (usize, usize),
98        pipe: Option<Pipe>,
99        envs: Vec<(OsString, OsString)>,
100        run_parallel: bool,
101    ) -> Self {
102        Self {
103            kind,
104            group_name: Some(group_name.to_owned()),
105            indices: Some(indices),
106            pipe,
107            envs,
108            run_parallel,
109        }
110    }
111
112    /// Run the `Assistant` by calling the benchmark binary with the needed arguments
113    ///
114    /// We don't run the assistant if `--load-baseline` was given on the command-line!
115    pub fn run(&self, config: &Config, module_path: &ModulePath) -> Result<Option<Child>> {
116        if config.meta.args.load_baseline.is_some() {
117            return Ok(None);
118        }
119
120        let id = self.kind.id();
121        let nocapture = config.meta.args.nocapture;
122
123        let mut command = Command::new(&config.bench_bin);
124        command.envs(self.envs.iter().cloned());
125        command.arg("--iai-run");
126
127        if let Some(group_name) = &self.group_name {
128            command.arg(group_name);
129        }
130
131        command.arg(&id);
132
133        if let Some((group_index, bench_index)) = &self.indices {
134            command.args([group_index.to_string(), bench_index.to_string()]);
135        }
136
137        nocapture.apply(&mut command);
138
139        match &self.pipe {
140            Some(Pipe::Stdout) => {
141                command.stdout(StdStdio::piped());
142            }
143            Some(Pipe::Stderr) => {
144                command.stderr(StdStdio::piped());
145            }
146            _ => {}
147        }
148
149        if self.pipe.is_some() || self.run_parallel {
150            let child = command
151                .spawn()
152                .map_err(|error| Error::LaunchError(config.bench_bin.clone(), error.to_string()))?;
153            return Ok(Some(child));
154        }
155
156        match nocapture {
157            NoCapture::False => {
158                let output = command
159                    .output()
160                    .map_err(|error| {
161                        Error::LaunchError(config.bench_bin.clone(), error.to_string())
162                    })
163                    .and_then(|output| {
164                        if output.status.success() {
165                            Ok(output)
166                        } else {
167                            let status = output.status;
168                            Err(Error::ProcessError((
169                                module_path.join(&id).to_string(),
170                                Some(output),
171                                status,
172                                None,
173                            )))
174                        }
175                    })?;
176
177                if log_enabled!(Level::Info) && !output.stdout.is_empty() {
178                    info!("{id} function in group '{module_path}': stdout:");
179                    write_all_to_stderr(&output.stdout);
180                }
181
182                if log_enabled!(Level::Info) && !output.stderr.is_empty() {
183                    info!("{id} function in group '{module_path}': stderr:");
184                    write_all_to_stderr(&output.stderr);
185                }
186            }
187            NoCapture::True | NoCapture::Stderr | NoCapture::Stdout => {
188                command
189                    .status()
190                    .map_err(|error| {
191                        Error::LaunchError(config.bench_bin.clone(), error.to_string())
192                    })
193                    .and_then(|status| {
194                        if status.success() {
195                            Ok(())
196                        } else {
197                            Err(Error::ProcessError((
198                                format!("{module_path}::{id}"),
199                                None,
200                                status,
201                                None,
202                            )))
203                        }
204                    })?;
205            }
206        };
207
208        Ok(None)
209    }
210}
211
212impl AssistantKind {
213    pub fn id(&self) -> String {
214        match self {
215            AssistantKind::Setup => "setup",
216            AssistantKind::Teardown => "teardown",
217        }
218        .to_owned()
219    }
220}
221
222impl ModulePath {
223    pub fn new(path: &str) -> Self {
224        Self(path.to_owned())
225    }
226
227    pub fn join(&self, path: &str) -> Self {
228        let new = format!("{}::{path}", self.0);
229        Self(new)
230    }
231
232    pub fn as_str(&self) -> &str {
233        &self.0
234    }
235
236    pub fn last(&self) -> Option<ModulePath> {
237        self.0
238            .rsplit_once("::")
239            .map(|(_, last)| ModulePath::new(last))
240    }
241
242    pub fn parent(&self) -> Option<ModulePath> {
243        self.0
244            .rsplit_once("::")
245            .map(|(prefix, _)| ModulePath::new(prefix))
246    }
247}
248
249impl Display for ModulePath {
250    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
251        f.write_str(&self.0)
252    }
253}
254
255impl From<ModulePath> for String {
256    fn from(value: ModulePath) -> Self {
257        value.to_string()
258    }
259}
260
261impl From<&ModulePath> for String {
262    fn from(value: &ModulePath) -> Self {
263        value.to_string()
264    }
265}
266
267impl Sandbox {
268    pub fn setup(inner: &api::Sandbox, meta: &Metadata) -> Result<Self> {
269        let enabled = inner.enabled.unwrap_or(defaults::SANDBOX_ENABLED);
270        let follow_symlinks = inner
271            .follow_symlinks
272            .unwrap_or(defaults::SANDBOX_FIXTURES_FOLLOW_SYMLINKS);
273        let current_dir = std::env::current_dir().map_err(|error| {
274            Error::SandboxError(format!("Failed to detect current directory: {error}"))
275        })?;
276
277        let temp_dir = if enabled {
278            debug!("Creating sandbox");
279
280            let temp_dir = tempfile::tempdir().map_err(|error| {
281                Error::SandboxError(format!("Failed creating temporary directory: {error}"))
282            })?;
283
284            for fixture in &inner.fixtures {
285                if fixture.is_relative() {
286                    let absolute_path = make_absolute(&meta.project_root, fixture);
287                    copy_directory(&absolute_path, temp_dir.path(), follow_symlinks)?;
288                } else {
289                    copy_directory(fixture, temp_dir.path(), follow_symlinks)?;
290                };
291            }
292
293            trace!(
294                "Changing current directory to sandbox directory: '{}'",
295                temp_dir.path().display()
296            );
297
298            let path = temp_dir.path();
299            std::env::set_current_dir(path).map_err(|error| {
300                Error::SandboxError(format!(
301                    "Failed setting current directory to sandbox directory: '{error}'"
302                ))
303            })?;
304            Some(temp_dir)
305        } else {
306            debug!(
307                "Sandbox disabled: Running benchmarks in current directory '{}'",
308                current_dir.display()
309            );
310            None
311        };
312
313        Ok(Self {
314            current_dir,
315            temp_dir,
316        })
317    }
318
319    pub fn reset(self) -> Result<()> {
320        if let Some(temp_dir) = self.temp_dir {
321            std::env::set_current_dir(&self.current_dir).map_err(|error| {
322                Error::SandboxError(format!("Failed to reset current directory: {error}"))
323            })?;
324
325            if log_enabled!(Level::Debug) {
326                debug!("Removing temporary workspace");
327                if let Err(error) = temp_dir.close() {
328                    debug!("Error trying to delete temporary workspace: {error}");
329                }
330            } else {
331                _ = temp_dir.close();
332            }
333        }
334
335        Ok(())
336    }
337}
iai_callgrind_runner/runner/common.rs

iai_callgrind_runner/runner/
common.rs
