1use std::sync::atomic::{AtomicUsize, Ordering};
2use std::sync::Arc;
3
4use tokio::sync::Notify;
5
6use criterion::measurement::WallTime;
7use criterion::{criterion_group, criterion_main, BenchmarkGroup, Criterion};
8
9fn rt() -> tokio::runtime::Runtime {
10 tokio::runtime::Builder::new_multi_thread()
11 .worker_threads(6)
12 .build()
13 .unwrap()
14}
15
16fn notify_waiters<const N_WAITERS: usize>(g: &mut BenchmarkGroup<WallTime>) {
17 let rt = rt();
18 let notify = Arc::new(Notify::new());
19 let counter = Arc::new(AtomicUsize::new(0));
20 for _ in 0..N_WAITERS {
21 rt.spawn({
22 let notify = notify.clone();
23 let counter = counter.clone();
24 async move {
25 loop {
26 notify.notified().await;
27 counter.fetch_add(1, Ordering::Relaxed);
28 }
29 }
30 });
31 }
32
33 const N_ITERS: usize = 500;
34 g.bench_function(N_WAITERS.to_string(), |b| {
35 b.iter(|| {
36 counter.store(0, Ordering::Relaxed);
37 loop {
38 notify.notify_waiters();
39 if counter.load(Ordering::Relaxed) >= N_ITERS {
40 break;
41 }
42 }
43 })
44 });
45}
46
47fn notify_one<const N_WAITERS: usize>(g: &mut BenchmarkGroup<WallTime>) {
48 let rt = rt();
49 let notify = Arc::new(Notify::new());
50 let counter = Arc::new(AtomicUsize::new(0));
51 for _ in 0..N_WAITERS {
52 rt.spawn({
53 let notify = notify.clone();
54 let counter = counter.clone();
55 async move {
56 loop {
57 notify.notified().await;
58 counter.fetch_add(1, Ordering::Relaxed);
59 }
60 }
61 });
62 }
63
64 const N_ITERS: usize = 500;
65 g.bench_function(N_WAITERS.to_string(), |b| {
66 b.iter(|| {
67 counter.store(0, Ordering::Relaxed);
68 loop {
69 notify.notify_one();
70 if counter.load(Ordering::Relaxed) >= N_ITERS {
71 break;
72 }
73 }
74 })
75 });
76}
77
78fn bench_notify_one(c: &mut Criterion) {
79 let mut group = c.benchmark_group("notify_one");
80 notify_one::<10>(&mut group);
81 notify_one::<50>(&mut group);
82 notify_one::<100>(&mut group);
83 notify_one::<200>(&mut group);
84 notify_one::<500>(&mut group);
85 group.finish();
86}
87
88fn bench_notify_waiters(c: &mut Criterion) {
89 let mut group = c.benchmark_group("notify_waiters");
90 notify_waiters::<10>(&mut group);
91 notify_waiters::<50>(&mut group);
92 notify_waiters::<100>(&mut group);
93 notify_waiters::<200>(&mut group);
94 notify_waiters::<500>(&mut group);
95 group.finish();
96}
97
98criterion_group!(
99 notify_waiters_simple,
100 bench_notify_one,
101 bench_notify_waiters
102);
103
104criterion_main!(notify_waiters_simple);
105