#![cfg(all(feature = "in-memory", not(feature = "metrics")))]
#![allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic
)]
use std::alloc::{GlobalAlloc, Layout, System};
use std::future::Future;
use std::pin::pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll, Waker};
use throttle_kit::{InMemoryBackend, KeyedRateLimiter, Quota, RateLimiter};
fn block_once<F: Future>(fut: F) -> F::Output {
let mut fut = pin!(fut);
let waker = Waker::noop();
let mut cx = Context::from_waker(waker);
match fut.as_mut().poll(&mut cx) {
Poll::Ready(output) => output,
Poll::Pending => panic!("check future must complete on its first poll"),
}
}
static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0);
struct Counting;
unsafe impl GlobalAlloc for Counting {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
unsafe { System.alloc_zeroed(layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
unsafe { System.realloc(ptr, layout, new_size) }
}
}
#[global_allocator]
static GLOBAL: Counting = Counting;
fn allocations() -> usize {
ALLOCATIONS.load(Ordering::Relaxed)
}
const WARM_KEY: &str = "zero-alloc-warm-key";
const ITERATIONS: usize = 100;
#[test]
fn hot_path_allocation_claims() {
let limiter = RateLimiter::new(Quota::per_second(1_000_000), InMemoryBackend::new());
assert!(block_once(limiter.check(WARM_KEY)).allowed);
let before = allocations();
for _ in 0..ITERATIONS {
assert!(block_once(limiter.check(WARM_KEY)).allowed);
}
assert_eq!(
allocations(),
before,
"warm-key check must not allocate (before: {before}, after: {})",
allocations()
);
let keyed = KeyedRateLimiter::new(Quota::per_second(1_000_000), InMemoryBackend::new());
assert!(block_once(keyed.check(WARM_KEY)).allowed);
let before = allocations();
for _ in 0..ITERATIONS {
assert!(block_once(keyed.check(WARM_KEY)).allowed);
}
assert_eq!(
allocations(),
before,
"keyed warm-key check must not allocate (before: {before}, after: {})",
allocations()
);
keyed.with_quota_for_key(WARM_KEY, Quota::per_second(1_000_000));
assert!(block_once(keyed.check(WARM_KEY)).allowed);
let before = allocations();
for _ in 0..ITERATIONS {
assert!(block_once(keyed.check(WARM_KEY)).allowed);
}
assert_eq!(
allocations(),
before,
"keyed warm-key check with an override must not allocate"
);
let limiter = RateLimiter::new(Quota::per_second(1_000_000), InMemoryBackend::new());
let before = allocations();
let key = format!("fresh-{before}");
assert!(block_once(limiter.check(&key)).allowed);
assert!(
allocations() > before,
"first-ever call for a key inserts its state and must allocate \
(counter sanity check)"
);
}