use cano::prelude::*;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use std::sync::Arc;
use std::time::Duration;
fn big_bucket() -> RateLimiterPolicy {
RateLimiterPolicy::per_second(u32::MAX / 2).with_max_tokens(u32::MAX / 2)
}
fn bench_try_acquire_success(c: &mut Criterion) {
let mut group = c.benchmark_group("rate_limiter_try_acquire_success");
let limiter = Arc::new(RateLimiter::new(big_bucket()));
group.bench_function("single_thread", |b| {
b.iter(|| {
let permit = limiter.try_acquire();
std::hint::black_box(&permit);
});
});
group.finish();
}
fn bench_try_acquire_rejected(c: &mut Criterion) {
let mut group = c.benchmark_group("rate_limiter_try_acquire_rejected");
let limiter = Arc::new(RateLimiter::new(
RateLimiterPolicy::new(1, Duration::from_secs(3600)).with_max_tokens(1),
));
let _ = limiter.try_acquire().expect("drain the only token");
group.bench_function("rejected", |b| {
b.iter(|| {
let permit = limiter.try_acquire();
std::hint::black_box(&permit);
});
});
group.finish();
}
fn bench_shared_concurrent(c: &mut Criterion) {
let mut group = c.benchmark_group("rate_limiter_shared_concurrent");
let rt = tokio::runtime::Runtime::new().unwrap();
for &n in &[16usize, 64, 256] {
let limiter = Arc::new(RateLimiter::new(big_bucket()));
group.bench_with_input(BenchmarkId::new("tasks", n), &n, |b, &n| {
b.to_async(&rt).iter(|| {
let limiter = Arc::clone(&limiter);
async move {
let mut handles = Vec::with_capacity(n);
for _ in 0..n {
let limiter = Arc::clone(&limiter);
handles.push(tokio::spawn(async move {
let permit = limiter.try_acquire();
std::hint::black_box(&permit);
}));
}
for h in handles {
let _ = h.await;
}
}
});
});
}
group.finish();
}
criterion_group!(
benches,
bench_try_acquire_success,
bench_try_acquire_rejected,
bench_shared_concurrent,
);
criterion_main!(benches);