use super::future::ResponseFuture;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio_util::sync::PollSemaphore;
use tower_service::Service;
use futures_core::ready;
use std::{
sync::Arc,
task::{Context, Poll},
};
#[derive(Debug)]
pub struct ConcurrencyLimit<T> {
inner: T,
semaphore: PollSemaphore,
permit: Option<OwnedSemaphorePermit>,
}
impl<T> ConcurrencyLimit<T> {
pub fn new(inner: T, max: usize) -> Self {
Self::with_semaphore(inner, Arc::new(Semaphore::new(max)))
}
pub fn with_semaphore(inner: T, semaphore: Arc<Semaphore>) -> Self {
ConcurrencyLimit {
inner,
semaphore: PollSemaphore::new(semaphore),
permit: None,
}
}
pub fn get_ref(&self) -> &T {
&self.inner
}
pub fn get_mut(&mut self) -> &mut T {
&mut self.inner
}
pub fn into_inner(self) -> T {
self.inner
}
}
impl<S, Request> Service<Request> for ConcurrencyLimit<S>
where
S: Service<Request>,
{
type Response = S::Response;
type Error = S::Error;
type Future = ResponseFuture<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if self.permit.is_none() {
self.permit = ready!(self.semaphore.poll_acquire(cx));
debug_assert!(
self.permit.is_some(),
"ConcurrencyLimit semaphore is never closed, so `poll_acquire` \
should never fail",
);
}
self.inner.poll_ready(cx)
}
fn call(&mut self, request: Request) -> Self::Future {
let permit = self
.permit
.take()
.expect("max requests in-flight; poll_ready must be called first");
let future = self.inner.call(request);
ResponseFuture::new(future, permit)
}
}
impl<T: Clone> Clone for ConcurrencyLimit<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
semaphore: self.semaphore.clone(),
permit: None,
}
}
}
#[cfg(feature = "load")]
#[cfg_attr(docsrs, doc(cfg(feature = "load")))]
impl<S> crate::load::Load for ConcurrencyLimit<S>
where
S: crate::load::Load,
{
type Metric = S::Metric;
fn load(&self) -> Self::Metric {
self.inner.load()
}
}