bench_rs/
bencher.rs

1use std::future::Future;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::time::Instant;
4
5use crate::{Stats, Step};
6use crate::fmt_thousands_sep;
7use crate::timing_future::TimingFuture;
8#[cfg(feature = "track-allocator")]
9use crate::track_allocator::GLOBAL;
10
11pub struct Bencher {
12    pub name: String,
13    pub count: usize,
14    pub steps: Vec<Step>,
15    pub bytes: usize,
16    pub n: usize,
17    pub poll: usize,
18    pub format_fn: fn(&Stats, &Bencher),
19
20    pub mem_track: (&'static AtomicUsize, &'static AtomicUsize)
21}
22
23impl Bencher {
24    #[cfg(feature = "track-allocator")]
25    pub fn new(name: impl AsRef<str>, count: usize, bytes: usize) -> Self {
26        Bencher {
27            name: name.as_ref().to_owned(),
28            count,
29            steps: Vec::with_capacity(count),
30            bytes,
31            n: 0,
32            poll: 0,
33            format_fn: |s, b| Self::default_format(s, b),
34
35            mem_track: (GLOBAL.counter(), GLOBAL.peak())
36        }
37    }
38
39    #[cfg(not(feature = "track-allocator"))]
40    pub fn new(name: impl AsRef<str>, count: usize, bytes: usize, counter: &'static AtomicUsize, peak: &'static AtomicUsize) -> Self {
41        Bencher {
42            name: name.as_ref().to_owned(),
43            count,
44            steps: Vec::with_capacity(count),
45            bytes,
46            n: 0,
47            poll: 0,
48            format_fn: |s, b| Self::default_format(s, b),
49
50            mem_track: (counter, peak)
51        }
52    }
53
54    // (time, memory_usage)
55    pub fn bench_once<T>(&self, f: &mut impl FnMut() -> T, n: usize) -> (u128, usize) {
56        let now = Instant::now();
57        self.reset_mem();
58
59        for _ in 0..n {
60            let _output = f();
61        }
62
63        (now.elapsed().as_nanos(), self.get_mem_peak())
64    }
65
66    pub fn iter<T>(&mut self, mut f: impl FnMut() -> T) {
67        let single = self.bench_once(&mut f, 1).0;
68        // 1_000_000ns : 1ms
69        self.n = (1_000_000 / single.max(1)).max(1) as usize;
70        (0..self.count).for_each(|_| {
71            let res = self.bench_once(&mut f, self.n);
72            self.steps.push(Step {
73                time: res.0 / self.n as u128,
74                mem: res.1 / self.n
75            })
76        });
77    }
78
79    pub fn async_iter<'a, T, Fut: Future<Output=T>>(&'a mut self, mut f: impl FnMut() -> Fut + 'a) -> impl Future + 'a {
80        async move {
81            let single = TimingFuture::new(f()).await.elapsed_time.as_nanos();
82            // 1_000_000ns : 1ms
83            self.n = (1_000_000 / single.max(1)).max(1) as usize;
84
85            let mut polls = Vec::with_capacity(self.count);
86
87            for _ in 0..self.count {
88                let mut mtime = 0u128;
89                self.reset_mem();
90                
91                for _ in 0..self.n {
92                    let tf = TimingFuture::new(f()).await;
93                    mtime += tf.elapsed_time.as_nanos();
94                    polls.push(tf.poll);
95                }
96
97                self.steps.push(Step {
98                    time: mtime / self.n as u128,
99                    mem: self.get_mem_peak() / self.n
100                });
101            }
102
103            self.poll = polls.iter().sum::<usize>() / polls.len();
104        }
105    }
106
107    pub fn finish(&self) {
108        let stats = Stats::from(&self.steps);
109        (self.format_fn)(&stats, self)
110    }
111
112    pub fn reset_mem(&self) {
113        self.mem_track.0.store(0, Ordering::SeqCst);
114        self.mem_track.1.store(0, Ordering::SeqCst);
115    }
116
117    pub fn get_mem_peak(&self) -> usize {
118        self.mem_track.1.load(Ordering::SeqCst)
119    }
120
121    fn default_format(stats: &Stats, bencher: &Bencher) {
122        bunt::println!(
123            "{[bg:white+blue+bold]} ... {[green+underline]} ns/iter (+/- {[red+underline]}) = {[yellow+underline]:.2} MB/s\
124            \n\t memory usage: {[green+underline]} bytes/iter (+/- {[red+underline]})\
125            \n\t @Total: {[magenta]} * {[white]} iters\
126            {[bold]}",
127             &bencher.name,
128             fmt_thousands_sep(stats.times_average, ','),
129             fmt_thousands_sep(stats.times_max - stats.times_min, ','),
130             (bencher.bytes as f64 * (1_000_000_000f64 / stats.times_average as f64)) / 1000f64 / 1000f64,
131
132             fmt_thousands_sep(stats.mem_average, ','),
133             fmt_thousands_sep(stats.mem_max - stats.mem_min, ','),
134
135             bencher.count,
136             bencher.n,
137
138             if bencher.poll > 0 {
139                format!(
140                    "\n\t @avg {} polls ",
141                    bencher.poll
142                 )
143             } else {
144                String::new()
145             },
146        );
147    }
148}