fluxbench_core/bencher/
async_iter.rs

1//! Async `iter*` methods.
2//!
3//! Each batch runs one `Runtime::block_on` over all its calls.
4
5use super::async_body::runtime_context_error;
6use super::config::MeasureError;
7use super::engine::SampleRaw;
8use super::state::Bencher;
9use crate::measure::Timer;
10use std::future::Future;
11use std::hint::black_box;
12use std::sync::Arc;
13use tokio::runtime::{Builder, Runtime};
14
15impl Bencher {
16    /// Attach the runtime that `iter_async` runs on.
17    ///
18    /// The `#[flux::bench]` wrapper attaches the runtime it builds.
19    pub fn attach_runtime(&mut self, rt: Arc<Runtime>) {
20        self.runtime = Some(rt);
21    }
22
23    /// Record a runtime build error. `finish` returns it as `MeasureError::AsyncRuntime`.
24    pub fn attach_runtime_error(&mut self, message: impl Into<String>) {
25        self.record_error(MeasureError::AsyncRuntime(message.into()));
26    }
27
28    /// Measure an async closure. Each batch awaits `n` calls inside one `block_on`.
29    ///
30    /// The timer covers creating and awaiting each future.
31    /// It runs on the attached runtime.
32    /// Without one, it builds a current-thread runtime once and keeps it.
33    /// Call it from sync code or from a body that `drive_async_body` runs.
34    /// A call inside a runtime context records `MeasureError::AsyncRuntime` and runs nothing.
35    #[inline]
36    pub fn iter_async<T, F, Fut>(&mut self, f: F)
37    where
38        F: FnMut() -> Fut,
39        Fut: Future<Output = T>,
40    {
41        if !self.begin_iter() {
42            return;
43        }
44        if let Some(error) = runtime_context_error("iter_async") {
45            self.record_error(error);
46            return;
47        }
48        match self.runtime_or_build() {
49            Ok(rt) => self.run_on_runtime(&rt, f),
50            Err(error) => self.record_error(error),
51        }
52    }
53
54    /// Measure an async closure. Same as `iter_async`.
55    #[deprecated(note = "use `iter_async`")]
56    #[inline]
57    pub fn iter_async_standalone<T, F, Fut>(&mut self, f: F)
58    where
59        F: FnMut() -> Fut,
60        Fut: Future<Output = T>,
61    {
62        self.iter_async(f);
63    }
64
65    /// Return the attached runtime, or build a current-thread runtime and attach it.
66    fn runtime_or_build(&mut self) -> Result<Arc<Runtime>, MeasureError> {
67        if let Some(rt) = &self.runtime {
68            return Ok(Arc::clone(rt));
69        }
70        let rt = Builder::new_current_thread()
71            .enable_all()
72            .build()
73            .map(Arc::new)
74            .map_err(|e| MeasureError::AsyncRuntime(e.to_string()))?;
75        self.runtime = Some(Arc::clone(&rt));
76        Ok(rt)
77    }
78
79    /// Drive warmup and measurement with one `rt.block_on` per batch.
80    fn run_on_runtime<T, F, Fut>(&mut self, rt: &Runtime, mut f: F)
81    where
82        F: FnMut() -> Fut,
83        Fut: Future<Output = T>,
84    {
85        self.run_loop(|b, n| {
86            let b = &*b;
87            let (ns, cycles, alloc) = rt.block_on(async {
88                b.alloc_reset();
89                let timer = Timer::start();
90                for _ in 0..n {
91                    black_box(f().await);
92                }
93                let (ns, cycles) = timer.stop();
94                (ns, cycles, b.alloc_read())
95            });
96            let mut raw = SampleRaw::default();
97            raw.add(ns, cycles, alloc);
98            raw
99        });
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::bencher::async_body::drive_async_body;
107    use crate::bencher::{MeasureConfig, run_benchmark_loop};
108    use std::cell::Cell;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110    use std::time::Duration;
111
112    fn fixed(samples: usize, ips: u64) -> MeasureConfig {
113        MeasureConfig {
114            warmup_ns: 0,
115            measurement_ns: 10_000_000,
116            sample_count: samples,
117            fixed_iters_per_sample: Some(ips),
118            min_iterations: None,
119            max_iterations: None,
120            track_allocations: false,
121        }
122    }
123
124    fn current_thread() -> Arc<Runtime> {
125        Arc::new(
126            Builder::new_current_thread()
127                .enable_all()
128                .build()
129                .expect("runtime builds"),
130        )
131    }
132
133    #[test]
134    fn iter_async_without_runtime_batches_calls() {
135        let calls = Cell::new(0_u64);
136        let result = run_benchmark_loop(Bencher::with_config(fixed(4, 8)), |b| {
137            b.iter_async(|| {
138                calls.set(calls.get() + 1);
139                async { 1_u64 }
140            });
141        })
142        .expect("measurement succeeds");
143        assert_eq!(result.iterations, 32);
144        assert_eq!(calls.get(), 32);
145    }
146
147    #[test]
148    fn iter_async_without_runtime_builds_one_lazily() {
149        let mut bencher = Bencher::with_config(fixed(3, 2));
150        assert!(bencher.runtime.is_none());
151        bencher.iter_async(|| async { tokio::time::sleep(Duration::from_micros(1)).await });
152        let rt = bencher.runtime.clone().expect("lazy runtime attached");
153        // The bencher and this clone hold the only references.
154        assert_eq!(Arc::strong_count(&rt), 2);
155        let result = bencher.finish().expect("measurement succeeds");
156        assert_eq!(result.iterations, 6);
157    }
158
159    #[test]
160    fn iter_async_current_thread() {
161        let rt = current_thread();
162        let mut bencher = Bencher::with_config(fixed(3, 4));
163        bencher.attach_runtime(Arc::clone(&rt));
164        drive_async_body(&rt, async {
165            bencher.iter_async(|| tokio::time::sleep(Duration::from_micros(1)));
166        })
167        .expect("body runs");
168        let result = bencher.finish().expect("measurement succeeds");
169        assert_eq!(result.iterations, 12);
170    }
171
172    #[test]
173    fn iter_async_multi_thread() {
174        let rt = Arc::new(
175            Builder::new_multi_thread()
176                .worker_threads(2)
177                .enable_all()
178                .build()
179                .expect("runtime builds"),
180        );
181        let mut bencher = Bencher::with_config(fixed(3, 5));
182        bencher.attach_runtime(Arc::clone(&rt));
183        drive_async_body(&rt, async {
184            bencher.iter_async(|| async { tokio::task::yield_now().await });
185        })
186        .expect("body runs");
187        let result = bencher.finish().expect("measurement succeeds");
188        assert_eq!(result.iterations, 15);
189    }
190
191    #[test]
192    fn runtime_built_once() {
193        let started = Arc::new(AtomicUsize::new(0));
194        let counter = Arc::clone(&started);
195        let rt = Arc::new(
196            Builder::new_multi_thread()
197                .worker_threads(1)
198                .enable_all()
199                .on_thread_start(move || {
200                    counter.fetch_add(1, Ordering::SeqCst);
201                })
202                .build()
203                .expect("runtime builds"),
204        );
205        let mut bencher = Bencher::with_config(fixed(4, 4));
206        bencher.attach_runtime(Arc::clone(&rt));
207        drive_async_body(&rt, async {
208            // The worker thread starts before it runs this task.
209            tokio::spawn(async {}).await.expect("task joins");
210            bencher.iter_async(|| async { tokio::task::yield_now().await });
211        })
212        .expect("body runs");
213        let held = bencher.runtime.clone().expect("runtime attached");
214        assert!(Arc::ptr_eq(&held, &rt));
215        drop(held);
216        bencher.finish().expect("measurement succeeds");
217        // One worker thread. No blocking-pool thread runs in this test.
218        assert_eq!(started.load(Ordering::SeqCst), 1);
219        assert_eq!(Arc::strong_count(&rt), 1);
220    }
221
222    #[test]
223    fn body_setup_awaits_io_on_current_thread() {
224        let rt = current_thread();
225        let mut bencher = Bencher::with_config(fixed(2, 3));
226        bencher.attach_runtime(Arc::clone(&rt));
227        drive_async_body(&rt, async {
228            let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
229                .await
230                .expect("listener binds");
231            let addr = listener.local_addr().expect("listener has an address");
232            tokio::time::sleep(Duration::from_micros(1)).await;
233            bencher.iter_async(|| async move { black_box(addr.port()) });
234        })
235        .expect("body runs");
236        let result = bencher.finish().expect("measurement succeeds");
237        assert_eq!(result.iterations, 6);
238    }
239
240    #[test]
241    fn spawn_in_body_runs_during_iter_async() {
242        let rt = current_thread();
243        let mut bencher = Bencher::with_config(fixed(2, 4));
244        bencher.attach_runtime(Arc::clone(&rt));
245        drive_async_body(&rt, async {
246            let (tx, mut rx) = tokio::sync::mpsc::channel::<u64>(1);
247            let echo = tokio::spawn(async move {
248                let mut count = 0_u64;
249                while rx.recv().await.is_some() {
250                    count += 1;
251                }
252                count
253            });
254            bencher.iter_async(|| {
255                let tx = tx.clone();
256                async move { tx.send(1).await.expect("echo task receives") }
257            });
258            drop(tx);
259            let received = echo.await.expect("echo task joins");
260            assert_eq!(received, 8);
261        })
262        .expect("body runs");
263        bencher.finish().expect("measurement succeeds");
264    }
265
266    #[test]
267    fn iter_async_inside_runtime_context_records_error() {
268        let outer = current_thread();
269        let result = outer.block_on(async {
270            run_benchmark_loop(Bencher::with_config(fixed(2, 2)), |b| {
271                b.iter_async(|| async {});
272            })
273        });
274        assert!(matches!(result, Err(MeasureError::AsyncRuntime(_))));
275    }
276
277    #[test]
278    fn iter_async_inside_task_records_error() {
279        let outer = current_thread();
280        let result = outer.block_on(async {
281            tokio::spawn(async {
282                let mut bencher = Bencher::with_config(fixed(2, 2));
283                bencher.iter_async(|| async {});
284                bencher.finish().err()
285            })
286            .await
287            .expect("task joins")
288        });
289        assert!(matches!(result, Some(MeasureError::AsyncRuntime(_))));
290    }
291
292    #[test]
293    fn attach_runtime_error_is_reported() {
294        let mut bencher = Bencher::with_config(fixed(2, 2));
295        bencher.attach_runtime_error("no runtime");
296        let err = bencher.finish().err();
297        assert_eq!(
298            err,
299            Some(MeasureError::AsyncRuntime("no runtime".to_string()))
300        );
301    }
302
303    #[test]
304    #[allow(deprecated)]
305    fn iter_async_standalone_runs_fixed_samples() {
306        let result = run_benchmark_loop(Bencher::with_config(fixed(2, 4)), |b| {
307            b.iter_async_standalone(|| async { 2_u64 });
308        })
309        .expect("measurement succeeds");
310        assert_eq!(result.samples.len(), 2);
311        assert_eq!(result.iterations, 8);
312    }
313}
fluxbench_core/bencher/async_iter.rs

fluxbench_core/bencher/
async_iter.rs
