Buffer

#[derive(Debug)]
pub struct Buffer<Req, F> {
    tx: PollSender<Message<Req, F>>,
    handle: Handle,
}

    pub fn pair<S>(service: S, bound: usize) -> (Self, Worker<S, Req>)
    where
        S: Service<Req, Future = F> + Send + 'static,
        F: Send,
        S::Error: Into<crate::BoxError> + Send + Sync,
        Req: Send + 'static,
    {
        let (tx, rx) = mpsc::channel(bound);
        let (handle, worker) = Worker::new(service, rx);
        let buffer = Self {
            tx: PollSender::new(tx),
            handle,
        };
        (buffer, worker)
    }

pub(crate) struct Message<Request, Fut> {
    pub(crate) request: Request,
    pub(crate) tx: Tx<Fut>,
    pub(crate) span: tracing::Span,
}

/// Response sender
pub(crate) type Tx<Fut> = oneshot::Sender<Result<Fut, ServiceError>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        // First, check if the worker is still alive.
        if self.tx.is_closed() {
            // If the inner service has errored, then we error here.
            return Poll::Ready(Err(self.get_worker_error()));
        }

        // Poll the sender to acquire a permit.
        self.tx
            .poll_reserve(cx)
            .map_err(|_| self.get_worker_error())
    }

    fn call(&mut self, request: Req) -> Self::Future {
        tracing::trace!("sending request to buffer worker");

        // get the current Span so that we can explicitly propagate it to the worker
        // if we didn't do this, events on the worker related to this span wouldn't be counted
        // towards that span since the worker would have no way of entering it.
        let span = tracing::Span::current();

        // If we've made it here, then a channel permit has already been
        // acquired, so we can freely allocate a oneshot.
        let (tx, rx) = oneshot::channel();

        match self.tx.send_item(Message { request, span, tx }) {
            Ok(_) => ResponseFuture::new(rx),
            // If the channel is closed, propagate the error from the worker.
            Err(_) => {
                tracing::trace!("buffer channel closed");
                ResponseFuture::failed(self.get_worker_error())
            }
        }
    }

        loop {
            match ready!(self.poll_next_msg(cx)) {
                Some((msg, first)) => {
                    let _guard = msg.span.enter();
                    if let Some(ref failed) = self.failed {
                        tracing::trace!("notifying caller about worker failure");
                        let _ = msg.tx.send(Err(failed.clone()));
                        continue;
                    }

                    // ...
                    match self.service.poll_ready(cx) {
                        Poll::Ready(Ok(())) => {
                            tracing::debug!(service.ready = true, message = "processing request");
                            let response = self.service.call(msg.request);

                            // Send the response future back to the sender.
                            //
                            // An error means the request had been canceled in-between
                            // our calls, the response future will just be dropped.
                            tracing::trace!("returning response future");
                            let _ = msg.tx.send(Ok(response));
                        }
                        Poll::Pending => {
                            tracing::trace!(service.ready = false, message = "delay");
                            // Put out current message back in its slot.
                            drop(_guard);
                            self.current_message = Some(msg);
                            return Poll::Pending;
                        }
                        Poll::Ready(Err(e)) => {
                            // ...
                            self.failed(error);
                            // ...
                        }
                    }
                }
                None => {
                    // No more more requests _ever_.
                    self.finish = true;
                    return Poll::Ready(());
                }
            }
        }

tower = { version = "=0.5.3", features = ["buffer", "util"] }
tokio = { version = "=1.54.1", features = ["full"] }

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tower::buffer::Buffer;
use tower::{Service, ServiceExt};

// An inner service that is not Clone. It reports which task runs each step.
struct Echo;

impl Service<u32> for Echo {
  type Response = u32;
  type Error = std::convert::Infallible;
  type Future = Pin<Box<dyn Future<Output = Result<u32, Self::Error>> + Send>>;

  fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    println!("inner poll_ready  in task {:?}", tokio::task::try_id());
    Poll::Ready(Ok(()))
  }

  fn call(&mut self, req: u32) -> Self::Future {
    println!("inner call        in task {:?}", tokio::task::try_id());
    Box::pin(async move {
      println!("inner future body in task {:?}", tokio::task::try_id());
      Ok(req * 2)
    })
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let svc = Buffer::new(Echo, 8);
  let caller = tokio::spawn(async move {
    println!("caller            is task {:?}", tokio::task::try_id());
    let mut svc = svc;
    let rsp = svc.ready().await.unwrap().call(21).await.unwrap();
    println!("response = {rsp}");
  });
  caller.await.unwrap();
}

caller            is task Some(Id(2))
inner poll_ready  in task Some(Id(1))
inner call        in task Some(Id(1))
inner future body in task Some(Id(2))
response = 42

        loop {
            match this.state.as_mut().project() {
                ResponseStateProj::Failed { error } => {
                    return Poll::Ready(Err(error.take().expect("polled after error")));
                }
                ResponseStateProj::Rx { rx } => match ready!(rx.poll(cx)) {
                    Ok(Ok(fut)) => this.state.set(ResponseState::Poll { fut }),
                    Ok(Err(e)) => return Poll::Ready(Err(e.into())),
                    Err(_) => return Poll::Ready(Err(Closed::new().into())),
                },
                ResponseStateProj::Poll { fut } => return fut.poll(cx).map_err(Into::into),
            }
        }

    fn failed(&mut self, error: crate::BoxError) {
        // ...
        // We do this by *first* exposing the error, *then* closing the channel used to
        // send more requests (so the client will see the error when the send fails), and *then*
        // sending the error to all outstanding requests.
        let error = ServiceError::new(error);

        let mut inner = self.handle.inner.lock().unwrap();

        if inner.is_some() {
            // Future::poll was called after we've already errored out!
            return;
        }

        *inner = Some(error.clone());
        drop(inner);

        self.rx.close();

        // By closing the mpsc::Receiver, we know that poll_next_msg will soon return Ready(None),
        // which will trigger the `self.finish == true` phase. We just need to make sure that any
        // requests that we receive before we've exhausted the receiver receive the error:
        self.failed = Some(error);
    }

tower = { version = "=0.5.3", features = ["buffer", "util"] }
tokio = { version = "=1.54.1", features = ["full"] }

use std::future::{ready, Ready};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use tower::buffer::{error::ServiceError, Buffer};
use tower::{BoxError, Service, ServiceExt};

// Pending until `fail` is set, then poll_ready returns an error.
#[derive(Default)]
struct Switch {
  fail: bool,
  waker: Option<Waker>,
}

struct Flaky(Arc<Mutex<Switch>>);

impl Service<u32> for Flaky {
  type Response = u32;
  type Error = BoxError;
  type Future = Ready<Result<u32, BoxError>>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    let mut s = self.0.lock().unwrap();
    if s.fail {
      return Poll::Ready(Err("backend down".into()));
    }
    s.waker = Some(cx.waker().clone());
    Poll::Pending
  }

  fn call(&mut self, req: u32) -> Self::Future {
    ready(Ok(req))
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let switch = Arc::new(Mutex::new(Switch::default()));
  let svc = Buffer::new(Flaky(switch.clone()), 1);

  // Fill the worker's hand and the single channel slot.
  let mut a = svc.clone();
  let fa = a.ready().await.unwrap().call(1);
  tokio::task::yield_now().await;
  let mut b = svc.clone();
  let fb = b.ready().await.unwrap().call(2);

  // c waits in poll_ready for a slot.
  let mut c = svc.clone();
  let waiter = tokio::spawn(async move { c.ready().await.map(|_| ()) });
  tokio::task::yield_now().await;
  println!("c finished before failure: {}", waiter.is_finished());

  let w = {
    let mut s = switch.lock().unwrap();
    s.fail = true;
    s.waker.take()
  };
  w.unwrap().wake();

  println!("a: {}", fa.await.unwrap_err());
  println!("b: {}", fb.await.unwrap_err());
  let e = waiter.await.unwrap().unwrap_err();
  println!("c: {e}");
  // The boxed error is a ServiceError; its source is the inner error.
  let source = std::error::Error::source(&*e).map(|s| s.to_string());
  println!("c is ServiceError: {}, source: {source:?}", e.is::<ServiceError>());
}

c finished before failure: false
a: buffered service failed: backend down
b: buffered service failed: backend down
c: buffered service failed: backend down
c is ServiceError: true, source: Some("backend down")

use tower::buffer::Buffer;
use tower::{service_fn, BoxError, Service, ServiceExt};

fn main() {
  // The worker is spawned on rt1, a runtime with its own thread.
  let rt1 = tokio::runtime::Builder::new_multi_thread()
    .worker_threads(1)
    .build()
    .unwrap();
  let mut svc = rt1.block_on(async {
    Buffer::new(service_fn(|x: u32| async move { Ok::<_, BoxError>(x) }), 4)
  });

  // The caller lives on rt2.
  let rt2 = tokio::runtime::Builder::new_current_thread().build().unwrap();
  rt2.block_on(async {
    let r = svc.ready().await.unwrap().call(1).await;
    println!("rt1 alive  : {:?}", r.map_err(|e| e.to_string()));
  });

  // Shutting rt1 down drops the worker task.
  drop(rt1);
  rt2.block_on(async {
    match svc.ready().await {
      Ok(_) => println!("rt1 dropped: ready"),
      Err(e) => println!("rt1 dropped: {e}"),
    }
  });
}

rt1 alive  : Ok(1)
rt1 dropped: buffer's worker closed unexpectedly

        if let Some(msg) = self.current_message.take() {
            // If the oneshot sender is closed, then the receiver is dropped,
            // and nobody cares about the response. If this is the case, we
            // should continue to the next request.
            if !msg.tx.is_closed() {
                tracing::trace!("resuming buffered request");
                return Poll::Ready(Some((msg, false)));
            }

            tracing::trace!("dropping cancelled buffered request");
        }

        // Get the next request
        while let Some(msg) = ready!(Pin::new(&mut self.rx).poll_recv(cx)) {
            if !msg.tx.is_closed() {
                tracing::trace!("processing new request");
                return Poll::Ready(Some((msg, true)));
            }
            // Otherwise, request is canceled, so pop the next one.
            tracing::trace!("dropping cancelled request");
        }

use std::future::{ready, Ready};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use tower::buffer::Buffer;
use tower::{BoxError, Service, ServiceExt};

// Not ready until someone opens the gate.
#[derive(Default)]
struct Gate {
  open: bool,
  waker: Option<Waker>,
}

struct Gated(Arc<Mutex<Gate>>);

impl Service<u32> for Gated {
  type Response = u32;
  type Error = BoxError;
  type Future = Ready<Result<u32, BoxError>>;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), BoxError>> {
    let mut g = self.0.lock().unwrap();
    if g.open {
      return Poll::Ready(Ok(()));
    }
    g.waker = Some(cx.waker().clone());
    Poll::Pending
  }

  fn call(&mut self, req: u32) -> Self::Future {
    println!("inner call({req})");
    ready(Ok(req))
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let gate = Arc::new(Mutex::new(Gate::default()));
  let svc = Buffer::new(Gated(gate.clone()), 8);

  let mut futs = Vec::new();
  for i in 1..=3 {
    let mut h = svc.clone();
    futs.push(h.ready().await.unwrap().call(i));
    tokio::task::yield_now().await; // let the worker pick it up
  }
  // Callers of 1 and 2 give up before the inner service is ready.
  let f3 = futs.pop().unwrap();
  drop(futs);

  let waker = {
    let mut g = gate.lock().unwrap();
    g.open = true;
    g.waker.take()
  };
  waker.unwrap().wake();
  println!("f3 = {}", f3.await.unwrap());
}

inner call(3)
f3 = 3

    /// When [`Buffer`]'s implementation of [`poll_ready`] returns [`Poll::Ready`], it reserves a
    /// slot in the channel for the forthcoming [`call`]. However, if this call doesn't arrive,
    /// this reserved slot may be held up for a long time. As a result, it's advisable to set
    /// `bound` to be at least the maximum number of concurrent requests the [`Buffer`] will see.
    /// If you do not, all the slots in the buffer may be held up by futures that have just called
    /// [`poll_ready`] but will not issue a [`call`], which prevents other senders from issuing new
    /// requests.

use std::time::Duration;
use tokio::time::timeout;
use tower::buffer::Buffer;
use tower::{service_fn, BoxError, Service, ServiceExt};

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  let svc = Buffer::new(service_fn(|x: u32| async move { Ok::<_, BoxError>(x) }), 2);

  // Two handles reserve a slot each, then go off and do something else.
  let mut h1 = svc.clone();
  let mut h2 = svc.clone();
  h1.ready().await.unwrap();
  h2.ready().await.unwrap();
  println!("h1, h2 reserved, not calling");

  // A third caller cannot even get a slot, though the inner service is idle.
  let mut c = svc.clone();
  let r = timeout(Duration::from_secs(5), c.ready()).await;
  println!("c.ready within 5s: {}", if r.is_ok() { "ok" } else { "timed out" });

  // Dropping a handle that holds a reservation gives the slot back.
  drop(h1);
  let r = timeout(Duration::from_secs(5), c.ready()).await;
  println!("after drop(h1): {}", if r.is_ok() { "ok" } else { "timed out" });
  println!("c.call(7) = {}", c.call(7).await.unwrap());
}

h1, h2 reserved, not calling
c.ready within 5s: timed out
after drop(h1): ok
c.call(7) = 7

accepted 2 requests, request 2 blocks

accepted 1 requests, request 1 blocks