LoadShed

#[derive(Debug)]
pub struct LoadShed<S> {
    inner: S,
    is_ready: bool,
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        // We check for readiness here, so that we can know in `call` if
        // the inner service is overloaded or not.
        self.is_ready = match self.inner.poll_ready(cx) {
            Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
            r => r.is_ready(),
        };

        // But we always report Ready, so that layers above don't wait until
        // the inner service is ready (the entire point of this layer!)
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: Req) -> Self::Future {
        if self.is_ready {
            // readiness only counts once, you need to check again!
            self.is_ready = false;
            ResponseFuture::called(self.inner.call(req))
        } else {
            ResponseFuture::overloaded()
        }
    }

    enum ResponseState<F> {
        Called {
            #[pin]
            fut: F
        },
        Overloaded,
    }

tower = { version = "=0.5.3", features = ["load-shed", "limit", "util"] }
tokio = { version = "=1.54.1", features = ["rt", "macros", "time", "test-util"] }

use std::time::Duration;
use tokio::time::{sleep, Instant};
use tower::{BoxError, Service, ServiceBuilder, ServiceExt};

async fn work(id: u32) -> Result<u32, BoxError> {
  sleep(Duration::from_millis(100)).await;
  Ok(id)
}

async fn run<S>(name: &str, svc: S)
where
  S: Service<u32, Response = u32, Error = BoxError> + Clone + Send + 'static,
  S::Future: Send,
{
  let start = Instant::now();
  let mut tasks = Vec::new();
  for id in 1..=4 {
    let mut svc = svc.clone();
    tasks.push(tokio::spawn(async move {
      let r = svc.ready().await?.call(id).await;
      Ok::<_, BoxError>((id, r, start.elapsed().as_millis()))
    }));
  }
  println!("{name}");
  for t in tasks {
    let (id, r, ms) = t.await.unwrap().unwrap();
    match r {
      Ok(v) => println!("  req {id}: ok {v} at {ms}ms"),
      Err(e) => println!("  req {id}: err \"{e}\" at {ms}ms"),
    }
  }
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  let a = ServiceBuilder::new()
    .load_shed()
    .concurrency_limit(2)
    .service_fn(work);
  run("A: LoadShed > ConcurrencyLimit(2)", a).await;

  let b = ServiceBuilder::new()
    .concurrency_limit(2)
    .load_shed()
    .service_fn(work);
  run("B: ConcurrencyLimit(2) > LoadShed", b).await;
}

A: LoadShed > ConcurrencyLimit(2)
  req 1: ok 1 at 100ms
  req 2: ok 2 at 100ms
  req 3: err "service overloaded" at 0ms
  req 4: err "service overloaded" at 0ms
B: ConcurrencyLimit(2) > LoadShed
  req 1: ok 1 at 100ms
  req 2: ok 2 at 100ms
  req 3: ok 3 at 200ms
  req 4: ok 4 at 200ms

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

use std::{convert::Infallible, future::poll_fn, sync::Arc, task::Poll};
use tokio::sync::Semaphore;
use tower::{
  limit::GlobalConcurrencyLimitLayer, load_shed::LoadShedLayer, service_fn, Service,
  ServiceBuilder,
};

async fn ready_now<S: Service<u32>>(svc: &mut S) -> bool {
  poll_fn(|cx| Poll::Ready(svc.poll_ready(cx).is_ready())).await
}

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

  let (mut a, mut b, mut c) = (svc.clone(), svc.clone(), svc.clone());

  ready_now(&mut a).await;
  let fut_a = a.call(1); // a holds the only permit
  ready_now(&mut b).await; // LoadShed says Ready, inner was Pending
  println!("b.call -> {:?}", b.call(2).await.map_err(|e| e.to_string()));

  drop(fut_a); // permit released
  println!("after fut_a dropped: available = {}", sem.available_permits());
  ready_now(&mut c).await;
  println!("c.call -> {:?}", c.call(3).await.map_err(|e| e.to_string()));
  ready_now(&mut b).await;
  println!("b.call -> {:?}", b.call(4).await.map_err(|e| e.to_string()));
}

b.call -> Err("service overloaded")
after fut_a dropped: available = 0
c.call -> Err("service overloaded")
b.call -> Ok(4)