use std::sync::atomic::{AtomicU32, Ordering};
pub struct AttemptBudget {
remaining: AtomicU32,
unlimited: bool,
}
impl AttemptBudget {
pub fn new(max: u32) -> Self {
Self {
remaining: AtomicU32::new(max),
unlimited: max == 0,
}
}
pub fn try_consume(&self) -> bool {
if self.unlimited {
return true;
}
self.remaining
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| {
if v > 0 { Some(v - 1) } else { None }
})
.is_ok()
}
pub fn remaining(&self) -> u32 {
if self.unlimited {
u32::MAX
} else {
self.remaining.load(Ordering::SeqCst)
}
}
}
impl std::fmt::Debug for AttemptBudget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AttemptBudget")
.field("remaining", &self.remaining())
.field("unlimited", &self.unlimited)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn consumes_until_exhausted() {
let b = AttemptBudget::new(3);
assert!(b.try_consume());
assert!(b.try_consume());
assert!(b.try_consume());
assert!(!b.try_consume()); assert_eq!(b.remaining(), 0);
}
#[test]
fn zero_means_unlimited() {
let b = AttemptBudget::new(0);
for _ in 0..1000 {
assert!(b.try_consume());
}
}
#[test]
fn remaining_tracks_consumption() {
let b = AttemptBudget::new(2);
assert_eq!(b.remaining(), 2);
b.try_consume();
assert_eq!(b.remaining(), 1);
b.try_consume();
assert_eq!(b.remaining(), 0);
}
}