feat(criterion-compat): add iter_manual to have finer control over the iteration loop

Author: GuillaumeLagrange

crates/criterion_compat/benches/criterion_integration/iter_manual.rs

@@ -0,0 +1,115 @@
+use codspeed_criterion_compat::{criterion_group, Criterion, IterManualOptions};
+
+fn iter_manual_basic(c: &mut Criterion) {
+    // rounds=5 is below criterion's sample_size >= 10 floor — proves
+    // iter_manual fully bypasses the adaptive outer sampler.
+    c.bench_function("iter_manual_basic", |b| {
+        b.iter_manual(IterManualOptions::new(5, 100), || {
+            let mut s = 0u64;
+            for i in 0..32 {
+                s = s.wrapping_add(i);
+            }
+            s
+        });
+    });
+}
+
+fn iter_manual_with_setup(c: &mut Criterion) {
+    c.bench_function("iter_manual_setup", |b| {
+        b.iter_manual_setup(
+            IterManualOptions {
+                rounds: 3,
+                iterations: 100,
+                warmup_rounds: 2,
+            },
+            || (0u64..256).collect::<Vec<_>>(),
+            |input| input.iter().copied().sum::<u64>(),
+        );
+    });
+}
+
+fn iter_manual_with_teardown(c: &mut Criterion) {
+    c.bench_function("iter_manual_setup_teardown", |b| {
+        b.iter_manual_setup_teardown(
+            IterManualOptions {
+                rounds: 7,
+                iterations: 50,
+                warmup_rounds: 1,
+            },
+            || (0u64..128).collect::<Vec<_>>(),
+            |input| input.iter().copied().sum::<u64>(),
+            |input| drop(input),
+        );
+    });
+}
+
+#[cfg(feature = "async_futures")]
+fn iter_manual_async_basic(c: &mut Criterion) {
+    use codspeed_criterion_compat::async_executor::FuturesExecutor;
+    c.bench_function("iter_manual_async_basic", |b| {
+        b.to_async(FuturesExecutor)
+            .iter_manual(IterManualOptions::new(5, 100), || async {
+                let mut s = 0u64;
+                for i in 0..32 {
+                    s = s.wrapping_add(i);
+                }
+                s
+            });
+    });
+}
+
+#[cfg(feature = "async_futures")]
+fn iter_manual_async_with_setup(c: &mut Criterion) {
+    use codspeed_criterion_compat::async_executor::FuturesExecutor;
+    c.bench_function("iter_manual_async_setup", |b| {
+        b.to_async(FuturesExecutor).iter_manual_setup(
+            IterManualOptions {
+                rounds: 3,
+                iterations: 100,
+                warmup_rounds: 2,
+            },
+            || (0u64..256).collect::<Vec<_>>(),
+            |input| {
+                let sum = input.iter().copied().sum::<u64>();
+                async move { sum }
+            },
+        );
+    });
+}
+
+#[cfg(feature = "async_futures")]
+fn iter_manual_async_with_teardown(c: &mut Criterion) {
+    use codspeed_criterion_compat::async_executor::FuturesExecutor;
+    c.bench_function("iter_manual_async_setup_teardown", |b| {
+        b.to_async(FuturesExecutor).iter_manual_setup_teardown(
+            IterManualOptions {
+                rounds: 7,
+                iterations: 50,
+                warmup_rounds: 1,
+            },
+            || (0u64..128).collect::<Vec<_>>(),
+            |input| {
+                let sum = input.iter().copied().sum::<u64>();
+                async move { sum }
+            },
+            |input| async move { drop(input) },
+        );
+    });
+}
+
+#[cfg(not(feature = "async_futures"))]
+fn iter_manual_async_basic(_c: &mut Criterion) {}
+#[cfg(not(feature = "async_futures"))]
+fn iter_manual_async_with_setup(_c: &mut Criterion) {}
+#[cfg(not(feature = "async_futures"))]
+fn iter_manual_async_with_teardown(_c: &mut Criterion) {}
+
+criterion_group!(
+    benches,
+    iter_manual_basic,
+    iter_manual_with_setup,
+    iter_manual_with_teardown,
+    iter_manual_async_basic,
+    iter_manual_async_with_setup,
+    iter_manual_async_with_teardown,
+);

crates/criterion_compat/benches/criterion_integration/mod.rs

@@ -1,6 +1,7 @@
 pub mod compare_functions;
 pub mod custom_measurement;
 // pub mod external_process;
+pub mod iter_manual;
 pub mod iter_with_large_drop;
 pub mod iter_with_large_setup;
 pub mod iter_with_setup;

crates/criterion_compat/benches/criterion_integration_main.rs

@@ -5,6 +5,7 @@ mod criterion_integration;
 criterion_main! {
     criterion_integration::compare_functions::fibonaccis,
     // criterion_integration::external_process::benches, FIXME: Currently doesn't work
+    criterion_integration::iter_manual::benches,
     criterion_integration::iter_with_large_drop::benches,
     criterion_integration::iter_with_large_setup::benches,
     criterion_integration::iter_with_setup::benches,

crates/criterion_compat/criterion_fork/src/bencher.rs

@@ -42,6 +42,10 @@ pub struct Bencher<'a, M: Measurement = WallTime> {
     pub(crate) value: M::Value,        // The measured value
     pub(crate) measurement: &'a M,     // Reference to the measurement object
     pub(crate) elapsed_time: Duration, // How much time did it take to perform the iteration? Used for the warmup period.
+    // CodSpeed addition: when `iter_manual*` runs, it drives the full
+    // benchmark itself and deposits the per-round results here. The outer
+    // sampler in `routine.rs` detects this and skips its adaptive logic.
+    pub(crate) codspeed_manual: Option<crate::codspeed_iter_manual::ManualMeasurement>,
 }
 impl<'a, M: Measurement> Bencher<'a, M> {
     /// Times a `routine` by executing it many times and timing the total elapsed time.
@@ -459,10 +463,12 @@ impl<'a, M: Measurement> Bencher<'a, M> {
 }
 
 /// Async/await variant of the Bencher struct.
+// CodSpeed addition: fields are `pub(crate)` so the codspeed_iter_manual module
+// can destructure this struct.
 #[cfg(feature = "async")]
 pub struct AsyncBencher<'a, 'b, A: AsyncExecutor, M: Measurement = WallTime> {
-    b: &'b mut Bencher<'a, M>,
-    runner: A,
+    pub(crate) b: &'b mut Bencher<'a, M>,
+    pub(crate) runner: A,
 }
 #[cfg(feature = "async")]
 impl<'a, 'b, A: AsyncExecutor, M: Measurement> AsyncBencher<'a, 'b, A, M> {

crates/criterion_compat/criterion_fork/src/codspeed_iter_manual.rs

@@ -0,0 +1,208 @@
+//! CodSpeed addition: manual control over benchmark iteration counts.
+//!
+//! `iter_manual*` lets the user pin down the exact number of measurement
+//! rounds, inner iterations per round, and warmup rounds. It bypasses
+//! Criterion's adaptive sampler entirely — see `routine.rs::sample` for the
+//! short-circuit.
+
+use std::time::Instant;
+
+use codspeed::instrument_hooks::InstrumentHooks;
+
+#[cfg(feature = "async")]
+use crate::async_executor::AsyncExecutor;
+use crate::black_box;
+use crate::measurement::Measurement;
+#[cfg(feature = "async")]
+use crate::AsyncBencher;
+use crate::Bencher;
+
+#[cfg(feature = "async")]
+use std::future::Future;
+
+/// Options for the [`iter_manual`](Bencher::iter_manual) family.
+#[derive(Debug, Clone, Copy)]
+pub struct IterManualOptions {
+    /// Number of measurement rounds (each produces one sample).
+    pub rounds: u64,
+    /// Number of routine invocations inside the measured region of each round.
+    pub iterations: u64,
+    /// Number of unmeasured warmup rounds run before measurement starts.
+    pub warmup_rounds: u64,
+}
+
+impl IterManualOptions {
+    /// Build options with the given rounds and iterations; `warmup_rounds` defaults to 0.
+    pub fn new(rounds: u64, iterations: u64) -> Self {
+        Self {
+            rounds,
+            iterations,
+            warmup_rounds: 0,
+        }
+    }
+}
+
+/// Captured output of a manual run. Stored on the `Bencher` and read by
+/// `routine.rs::sample` to short-circuit the adaptive sampler.
+pub(crate) struct ManualMeasurement {
+    /// One entry per measurement round, in nanoseconds (or whatever the
+    /// measurement's `to_f64` returns).
+    pub samples: Vec<f64>,
+    /// Number of inner iterations per round.
+    pub iterations: u64,
+}
+
+impl<'a, M: Measurement> Bencher<'a, M> {
+    /// Run `routine` exactly `opts.iterations` times inside each of `opts.rounds`
+    /// measurement rounds, optionally preceded by `opts.warmup_rounds` unmeasured rounds.
+    ///
+    /// Criterion's adaptive sampler is bypassed for this benchmark.
+    #[inline(never)]
+    pub fn iter_manual<O, R>(&mut self, opts: IterManualOptions, mut routine: R)
+    where
+        R: FnMut() -> O,
+    {
+        self.iter_manual_setup_teardown(opts, || (), |_| routine(), |_| ());
+    }
+
+    /// Like [`iter_manual`](Self::iter_manual), with a `setup` closure producing
+    /// fresh input for each round. The routine borrows the input mutably.
+    #[inline(never)]
+    pub fn iter_manual_setup<I, O, S, R>(&mut self, opts: IterManualOptions, setup: S, routine: R)
+    where
+        S: FnMut() -> I,
+        R: FnMut(&mut I) -> O,
+    {
+        self.iter_manual_setup_teardown(opts, setup, routine, |_| ());
+    }
+
+    /// Like [`iter_manual_setup`](Self::iter_manual_setup), with a `teardown`
+    /// closure called after each round, outside the measured region.
+    #[inline(never)]
+    pub fn iter_manual_setup_teardown<I, O, S, R, T>(
+        &mut self,
+        opts: IterManualOptions,
+        mut setup: S,
+        mut routine: R,
+        mut teardown: T,
+    ) where
+        S: FnMut() -> I,
+        R: FnMut(&mut I) -> O,
+        T: FnMut(I),
+    {
+        self.iterated = true;
+
+        let bench_start = InstrumentHooks::current_timestamp();
+        let time_start = Instant::now();
+
+        for _ in 0..opts.warmup_rounds {
+            let mut input = black_box(setup());
+            for _ in 0..opts.iterations {
+                black_box(routine(&mut input));
+            }
+            teardown(input);
+        }
+
+        let mut samples = Vec::with_capacity(opts.rounds as usize);
+        for _ in 0..opts.rounds {
+            let mut input = black_box(setup());
+            let start = self.measurement.start();
+            for _ in 0..opts.iterations {
+                black_box(routine(&mut input));
+            }
+            let value = self.measurement.end(start);
+            teardown(input);
+            samples.push(self.measurement.to_f64(&value));
+        }
+
+        self.elapsed_time = time_start.elapsed();
+        let bench_end = InstrumentHooks::current_timestamp();
+        InstrumentHooks::instance().add_benchmark_timestamps(bench_start, bench_end);
+
+        self.codspeed_manual = Some(ManualMeasurement {
+            samples,
+            iterations: opts.iterations,
+        });
+    }
+}
+
+#[cfg(feature = "async")]
+impl<'a, 'b, A: AsyncExecutor, M: Measurement> AsyncBencher<'a, 'b, A, M> {
+    /// Async/await variant of [`Bencher::iter_manual`].
+    #[inline(never)]
+    pub fn iter_manual<O, R, F>(&mut self, opts: IterManualOptions, mut routine: R)
+    where
+        R: FnMut() -> F,
+        F: Future<Output = O>,
+    {
+        self.iter_manual_setup_teardown(opts, || (), |_| routine(), |_| std::future::ready(()));
+    }
+
+    /// Async/await variant of [`Bencher::iter_manual_setup`].
+    #[inline(never)]
+    pub fn iter_manual_setup<I, O, S, R, F>(
+        &mut self,
+        opts: IterManualOptions,
+        setup: S,
+        routine: R,
+    ) where
+        S: FnMut() -> I,
+        R: FnMut(&mut I) -> F,
+        F: Future<Output = O>,
+    {
+        self.iter_manual_setup_teardown(opts, setup, routine, |_| std::future::ready(()));
+    }
+
+    /// Async/await variant of [`Bencher::iter_manual_setup_teardown`].
+    #[inline(never)]
+    pub fn iter_manual_setup_teardown<I, O, S, R, T, RF, TF>(
+        &mut self,
+        opts: IterManualOptions,
+        mut setup: S,
+        mut routine: R,
+        mut teardown: T,
+    ) where
+        S: FnMut() -> I,
+        R: FnMut(&mut I) -> RF,
+        T: FnMut(I) -> TF,
+        RF: Future<Output = O>,
+        TF: Future<Output = ()>,
+    {
+        let AsyncBencher { b, runner } = self;
+        runner.block_on(async {
+            b.iterated = true;
+
+            let bench_start = InstrumentHooks::current_timestamp();
+            let time_start = Instant::now();
+
+            for _ in 0..opts.warmup_rounds {
+                let mut input = black_box(setup());
+                for _ in 0..opts.iterations {
+                    black_box(routine(&mut input).await);
+                }
+                teardown(input).await;
+            }
+
+            let mut samples = Vec::with_capacity(opts.rounds as usize);
+            for _ in 0..opts.rounds {
+                let mut input = black_box(setup());
+                let start = b.measurement.start();
+                for _ in 0..opts.iterations {
+                    black_box(routine(&mut input).await);
+                }
+                let value = b.measurement.end(start);
+                teardown(input).await;
+                samples.push(b.measurement.to_f64(&value));
+            }
+
+            b.elapsed_time = time_start.elapsed();
+            let bench_end = InstrumentHooks::current_timestamp();
+            InstrumentHooks::instance().add_benchmark_timestamps(bench_start, bench_end);
+
+            b.codspeed_manual = Some(ManualMeasurement {
+                samples,
+                iterations: opts.iterations,
+            });
+        });
+    }
+}

crates/criterion_compat/criterion_fork/src/lib.rs

@@ -53,6 +53,8 @@ mod benchmark;
 mod benchmark_group;
 pub mod async_executor;
 mod bencher;
+// CodSpeed addition: manual iteration control. See `codspeed_iter_manual.rs`.
+mod codspeed_iter_manual;
 mod connection;
 #[cfg(feature = "csv_output")]
 mod csv_report;
@@ -99,6 +101,8 @@ use crate::report::{BencherReport, CliReport, CliVerbosity, Report, ReportContex
 pub use crate::bencher::AsyncBencher;
 pub use crate::bencher::Bencher;
 pub use crate::benchmark_group::{BenchmarkGroup, BenchmarkId};
+// CodSpeed addition.
+pub use crate::codspeed_iter_manual::IterManualOptions;
 
 static DEBUG_ENABLED: Lazy<bool> = Lazy::new(|| std::env::var_os("CRITERION_DEBUG").is_some());
 static GNUPLOT_VERSION: Lazy<Result<Version, VersionError>> = Lazy::new(criterion_plot::version);