From 5de748aa9dff0a1bf8773cc08419ccca20af3394 Mon Sep 17 00:00:00 2001
From: John Nunley <jtnunley01@gmail.com>
Date: Mon, 10 Jul 2023 12:29:11 -0700
Subject: [PATCH] bench: Add benchmarks

---
 Cargo.toml       |  6 +++
 benches/waker.rs | 98 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)
 create mode 100644 benches/waker.rs

diff --git a/Cargo.toml b/Cargo.toml
index bf220d2..900d9c8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,4 +22,10 @@ exclude = ["/.*"]
 portable-atomic = { version = "1", optional = true, default-features = false }
 
 [dev-dependencies]
+criterion = { version = "0.4.0", default-features = false, features = ["cargo_bench_support"] }
 futures = "0.3.5"
+rayon = "1.7.0"
+
+[[bench]]
+name = "waker"
+harness = false
diff --git a/benches/waker.rs b/benches/waker.rs
new file mode 100644
index 0000000..199d82d
--- /dev/null
+++ b/benches/waker.rs
@@ -0,0 +1,98 @@
+use atomic_waker::AtomicWaker;
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkGroup, Criterion};
+
+use std::sync::Arc;
+use std::task::{Wake, Waker};
+
+enum Contention {
+    Low,
+    High,
+}
+
+impl Contention {
+    fn bench<M>(
+        self,
+        group: &mut BenchmarkGroup<'_, M>,
+        f: impl Fn(&AtomicWaker) + Send + Sync + Clone,
+    ) where
+        M: criterion::measurement::Measurement,
+    {
+        let waker = AtomicWaker::new();
+
+        match self {
+            Self::Low => {
+                // Run the benchmark with low contention.
+                group.bench_function("low contention", move |b| {
+                    b.iter(|| {
+                        for _ in 0..500 {
+                            f(&waker);
+                        }
+                    })
+                });
+            }
+
+            Self::High => {
+                // Run the benchmark with high contention.
+                group.bench_function("high contention", move |b| {
+                    b.iter(|| {
+                        use rayon::prelude::*;
+                        (0..100_000).into_par_iter().for_each(|_| f(&waker));
+                    })
+                });
+            }
+        }
+    }
+}
+
+fn run_lo_hi(c: &mut Criterion, name: &str, f: impl Fn(&AtomicWaker) + Send + Sync + Clone) {
+    let mut group = c.benchmark_group(name);
+
+    Contention::Low.bench(&mut group, f.clone());
+    Contention::High.bench(&mut group, f);
+
+    group.finish();
+}
+
+fn store_and_wake(c: &mut Criterion) {
+    run_lo_hi(c, "store and wake", |waker| {
+        let noop_waker = noop_waker();
+        waker.register(&noop_waker);
+        waker.wake();
+    });
+
+    run_lo_hi(c, "store and take", |waker| {
+        let noop_waker = noop_waker();
+        waker.register(&noop_waker);
+        black_box(waker.take());
+    });
+}
+
+fn wake_without_store(c: &mut Criterion) {
+    run_lo_hi(c, "wake without store", |waker| {
+        waker.wake();
+    });
+
+    run_lo_hi(c, "take without store", |waker| {
+        black_box(waker.take());
+    });
+}
+
+criterion_group!(benches, store_and_wake, wake_without_store);
+
+criterion_main!(benches);
+
+fn noop_waker() -> Waker {
+    struct Noop;
+
+    impl Wake for Noop {
+        fn wake(self: Arc<Self>) {
+            // Do nothing
+        }
+
+        fn wake_by_ref(self: &Arc<Self>) {
+            // Do nothing
+        }
+    }
+
+    Waker::from(Arc::new(Noop))
+}