ConcurrencyLimit

#[derive(Debug)]
pub struct ConcurrencyLimit<T> {
    inner: T,
    semaphore: PollSemaphore,
    /// The currently acquired semaphore permit, if there is sufficient
    /// concurrency to send a new request.
    ///
    /// The permit is acquired in `poll_ready`, and taken in `call` when sending
    /// a new request.
    permit: Option<OwnedSemaphorePermit>,
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        // If we haven't already acquired a permit from the semaphore, try to
        // acquire one first.
        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",
            );
        }

        // Once we've acquired a permit (or if we already had one), poll the
        // inner service.
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // Take the permit
        let permit = self
            .permit
            .take()
            .expect("max requests in-flight; poll_ready must be called first");

        // Call the inner service
        let future = self.inner.call(request);

        ResponseFuture::new(future, permit)
    }

    pub struct ResponseFuture<T> {
        #[pin]
        inner: T,
        // Keep this around so that it is dropped when the future completes
        _permit: OwnedSemaphorePermit,
    }

impl<T: Clone> Clone for ConcurrencyLimit<T> {
    fn clone(&self) -> Self {
        // Since we hold an `OwnedSemaphorePermit`, we can't derive `Clone`.
        // Instead, when cloning the service, create a new service with the
        // same semaphore, but with the permit in the un-acquired state.
        Self {
            inner: self.inner.clone(),
            semaphore: self.semaphore.clone(),
            permit: None,
        }
    }
}

tower = { version = "=0.5.3", features = ["limit", "util"] }
tokio = { version = "=1.54.1", features = ["rt", "macros", "sync"] }

use std::{
  convert::Infallible,
  future::poll_fn,
  panic::{catch_unwind, AssertUnwindSafe},
  sync::Arc,
  task::Poll,
};
use tokio::sync::Semaphore;
use tower::{limit::GlobalConcurrencyLimitLayer, service_fn, Layer, Service};

async fn poll_ready_once<S: Service<u32>>(svc: &mut S) -> &'static str {
  poll_fn(|cx| {
    Poll::Ready(match svc.poll_ready(cx) {
      Poll::Ready(Ok(())) => "Ready",
      Poll::Ready(Err(_)) => "Err",
      Poll::Pending => "Pending",
    })
  })
  .await
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let sem = Arc::new(Semaphore::new(2));
  let inner = service_fn(|x: u32| async move { Ok::<_, Infallible>(x) });
  let mut a = GlobalConcurrencyLimitLayer::with_semaphore(sem.clone()).layer(inner);

  println!("1 start:                available = {}", sem.available_permits());
  println!("  a.poll_ready -> {}", poll_ready_once(&mut a).await);
  println!("2 a ready, no call:     available = {}", sem.available_permits());
  println!("  a.poll_ready -> {}", poll_ready_once(&mut a).await);
  println!("3 a polled again:       available = {}", sem.available_permits());
  let fut = a.call(1);
  println!("4 a called, fut alive:  available = {}", sem.available_permits());

  let mut b = a.clone();
  println!("  b.poll_ready -> {}", poll_ready_once(&mut b).await);
  println!("5 clone b ready:        available = {}", sem.available_permits());
  let mut c = a.clone();
  println!("  c.poll_ready -> {}", poll_ready_once(&mut c).await);

  drop(fut);
  println!("6 fut dropped:          available = {}", sem.available_permits());
  println!("  c.poll_ready -> {}", poll_ready_once(&mut c).await);

  let mut d = a.clone();
  std::panic::set_hook(Box::new(|_| {})); // keep the default hook quiet
  let r = catch_unwind(AssertUnwindSafe(|| {
    let _ = d.call(2);
  }));
  let p = r.unwrap_err();
  let msg = p
    .downcast_ref::<String>()
    .cloned()
    .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string()));
  println!("7 d.call without poll_ready panics: {msg:?}");
}

#[derive(Debug, Clone)]
pub struct GlobalConcurrencyLimitLayer {
    semaphore: Arc<Semaphore>,
}

1 start:                available = 2
  a.poll_ready -> Ready
2 a ready, no call:     available = 1
  a.poll_ready -> Ready
3 a polled again:       available = 1
4 a called, fut alive:  available = 1
  b.poll_ready -> Ready
5 clone b ready:        available = 0
  c.poll_ready -> Pending
6 fut dropped:          available = 0
  c.poll_ready -> Ready
7 d.call without poll_ready panics: Some("max requests in-flight; poll_ready must be called first")