#![cfg(loom)]
use std::time::Duration;
use loom::sync::atomic::{AtomicU64, Ordering};
use loom::sync::Arc;
use loom::thread;
use throttle_core::{ManualClock, Throttle};
#[test]
fn two_threads_never_overgrant() {
loom::model(|| {
let clock = Arc::new(ManualClock::new());
let throttle = Arc::new(Throttle::with_clock(2, 0, clock));
let granted = Arc::new(AtomicU64::new(0));
let t1 = {
let t = throttle.clone();
let g = granted.clone();
thread::spawn(move || {
if t.try_acquire(1) {
g.fetch_add(1, Ordering::SeqCst);
}
if t.try_acquire(1) {
g.fetch_add(1, Ordering::SeqCst);
}
})
};
let t2 = {
let t = throttle.clone();
let g = granted.clone();
thread::spawn(move || {
if t.try_acquire(1) {
g.fetch_add(1, Ordering::SeqCst);
}
if t.try_acquire(1) {
g.fetch_add(1, Ordering::SeqCst);
}
})
};
t1.join().unwrap();
t2.join().unwrap();
let g = granted.load(Ordering::SeqCst);
assert_eq!(g, 2, "expected exactly 2 grants under no-refill, got {g}");
assert_eq!(throttle.available(), 0);
});
}
#[test]
fn refill_grants_resume_after_clock_advance() {
loom::model(|| {
let clock = Arc::new(ManualClock::new());
let throttle = Arc::new(Throttle::with_clock(1, 1_000, clock.clone()));
assert!(throttle.try_acquire(1));
assert!(!throttle.try_acquire(1));
let advancer = {
let c = clock.clone();
thread::spawn(move || {
c.advance(Duration::from_millis(10));
})
};
advancer.join().unwrap();
assert!(
throttle.try_acquire(1),
"refill did not become visible after clock advance"
);
});
}
#[test]
fn try_acquire_zero_is_always_true() {
loom::model(|| {
let clock = Arc::new(ManualClock::new());
let throttle = Arc::new(Throttle::with_clock(1, 0, clock));
let racer = {
let t = throttle.clone();
thread::spawn(move || {
let _ = t.try_acquire(1);
})
};
assert!(throttle.try_acquire(0));
racer.join().unwrap();
});
}
#[test]
fn capacity_one_serializes_grants() {
loom::model(|| {
let clock = Arc::new(ManualClock::new());
let throttle = Arc::new(Throttle::with_clock(1, 0, clock));
let winners = Arc::new(AtomicU64::new(0));
let t1 = {
let t = throttle.clone();
let w = winners.clone();
thread::spawn(move || {
if t.try_acquire(1) {
w.fetch_add(1, Ordering::SeqCst);
}
})
};
let t2 = {
let t = throttle.clone();
let w = winners.clone();
thread::spawn(move || {
if t.try_acquire(1) {
w.fetch_add(1, Ordering::SeqCst);
}
})
};
t1.join().unwrap();
t2.join().unwrap();
let w = winners.load(Ordering::SeqCst);
assert_eq!(
w, 1,
"capacity-1 bucket must hand out exactly 1 grant, got {w}"
);
assert_eq!(throttle.available(), 0);
});
}