#[cfg(feature = "sync_reader")]
pub use sync::ReadLimiter;
#[cfg(feature = "sync_reader")]
mod sync {
use crate::{Error, Result};
use alloc::string::String;
use core::sync::atomic::{AtomicUsize, Ordering};
pub struct ReadLimiter {
limit: AtomicUsize,
error_on_limit_exceeded: bool,
}
impl ReadLimiter {
pub fn new(limit: Option<usize>) -> Self {
match limit {
Some(value) => Self {
limit: AtomicUsize::new(value),
error_on_limit_exceeded: true,
},
None => Self {
limit: AtomicUsize::new(usize::MAX),
error_on_limit_exceeded: false,
},
}
}
#[inline]
pub fn can_read(&self, amount: usize) -> Result<()> {
let current = self.limit.load(Ordering::Relaxed);
if amount > current && self.error_on_limit_exceeded {
return Err(Error::failed(String::from("read limit exceeded")));
} else {
self.limit
.store(current.wrapping_sub(amount), Ordering::Relaxed);
}
Ok(())
}
}
}
#[cfg(not(feature = "sync_reader"))]
pub use unsync::ReadLimiter;
#[cfg(not(feature = "sync_reader"))]
mod unsync {
use crate::{Error, Result};
use alloc::string::String;
use core::cell::Cell;
pub struct ReadLimiter {
limit: Cell<usize>,
error_on_limit_exceeded: bool,
}
impl ReadLimiter {
pub fn new(limit: Option<usize>) -> Self {
match limit {
Some(value) => Self {
limit: Cell::new(value),
error_on_limit_exceeded: true,
},
None => Self {
limit: Cell::new(usize::MAX),
error_on_limit_exceeded: false,
},
}
}
#[inline]
pub fn can_read(&self, amount: usize) -> Result<()> {
let current = self.limit.get();
if amount > current && self.error_on_limit_exceeded {
Err(Error::failed(String::from("read limit exceeded")))
} else {
self.limit.set(current.wrapping_sub(amount));
Ok(())
}
}
}
}