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,43 @@
+use codspeed_criterion_compat::{criterion_group, Criterion, IterManualOptions};
+
+fn iter_manual_simple(c: &mut Criterion) {
+    c.bench_function("iter_manual_simple", |b| {
+        b.iter_manual(
+            IterManualOptions::new().rounds(3).iters(5).warmup(1),
+            || std::thread::sleep(std::time::Duration::from_millis(100)),
+        );
+    });
+}
+
+fn iter_manual_with_external_setup(c: &mut Criterion) {
+    c.bench_function("iter_manual_with_external_setup", |b| {
+        // Setup deliberately does a chunk of work so it stands out in flamegraphs.
+        // The measured region should NOT include this work.
+        let input: Vec<u64> = (0..10_000u64).map(|i| i.wrapping_mul(31)).collect();
+        b.iter_manual(
+            IterManualOptions::new().rounds(3).iters(50).warmup(1),
+            || input.iter().copied().sum::<u64>(),
+        );
+    });
+}
+
+#[cfg(feature = "async_futures")]
+fn iter_manual_async(c: &mut Criterion) {
+    use codspeed_criterion_compat::async_executor::FuturesExecutor;
+    c.bench_function("iter_manual_async", |b| {
+        b.to_async(FuturesExecutor).iter_manual(
+            IterManualOptions::new().rounds(3).iters(100).warmup(2),
+            || async { (0u64..256).sum::<u64>() },
+        );
+    });
+}
+
+#[cfg(not(feature = "async_futures"))]
+fn iter_manual_async(_c: &mut Criterion) {}
+
+criterion_group!(
+    benches,
+    iter_manual_simple,
+    iter_manual_with_external_setup,
+    iter_manual_async,
+);

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,188 @@
+//! CodSpeed addition: manual control over benchmark sampling.
+//!
+//! `iter_manual` lets the user pin down the exact number of measurement rounds
+//! and iterations per round, bypassing criterion's adaptive sampler. See
+//! `routine.rs::sample` for the short-circuit that picks up the result.
+
+use std::time::{Duration, 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 [`Bencher::iter_manual`].
+#[derive(Debug, Clone, Copy)]
+pub struct IterManualOptions {
+    rounds: u64,
+    iters: u64,
+    warmup_rounds: u64,
+}
+
+impl Default for IterManualOptions {
+    fn default() -> Self {
+        Self {
+            rounds: 1,
+            iters: 1,
+            warmup_rounds: 0,
+        }
+    }
+}
+
+impl IterManualOptions {
+    /// Start with defaults: 1 round, 1 iteration per round, 0 warmup rounds.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Number of measurement rounds (each produces one sample).
+    #[must_use]
+    pub fn rounds(mut self, rounds: u64) -> Self {
+        self.rounds = rounds;
+        self
+    }
+
+    /// Number of routine invocations inside a measurement round.
+    #[must_use]
+    pub fn iters(mut self, iters: u64) -> Self {
+        self.iters = iters;
+        self
+    }
+
+    /// Number of unmeasured warmup rounds run before measurement starts.
+    #[must_use]
+    pub fn warmup(mut self, warmup_rounds: u64) -> Self {
+        self.warmup_rounds = warmup_rounds;
+        self
+    }
+}
+
+/// 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 the units of `Measurement::to_f64`.
+    pub samples: Vec<f64>,
+    /// Number of routine invocations per round.
+    pub iterations: u64,
+}
+
+impl<'a, M: Measurement> Bencher<'a, M> {
+    /// Run `routine` with a precise schedule: `opts.rounds` measurement rounds
+    /// (each producing one sample), each consisting of `opts.iters` calls to
+    /// `routine`. Optionally preceded by `opts.warmup_rounds` unmeasured rounds
+    /// of the same shape.
+    ///
+    /// Criterion's adaptive sampler is bypassed for this benchmark.
+    #[inline(never)]
+    pub fn iter_manual<O, R>(&mut self, opts: IterManualOptions, routine: R)
+    where
+        R: FnMut() -> O,
+    {
+        self.__codspeed_root_frame__iter_manual(opts, routine);
+    }
+
+    #[inline(never)]
+    #[allow(missing_docs, non_snake_case)]
+    pub fn __codspeed_root_frame__iter_manual<O, R>(
+        &mut self,
+        opts: IterManualOptions,
+        mut routine: R,
+    ) where
+        R: FnMut() -> O,
+    {
+        self.iterated = true;
+
+        for _ in 0..opts.warmup_rounds {
+            for _ in 0..opts.iters {
+                black_box(routine());
+            }
+        }
+
+        self.elapsed_time = Duration::ZERO;
+        let mut samples = Vec::with_capacity(opts.rounds as usize);
+        for _ in 0..opts.rounds {
+            let bench_start = InstrumentHooks::current_timestamp();
+            let round_start = Instant::now();
+            let start = self.measurement.start();
+            for _ in 0..opts.iters {
+                black_box(routine());
+            }
+            let value = self.measurement.end(start);
+            self.elapsed_time += round_start.elapsed();
+            let bench_end = InstrumentHooks::current_timestamp();
+            InstrumentHooks::instance().add_benchmark_timestamps(bench_start, bench_end);
+
+            samples.push(self.measurement.to_f64(&value));
+        }
+
+        self.codspeed_manual = Some(ManualMeasurement {
+            samples,
+            iterations: opts.iters,
+        });
+    }
+}
+
+#[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, routine: R)
+    where
+        R: FnMut() -> F,
+        F: Future<Output = O>,
+    {
+        self.__codspeed_root_frame__iter_manual(opts, routine);
+    }
+
+    #[inline(never)]
+    #[allow(missing_docs, non_snake_case)]
+    pub fn __codspeed_root_frame__iter_manual<O, R, F>(
+        &mut self,
+        opts: IterManualOptions,
+        mut routine: R,
+    ) where
+        R: FnMut() -> F,
+        F: Future<Output = O>,
+    {
+        let AsyncBencher { b, runner } = self;
+        runner.block_on(async {
+            b.iterated = true;
+
+            for _ in 0..opts.warmup_rounds {
+                for _ in 0..opts.iters {
+                    black_box(routine().await);
+                }
+            }
+
+            b.elapsed_time = Duration::ZERO;
+            let mut samples = Vec::with_capacity(opts.rounds as usize);
+            for _ in 0..opts.rounds {
+                let bench_start = InstrumentHooks::current_timestamp();
+                let round_start = Instant::now();
+                let start = b.measurement.start();
+                for _ in 0..opts.iters {
+                    black_box(routine().await);
+                }
+                let value = b.measurement.end(start);
+                b.elapsed_time += round_start.elapsed();
+                let bench_end = InstrumentHooks::current_timestamp();
+                InstrumentHooks::instance().add_benchmark_timestamps(bench_start, bench_end);
+
+                samples.push(b.measurement.to_f64(&value));
+            }
+
+            b.codspeed_manual = Some(ManualMeasurement {
+                samples,
+                iterations: opts.iters,
+            });
+        });
+    }
+}

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);