Add support for benchmarking async functions (#4823)

Author: Spxg

CHANGELOG.md

@@ -29,6 +29,7 @@
 
 * Added benchmark support to `wasm-bindgen-test`.
   [#4812](https://github.com/wasm-bindgen/wasm-bindgen/pull/4812)
+  [#4823](https://github.com/wasm-bindgen/wasm-bindgen/pull/4823)
 
 ### Fixed
 

benches/js_value_from.rs

@@ -3,7 +3,7 @@ use wasm_bindgen_test::{wasm_bindgen_bench, Criterion};
 
 #[wasm_bindgen_bench]
 fn bench_js_value_from(c: &mut Criterion) {
-    c.bench_function("bench_from_str", |b| {
+    c.bench_function("bench_js_value_from_str", |b| {
         b.iter(|| JsValue::from_str("42"));
     });
 }

crates/cli/src/wasm_bindgen_test_runner.rs

@@ -361,10 +361,14 @@ fn rmain(cli: Cli) -> anyhow::Result<()> {
         PathBuf::from(path)
     } else {
         // such as `js-sys/target/wbg_benchmark.json`
-        env::current_dir()
+        let path = env::current_dir()
             .context("Failed to get current dir")?
-            .join("target")
-            .join("wbg_benchmark.json")
+            .join("target");
+        // crates in the workspace that do not have a target dir.
+        if cli.bench {
+            fs::create_dir_all(&path)?;
+        }
+        path.join("wbg_benchmark.json")
     };
 
     // The debug here means adding some assertions and some error messages to the generated js

crates/cli/tests/reference/js-namespace-export.bg.js

@@ -212,7 +212,7 @@ function concat(a, b) {
  * @param {string} s
  * @returns {string}
  */
-function uppercase2(s) {
+function uppercase(s) {
     let deferred2_0;
     let deferred2_1;
     try {
@@ -231,7 +231,7 @@ const _default = {};
 _default.Counter = Counter;
 _default.concat = concat;
 _default.uppercase = {};
-_default.uppercase.uppercase = uppercase2;
+_default.uppercase.uppercase = uppercase;
 export default _default;
 
 /**
@@ -402,7 +402,7 @@ types.http.HttpStatus = HttpStatus;
  * @param {string} s
  * @returns {string}
  */
-function uppercase(s) {
+function uppercase2(s) {
     let deferred2_0;
     let deferred2_1;
     try {
@@ -419,7 +419,7 @@ function uppercase(s) {
 
 export const utils = {};
 utils.string = {};
-utils.string.uppercase = uppercase;
+utils.string.uppercase = uppercase2;
 
 export function __wbg___wbindgen_throw_b855445ff6a94295(arg0, arg1) {
     throw new Error(getStringFromWasm0(arg0, arg1));

crates/cli/tests/reference/js-namespace-export.d.ts

@@ -25,13 +25,13 @@ declare class Counter {
 
 declare function concat(a: string, b: string): string;
 
-declare function uppercase2(s: string): string;
+declare function uppercase(s: string): string;
 
 declare let _default: {
   Counter: typeof Counter,
   concat: typeof concat,
   uppercase: {
-    uppercase: typeof uppercase2,
+    uppercase: typeof uppercase,
   },
 };
 export default _default;
@@ -93,10 +93,10 @@ export let types: {
   },
 };
 
-declare function uppercase(s: string): string;
+declare function uppercase2(s: string): string;
 
 export let utils: {
   string: {
-    uppercase: typeof uppercase,
+    uppercase: typeof uppercase2,
   },
 };

crates/futures/benches/promise.rs

@@ -0,0 +1,13 @@
+use js_sys::Promise;
+use wasm_bindgen::JsValue;
+use wasm_bindgen_futures::JsFuture;
+use wasm_bindgen_test::{wasm_bindgen_bench, Criterion};
+
+#[wasm_bindgen_bench]
+async fn bench_promise(c: &mut Criterion) {
+    c.bench_async_function("bench_promise_to_future", |b| {
+        let f = b.iter_future(|| JsFuture::from(Promise::resolve(&JsValue::from(42))));
+        Box::pin(f)
+    })
+    .await;
+}

crates/test/Cargo.toml

@@ -21,6 +21,7 @@ wasm-bindgen-futures = { path = '../futures', version = '=0.4.55', default-featu
 wasm-bindgen-test-macro = { path = '../test-macro', version = '=0.3.55' }
 
 # benchmark required start
+async-trait = "0.1.89"
 cast = "0.3"
 libm = "0.2.11"
 nu-ansi-term = { version = "0.50", default-features = false }

crates/test/src/rt/criterion/analysis.rs

@@ -15,7 +15,7 @@ use super::{baseline, compare, Criterion, SavedSample};
 use alloc::vec::Vec;
 
 // Common analysis procedure
-pub(crate) fn common<M: Measurement>(
+pub(crate) async fn common<M: Measurement>(
     id: &BenchmarkId,
     routine: &mut dyn Routine<M>,
     config: &BenchmarkConfig,
@@ -24,7 +24,9 @@ pub(crate) fn common<M: Measurement>(
     criterion.report.benchmark_start(id);
 
     let (sampling_mode, iters, times);
-    let sample = routine.sample(&criterion.measurement, id, config, criterion);
+    let sample = routine
+        .sample(&criterion.measurement, id, config, criterion)
+        .await;
     sampling_mode = sample.0;
     iters = sample.1;
     times = sample.2;

crates/test/src/rt/criterion/bencher.rs

@@ -1,4 +1,5 @@
 use super::measurement::Measurement;
+use core::future::Future;
 use core::hint::black_box;
 use core::time::Duration;
 
@@ -64,4 +65,33 @@ impl<'a, M: Measurement> Bencher<'a, M> {
         self.value = end;
         self.elapsed_time = end;
     }
+
+    /// Times a `routine` by executing it many times and timing the total elapsed time.
+    ///
+    /// Prefer this timing loop when `routine` returns a value that doesn't have a destructor.
+    ///
+    /// # Timing model
+    ///
+    /// Note that the `Bencher` also times the time required to destroy the output of `routine()`.
+    /// Therefore prefer this timing loop when the runtime of `mem::drop(O)` is negligible compared
+    /// to the runtime of the `routine`.
+    ///
+    /// ```text
+    /// elapsed = Instant::now + iters * (routine + mem::drop(O) + Range::next)
+    /// ```
+    #[inline(never)]
+    pub async fn iter_future<O, R, Fut>(&mut self, mut routine: R)
+    where
+        R: FnMut() -> Fut,
+        Fut: Future<Output = O>,
+    {
+        self.iterated = true;
+        let start = self.measurement.start();
+        for _ in 0..self.iters {
+            black_box(routine().await);
+        }
+        let end = self.measurement.end(start);
+        self.value = end;
+        self.elapsed_time = end;
+    }
 }

crates/test/src/rt/criterion/mod.rs

@@ -36,10 +36,15 @@ mod report;
 mod routine;
 mod stats;
 
+use core::future::Future;
+use core::pin::Pin;
+use core::ptr;
+use core::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
 use core::time::Duration;
 use libm::{ceil, sqrt};
 use serde::{Deserialize, Serialize};
 
+use alloc::boxed::Box;
 use alloc::vec;
 use alloc::vec::Vec;
 use benchmark::BenchmarkConfig;
@@ -249,6 +254,7 @@ impl<M: Measurement> Criterion<M> {
         self
     }
 }
+
 impl<M> Criterion<M>
 where
     M: Measurement + 'static,
@@ -264,7 +270,7 @@ where
     /// fn bench(c: &mut Criterion) {
     ///     // Setup (construct data, allocate memory, etc)
     ///     c.bench_function(
-    ///         "function_name",
+    ///         "bench desc",
     ///         |b| b.iter(|| {
     ///             // Code to benchmark goes here
     ///         }),
@@ -274,9 +280,69 @@ where
     pub fn bench_function<F>(&mut self, desc: &str, f: F) -> &mut Criterion<M>
     where
         F: FnMut(&mut Bencher<'_, M>),
+    {
+        const NOOP: RawWaker = {
+            const VTABLE: RawWakerVTable = RawWakerVTable::new(
+                // Cloning just returns a new no-op raw waker
+                |_| NOOP,
+                // `wake` does nothing
+                |_| {},
+                // `wake_by_ref` does nothing
+                |_| {},
+                // Dropping does nothing as we don't allocate anything
+                |_| {},
+            );
+            RawWaker::new(ptr::null(), &VTABLE)
+        };
+
+        // bench_function never be pending
+        fn block_on(f: impl Future<Output = ()>) {
+            let waker = unsafe { Waker::from_raw(NOOP) };
+            let mut ctx = Context::from_waker(&waker);
+            match core::pin::pin!(f).poll(&mut ctx) {
+                Poll::Ready(_) => (),
+                // sync functions not be pending
+                Poll::Pending => unreachable!(),
+            }
+        }
+
+        let id = report::BenchmarkId::new(desc.into());
+        block_on(analysis::common(
+            &id,
+            &mut routine::Function::new(f),
+            &self.config,
+            self,
+        ));
+
+        self
+    }
+
+    /// Benchmarks a future.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use wasm_bindgen_test::{Criterion, wasm_bindgen_bench};
+    ///
+    /// #[wasm_bindgen_bench]
+    /// async fn bench(c: &mut Criterion) {
+    ///     // Setup (construct data, allocate memory, etc)
+    ///     c.bench_async_function(
+    ///         "bench desc",
+    ///         Box::pin(
+    ///             b.iter_future(|| async {
+    ///                 // Code to benchmark goes here
+    ///             })
+    ///         )
+    ///     ).await;
+    /// }
+    /// ```
+    pub async fn bench_async_function<F>(&mut self, desc: &str, f: F) -> &mut Criterion<M>
+    where
+        for<'b> F: FnMut(&'b mut Bencher<'_, M>) -> Pin<Box<dyn Future<Output = ()> + 'b>>,
     {
         let id = report::BenchmarkId::new(desc.into());
-        analysis::common(&id, &mut routine::Function::new(f), &self.config, self);
+        analysis::common(&id, &mut routine::AsyncFunction::new(f), &self.config, self).await;
         self
     }
 }