use std::sync::Arc;
use super::ConcurrencyLimit;
use tokio::sync::Semaphore;
use tower_layer::Layer;
#[derive(Debug, Clone)]
pub struct ConcurrencyLimitLayer {
max: usize,
}
impl ConcurrencyLimitLayer {
pub fn new(max: usize) -> Self {
ConcurrencyLimitLayer { max }
}
}
impl<S> Layer<S> for ConcurrencyLimitLayer {
type Service = ConcurrencyLimit<S>;
fn layer(&self, service: S) -> Self::Service {
ConcurrencyLimit::new(service, self.max)
}
}
#[derive(Debug, Clone)]
pub struct GlobalConcurrencyLimitLayer {
semaphore: Arc<Semaphore>,
}
impl GlobalConcurrencyLimitLayer {
pub fn new(max: usize) -> Self {
Self::with_semaphore(Arc::new(Semaphore::new(max)))
}
pub fn with_semaphore(semaphore: Arc<Semaphore>) -> Self {
GlobalConcurrencyLimitLayer { semaphore }
}
}
impl<S> Layer<S> for GlobalConcurrencyLimitLayer {
type Service = ConcurrencyLimit<S>;
fn layer(&self, service: S) -> Self::Service {
ConcurrencyLimit::with_semaphore(service, self.semaphore.clone())
}
}