blob: 1966d57850a0bd5775d5ff85032f0dd43ddc88ee [file] [log] [blame]
use std::sync::Arc;
use super::ConcurrencyLimit;
use tokio::sync::Semaphore;
use tower_layer::Layer;
/// Enforces a limit on the concurrent number of requests the underlying
/// service can handle.
#[derive(Debug, Clone)]
pub struct ConcurrencyLimitLayer {
max: usize,
}
impl ConcurrencyLimitLayer {
/// Create a new concurrency limit layer.
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)
}
}
/// Enforces a limit on the concurrent number of requests the underlying
/// service can handle.
///
/// Unlike [`ConcurrencyLimitLayer`], which enforces a per-service concurrency
/// limit, this layer accepts a owned semaphore (`Arc<Semaphore>`) which can be
/// shared across multiple services.
///
/// Cloning this layer will not create a new semaphore.
#[derive(Debug, Clone)]
pub struct GlobalConcurrencyLimitLayer {
semaphore: Arc<Semaphore>,
}
impl GlobalConcurrencyLimitLayer {
/// Create a new `GlobalConcurrencyLimitLayer`.
pub fn new(max: usize) -> Self {
Self::with_semaphore(Arc::new(Semaphore::new(max)))
}
/// Create a new `GlobalConcurrencyLimitLayer` from a `Arc<Semaphore>`
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())
}
}