criterion/
program.rs

1use std::fmt;
2use std::io::BufReader;
3use std::marker::PhantomData;
4use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, Stdio};
5use std::time::{Duration, Instant};
6
7use routine::Routine;
8use DurationExt;
9
10// A two-way channel to the standard streams of a child process
11pub struct Program {
12    buffer: String,
13    stdin: ChildStdin,
14    // NB Don't move the `stdin` field, because it must be dropped first
15    _child: Child,
16    stderr: ChildStderr,
17    stdout: BufReader<ChildStdout>,
18}
19
20impl Program {
21    pub fn spawn(cmd: &mut Command) -> Program {
22        cmd.stderr(Stdio::piped());
23        cmd.stdin(Stdio::piped());
24        cmd.stdout(Stdio::piped());
25
26        let mut child = match cmd.spawn() {
27            Err(e) => panic!("`{:?}`: {}", cmd, e),
28            Ok(child) => child,
29        };
30
31        Program {
32            buffer: String::new(),
33            stderr: child.stderr.take().unwrap(),
34            stdin: child.stdin.take().unwrap(),
35            stdout: BufReader::new(child.stdout.take().unwrap()),
36            _child: child,
37        }
38    }
39
40    pub fn send<T>(&mut self, line: T) -> &mut Program
41    where
42        T: fmt::Display,
43    {
44        use std::io::Write;
45
46        match writeln!(&mut self.stdin, "{}", line) {
47            Err(e) => panic!("`write into child stdin`: {}", e),
48            Ok(_) => self,
49        }
50    }
51
52    pub fn recv(&mut self) -> &str {
53        use std::io::{BufRead, Read};
54
55        self.buffer.clear();
56
57        match self.stdout.read_line(&mut self.buffer) {
58            Err(e) => {
59                self.buffer.clear();
60
61                match self.stderr.read_to_string(&mut self.buffer) {
62                    Err(e) => {
63                        panic!("`read from child stderr`: {}", e);
64                    }
65                    Ok(_) => {
66                        println!("stderr:\n{}", self.buffer);
67                    }
68                }
69
70                panic!("`read from child stdout`: {}", e);
71            }
72            Ok(_) => &self.buffer,
73        }
74    }
75
76    fn bench(&mut self, iters: &[u64]) -> Vec<f64> {
77        let mut n = 0;
78        for iters in iters {
79            self.send(iters);
80            n += 1;
81        }
82
83        (0..n)
84            .map(|_| {
85                let msg = self.recv();
86                let msg = msg.trim();
87
88                let elapsed: u64 = msg.parse().expect("Couldn't parse program output");
89                elapsed as f64
90            })
91            .collect()
92    }
93
94    fn warm_up(&mut self, how_long_ns: Duration) -> (u64, u64) {
95        let mut iters = 1;
96
97        let mut total_iters = 0;
98        let start = Instant::now();
99        loop {
100            self.send(iters).recv();
101
102            total_iters += iters;
103            let elapsed = start.elapsed();
104            if elapsed > how_long_ns {
105                return (elapsed.to_nanos(), total_iters);
106            }
107
108            iters *= 2;
109        }
110    }
111}
112
113impl Routine<()> for Command {
114    fn start(&mut self, _: &()) -> Option<Program> {
115        Some(Program::spawn(self))
116    }
117
118    fn bench(&mut self, program: &mut Option<Program>, iters: &[u64], _: &()) -> Vec<f64> {
119        let program = program.as_mut().unwrap();
120        program.bench(iters)
121    }
122
123    fn warm_up(
124        &mut self,
125        program: &mut Option<Program>,
126        how_long_ns: Duration,
127        _: &(),
128    ) -> (u64, u64) {
129        let program = program.as_mut().unwrap();
130        program.warm_up(how_long_ns)
131    }
132}
133
134pub struct CommandFactory<F, T>
135where
136    F: FnMut(&T) -> Command + 'static,
137{
138    f: F,
139    _phantom: PhantomData<T>,
140}
141impl<F, T> CommandFactory<F, T>
142where
143    F: FnMut(&T) -> Command + 'static,
144{
145    pub fn new(f: F) -> CommandFactory<F, T> {
146        CommandFactory {
147            f,
148            _phantom: PhantomData,
149        }
150    }
151}
152
153impl<F, T> Routine<T> for CommandFactory<F, T>
154where
155    F: FnMut(&T) -> Command + 'static,
156{
157    fn start(&mut self, parameter: &T) -> Option<Program> {
158        let mut command = (self.f)(parameter);
159        Some(Program::spawn(&mut command))
160    }
161
162    fn bench(&mut self, program: &mut Option<Program>, iters: &[u64], _: &T) -> Vec<f64> {
163        let program = program.as_mut().unwrap();
164        program.bench(iters)
165    }
166
167    fn warm_up(
168        &mut self,
169        program: &mut Option<Program>,
170        how_long_ns: Duration,
171        _: &T,
172    ) -> (u64, u64) {
173        let program = program.as_mut().unwrap();
174        program.warm_up(how_long_ns)
175    }
176}