feat: add AsyncBencher support

Author: art049

crates/criterion_compat/Cargo.toml

@@ -20,6 +20,22 @@ criterion = "0.4"
 codspeed = { path = "../codspeed", version = "1.0.0" }
 colored = "2.0.0"
 
+futures = { version = "0.3", default_features = false, optional = true }
+smol = { version = "1.2", default-features = false, optional = true }
+tokio = { version = "1.0", default-features = false, features = [
+  "rt",
+], optional = true }
+async-std = { version = "1.9", optional = true }
+
+[features]
+async = ["futures", "criterion/async"]
+
+async_futures = ["criterion/async_futures", "futures/executor", "async"]
+async_smol    = ["criterion/async_smol", "smol", "async"]
+async_tokio   = ["criterion/async_tokio", "tokio", "async"]
+async_std     = ["criterion/async_std", "async-std", "async"]
+
+
 [[bench]]
 name = "criterion_example"
 harness = false

crates/criterion_compat/README.md

@@ -74,6 +74,6 @@ Checked: benches/bencher_example.rs::fib_20 (group: benches)
     Finished running 1 benchmark suite(s)
 ```
 
-### Not supported yet:
+### Not supported:
 
-- AsyncBencher
+- `iter_custom`

crates/criterion_compat/benches/criterion_integration/async_measurement_overhead.rs

@@ -0,0 +1,48 @@
+use criterion::{async_executor::FuturesExecutor, criterion_group, BatchSize, Criterion};
+
+fn some_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("async overhead");
+    group.bench_function("iter", |b| b.to_async(FuturesExecutor).iter(|| async { 1 }));
+    group.bench_function("iter_with_setup", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_with_setup(|| (), |_| async { 1 })
+    });
+    group.bench_function("iter_with_large_setup", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_with_large_setup(|| (), |_| async { 1 })
+    });
+    group.bench_function("iter_with_large_drop", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_with_large_drop(|| async { 1 })
+    });
+    group.bench_function("iter_batched_small_input", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_batched(|| (), |_| async { 1 }, BatchSize::SmallInput)
+    });
+    group.bench_function("iter_batched_large_input", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_batched(|| (), |_| async { 1 }, BatchSize::LargeInput)
+    });
+    group.bench_function("iter_batched_per_iteration", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_batched(|| (), |_| async { 1 }, BatchSize::PerIteration)
+    });
+    group.bench_function("iter_batched_ref_small_input", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_batched_ref(|| (), |_| async { 1 }, BatchSize::SmallInput)
+    });
+    group.bench_function("iter_batched_ref_large_input", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_batched_ref(|| (), |_| async { 1 }, BatchSize::LargeInput)
+    });
+    group.bench_function("iter_batched_ref_per_iteration", |b| {
+        b.to_async(FuturesExecutor).iter_batched_ref(
+            || (),
+            |_| async { 1 },
+            BatchSize::PerIteration,
+        )
+    });
+    group.finish();
+}
+
+criterion_group!(benches, some_benchmark);

crates/criterion_compat/benches/criterion_integration/mod.rs

@@ -8,3 +8,14 @@ pub mod measurement_overhead;
 pub mod sampling_mode;
 pub mod special_characters;
 pub mod with_inputs;
+
+#[cfg(feature = "async_futures")]
+pub mod async_measurement_overhead;
+
+#[cfg(not(feature = "async_futures"))]
+pub mod async_measurement_overhead {
+    use criterion::{criterion_group, Criterion};
+    fn some_benchmark(_c: &mut Criterion) {}
+
+    criterion_group!(benches, some_benchmark);
+}

crates/criterion_compat/benches/criterion_integration_main.rs

@@ -13,4 +13,5 @@ criterion_main! {
     criterion_integration::measurement_overhead::benches,
     criterion_integration::custom_measurement::benches,
     criterion_integration::sampling_mode::benches,
+    criterion_integration::async_measurement_overhead::benches,
 }

crates/criterion_compat/src/compat/bencher.rs

@@ -4,6 +4,11 @@ use codspeed::codspeed::{black_box, CodSpeed};
 use colored::Colorize;
 use criterion::BatchSize;
 
+#[cfg(feature = "async")]
+use criterion::async_executor::AsyncExecutor;
+#[cfg(feature = "async")]
+use std::future::Future;
+
 pub struct Bencher {
     codspeed: Rc<RefCell<CodSpeed>>,
     uri: String,
@@ -91,4 +96,111 @@ impl Bencher {
         drop(black_box(output));
         drop(black_box(input));
     }
+
+    #[cfg(feature = "async")]
+    pub fn to_async<A: AsyncExecutor>(&mut self, runner: A) -> AsyncBencher<A, M> {
+        AsyncBencher { b: self, runner }
+    }
+}
+
+pub struct AsyncBencher<A: AsyncExecutor> {
+    b: &mut Bencher,
+    runner: A,
+}
+
+impl AsyncBencher {
+    #[inline(never)]
+    pub fn iter<O, R, F>(&mut self, mut routine: R)
+    where
+        R: FnMut() -> F,
+        F: Future<Output = O>,
+    {
+        let AsyncBencher { b, runner } = self;
+        runner.block_on(async {
+            let mut codspeed = b.codspeed.borrow_mut();
+            codspeed.start_benchmark(b.uri.as_str());
+            black_box(routine().await);
+            codspeed.end_benchmark();
+        });
+    }
+
+    #[inline(never)]
+    pub fn iter_custom<R, F>(&mut self, mut routine: R)
+    where
+        R: FnMut(u64) -> F,
+        F: Future<Output = M::Value>,
+    {
+        println!(
+            "{} {} (CodSpeed doesn't support custom iterations)",
+            "Skipping:".to_string().yellow(),
+            self.uri.yellow(),
+        );
+    }
+
+    #[doc(hidden)]
+    pub fn iter_with_setup<I, O, S, R, F>(&mut self, setup: S, routine: R)
+    where
+        S: FnMut() -> I,
+        R: FnMut(I) -> F,
+        F: Future<Output = O>,
+    {
+        self.iter_batched(setup, routine, BatchSize::PerIteration);
+    }
+
+    pub fn iter_with_large_drop<O, R, F>(&mut self, mut routine: R)
+    where
+        R: FnMut() -> F,
+        F: Future<Output = O>,
+    {
+        self.iter_batched(|| (), |_| routine(), BatchSize::SmallInput);
+    }
+
+    #[doc(hidden)]
+    pub fn iter_with_large_setup<I, O, S, R, F>(&mut self, setup: S, routine: R)
+    where
+        S: FnMut() -> I,
+        R: FnMut(I) -> F,
+        F: Future<Output = O>,
+    {
+        self.iter_batched(setup, routine, BatchSize::NumBatches(1));
+    }
+
+    #[inline(never)]
+    pub fn iter_batched<I, O, S, R, F>(&mut self, mut setup: S, mut routine: R, size: BatchSize)
+    where
+        S: FnMut() -> I,
+        R: FnMut(I) -> F,
+        F: Future<Output = O>,
+    {
+        let AsyncBencher { b, runner } = self;
+        runner.block_on(async {
+            let mut codspeed = b.codspeed.borrow_mut();
+            let input = black_box(setup());
+            codspeed.start_benchmark(b.uri.as_str());
+            let output = routine(input);
+            codspeed.end_benchmark();
+            drop(black_box(output));
+        })
+    }
+
+    #[inline(never)]
+    pub fn iter_batched_ref<I, O, S, R, F>(&mut self, mut setup: S, mut routine: R, size: BatchSize)
+    where
+        S: FnMut() -> I,
+        R: FnMut(&mut I) -> F,
+        F: Future<Output = O>,
+    {
+        let AsyncBencher { b, runner } = self;
+        runner.block_on(async {
+            let mut codspeed = b.codspeed.borrow_mut();
+            let mut input = black_box(setup());
+
+            codspeed.start_benchmark(b.uri.as_str());
+            let output = routine(&mut input);
+            codspeed.end_benchmark();
+
+            drop(black_box(output));
+            drop(black_box(input));
+        });
+    }
 }