zhenyi
47 lines
1.1 KiB
Rust
47 lines
1.1 KiB
Rust
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
|
|
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct HandlerLimitError;
|
|
|
|
#[derive(Clone)]
|
|
pub struct HandlerLimiter {
|
|
sem: Arc<Semaphore>,
|
|
max_inflight: usize,
|
|
rejected: Arc<AtomicU64>,
|
|
}
|
|
|
|
impl HandlerLimiter {
|
|
pub fn new(max_inflight: usize) -> Self {
|
|
Self {
|
|
sem: Arc::new(Semaphore::new(max_inflight)),
|
|
max_inflight,
|
|
rejected: Arc::new(AtomicU64::new(0)),
|
|
}
|
|
}
|
|
|
|
pub fn try_acquire(&self) -> Result<OwnedSemaphorePermit, HandlerLimitError> {
|
|
match self.sem.clone().try_acquire_owned() {
|
|
Ok(permit) => Ok(permit),
|
|
Err(_) => {
|
|
self.rejected.fetch_add(1, Ordering::Relaxed);
|
|
Err(HandlerLimitError)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn inflight(&self) -> usize {
|
|
self.max_inflight - self.sem.available_permits()
|
|
}
|
|
|
|
pub fn available(&self) -> usize {
|
|
self.sem.available_permits()
|
|
}
|
|
|
|
pub fn rejected_total(&self) -> u64 {
|
|
self.rejected.load(Ordering::Relaxed)
|
|
}
|
|
}
|