Repository: tower-rs/tower
Pull request #635 (2022-02-01 22:42:05)

rewrite `buffer` to use bounded MPSC
## Motivation

Currently, `tower::buffer` uses `tokio::sync::mpsc`'s _unbounded_
channel plus a `tokio::sync::Semaphore`, in order to re-implement a
bounded channel. This was necessary because when this code was updated
to the latest version of `tokio`, there was no way to reserve a
non-borrowed send permit from a `Sender`. Thus, it was necessary to use
the `Semaphore` for the future that is polled in `poll_ready` to acquire
send capacity, since a `Permit` from the `Sender` could not be stored in
the struct until it's consumed in `call`.

This code is Not Ideal. Reimplementing the bounded channel makes the
implementation more complicated, and means that there is a bunch of
extra stuff we have to do to e.g. propagate cancellations/service errors
to tasks waiting on `poll_ready`. The bounded MPSC would solve most of
this for us. It might also be a bit more efficient, since we would only
have a single reference-counted heap allocation (the `Sender`), rather
than two (the `Sender` _and_ the `Arc<Semaphore>`).

Now that `tokio::sync::mpsc::Sender` has a `reserve_owned` method, we no
longer need to use the semaphore, and can simplify the buffer
implementation a bit.

## Solution

This branch changes the buffer to use only a single bounded MPSC, rather
than an unbounded MPSC and a semaphore. The bounded MPSC internally
manages its semaphore, so we can now remove a lot of complexity in the
current implementation.

I had to change some of the integration tests slightly as part of this
change. This is because the buffer implementation using semaphore
permits is _very subtly_ different from one using a bounded channel. In
the `Semaphore`-based implementation, a semaphore permit is stored in
the `Message` struct sent over the channel. This is so that the capacity
is used as long as the message is in flight. However, when the worker
task is processing a message that's been recieved from the channel, the
permit is still not dropped. Essentially, the one message actively held
by the worker task _also_ occupies one "slot" of capacity, so the actual
channel capacity is one less than the value passed to the constructor,
_once the first request has been sent to the worker_. The bounded MPSC
changed this behavior so that capacity is only occupied while a request
is actually in the channel, which broke some tests that relied on the
old (and technically wrong) behavior.

## Notes

There is one sort of significant issue with this change, which is why
it's currently a draft. The issue is that this unfortunately requires
adding `Send` and `'static` bounds to the `T::Future` and `Request`
types. This is because they occur within the `Message` type that's sent
over the channel, and a MPSC's `OwnedPermit<T>` for a message of type
`T` is only `Send` or `'static` when `T` is. The `ReusableBoxFuture`
that the `reserve_owned` future is stored in requires the future be
`Send + 'static`, which it is not when `OwnedPermit` isn't.

I don't believe that it's actually _necessary_ for the `OwnedPermit<T>`
type to require `T: Send` in order to be `Send`, or `T: 'static` in
order to be valid for the `'static` lifetime. An `OwnedPermit` never
actually contains an instance of type `T`, it just represents the
_capacity_ to send that type to the channel. The channel itself will
actually contain the values of type `T`. Therefore, it's possible this
could be changed upstream in Tokio, although I haven't looked into it
yet.

In practice, these bounds shouldn't really be a problem, but adding them
is a breaking change. We could alternatively just wait for `tower` 0.5
to merge this PR.

hawkw: Now that `tokio-util` 0.7 is released, this code can be rewritten to just use the `PollSender` type. This is because the semantics of `PollSender` was changed to have a `poll_reserve` rather than `poll_send_done`, which is the required behavior for `Buffer`. The implementation of `PollSender` now looks more or less identical to the code I had written using `ReusableBoxFuture` and `reserve_owned`. Therefore, we can now just use the `tokio_util::sync::PollSender` type, and the buffer implementation is now much simpler.
LucioFranco: @hawkw is this still blocked on the breaking change?
hawkw: > @hawkw is this still blocked on the breaking change?

yes, this introduces [trait bounds](https://github.com/tower-rs/tower/pull/635/files#diff-35398bbf39b455b121231a971e231cfc669556eda1c98705f367aa774f3da86cR92-R93) that were previously not required. It's _possible_ an upstream change in `tokio-util` could make them unnecessary, but I'm not really sure how easy that would be yet (haven't looked too closely at it).
LucioFranco: I vote that we go ahead with a breaking change for this? 
LucioFranco: cc @davidpdrsn 
hawkw: > I vote that we go ahead with a breaking change for this?

Yeah, I'm hoping to prepare an 0.5 in the next couple days. we should ship another 0.4 release first though, which is why it's a draft.

I do also want to look at `tokio-util` to double check if it's possible to not introduce those bounds, because it is kind of a bummer that `buffer` now cannot be used with `!Sync` services...

Patch:
diff --git a/tower/src/buffer/message.rs b/tower/src/buffer/message.rs
index 2b2fcde..6d13aa1 100644
--- a/tower/src/buffer/message.rs
+++ b/tower/src/buffer/message.rs
@@ -1,5 +1,5 @@
 use super::error::ServiceError;
-use tokio::sync::{oneshot, OwnedSemaphorePermit};
+use tokio::sync::oneshot;
 
 /// Message sent over buffer
 #[derive(Debug)]
@@ -7,7 +7,6 @@ pub(crate) struct Message<Request, Fut> {
     pub(crate) request: Request,
     pub(crate) tx: Tx<Fut>,
     pub(crate) span: tracing::Span,
-    pub(super) _permit: OwnedSemaphorePermit,
 }
 
 /// Response sender
diff --git a/tower/src/buffer/service.rs b/tower/src/buffer/service.rs
index 9b690f0..c685984 100644
--- a/tower/src/buffer/service.rs
+++ b/tower/src/buffer/service.rs
@@ -4,11 +4,9 @@ use super::{
     worker::{Handle, Worker},
 };
 
-use futures_core::ready;
-use std::sync::Arc;
 use std::task::{Context, Poll};
-use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
-use tokio_util::sync::PollSemaphore;
+use tokio::sync::{mpsc, oneshot};
+use tokio_util::sync::PollSender;
 use tower_service::Service;
 
 /// Adds an mpsc buffer in front of an inner service.
@@ -19,23 +17,7 @@ pub struct Buffer<T, Request>
 where
     T: Service<Request>,
 {
-    // Note: this actually _is_ bounded, but rather than using Tokio's bounded
-    // channel, we use Tokio's semaphore separately to implement the bound.
-    tx: mpsc::UnboundedSender<Message<Request, T::Future>>,
-    // When the buffer's channel is full, we want to exert backpressure in
-    // `poll_ready`, so that callers such as load balancers could choose to call
-    // another service rather than waiting for buffer capacity.
-    //
-    // Unfortunately, this can't be done easily using Tokio's bounded MPSC
-    // channel, because it doesn't expose a polling-based interface, only an
-    // `async fn ready`, which borrows the sender. Therefore, we implement our
-    // own bounded MPSC on top of the unbounded channel, using a semaphore to
-    // limit how many items are in the channel.
-    semaphore: PollSemaphore,
-    // The current semaphore permit, if one has been acquired.
-    //
-    // This is acquired in `poll_ready` and taken in `call`.
-    permit: Option<OwnedSemaphorePermit>,
+    tx: PollSender<Message<Request, T::Future>>,
     handle: Handle,
 }
 
@@ -87,15 +69,13 @@ where
         T: Send + 'static,
         T::Error: Send + Sync,
         Request: Send + 'static,
+        T::Future: Send + 'static,
     {
-        let (tx, rx) = mpsc::unbounded_channel();
-        let semaphore = Arc::new(Semaphore::new(bound));
-        let (handle, worker) = Worker::new(service, rx, &semaphore);
-        let buffer = Buffer {
-            tx,
+        let (tx, rx) = mpsc::channel(bound);
+        let (handle, worker) = Worker::new(service, rx);
+        let buffer = Self {
+            tx: PollSender::new(tx),
             handle,
-            semaphore: PollSemaphore::new(semaphore),
-            permit: None,
         };
         (buffer, worker)
     }
@@ -109,6 +89,8 @@ impl<T, Request> Service<Request> for Buffer<T, Request>
 where
     T: Service<Request>,
     T::Error: Into<crate::BoxError>,
+    T::Future: Send + 'static,
+    Request: Send + 'static,
 {
     type Response = T::Response;
     type Error = crate::BoxError;
@@ -121,47 +103,31 @@ where
             return Poll::Ready(Err(self.get_worker_error()));
         }
 
-        // Then, check if we've already acquired a permit.
-        if self.permit.is_some() {
-            // We've already reserved capacity to send a request. We're ready!
-            return Poll::Ready(Ok(()));
-        }
-
-        // Finally, if we haven't already acquired a permit, poll the semaphore
-        // to acquire one. If we acquire a permit, then there's enough buffer
-        // capacity to send a new request. Otherwise, we need to wait for
-        // capacity.
-        let permit =
-            ready!(self.semaphore.poll_acquire(cx)).ok_or_else(|| self.get_worker_error())?;
-        self.permit = Some(permit);
-
-        Poll::Ready(Ok(()))
+        // Poll the sender to acquire a permit.
+        self.tx
+            .poll_reserve(cx)
+            .map_err(|_| self.get_worker_error())
     }
 
     fn call(&mut self, request: Request) -> Self::Future {
         tracing::trace!("sending request to buffer worker");
-        let _permit = self
-            .permit
-            .take()
-            .expect("buffer full; poll_ready must be called first");
 
         // 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 semaphore permit has already been
+        // 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(Message {
-            request,
-            span,
-            tx,
-            _permit,
-        }) {
-            Err(_) => ResponseFuture::failed(self.get_worker_error()),
+        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())
+            }
         }
     }
 }
@@ -169,15 +135,13 @@ where
 impl<T, Request> Clone for Buffer<T, Request>
 where
     T: Service<Request>,
+    Request: Send + 'static,
+    T::Future: Send + 'static,
 {
     fn clone(&self) -> Self {
         Self {
-            tx: self.tx.clone(),
             handle: self.handle.clone(),
-            semaphore: self.semaphore.clone(),
-            // The new clone hasn't acquired a permit yet. It will when it's
-            // next polled ready.
-            permit: None,
+            tx: self.tx.clone(),
         }
     }
 }
diff --git a/tower/src/buffer/worker.rs b/tower/src/buffer/worker.rs
index fe7ea55..7f4416d 100644
--- a/tower/src/buffer/worker.rs
+++ b/tower/src/buffer/worker.rs
@@ -3,13 +3,13 @@ use super::{
     message::Message,
 };
 use futures_core::ready;
-use std::sync::{Arc, Mutex, Weak};
+use std::sync::{Arc, Mutex};
 use std::{
     future::Future,
     pin::Pin,
     task::{Context, Poll},
 };
-use tokio::sync::{mpsc, Semaphore};
+use tokio::sync::mpsc;
 use tower_service::Service;
 
 pin_project_lite::pin_project! {
@@ -26,19 +26,11 @@ pin_project_lite::pin_project! {
         T: Service<Request>,
     {
         current_message: Option<Message<Request, T::Future>>,
-        rx: mpsc::UnboundedReceiver<Message<Request, T::Future>>,
+        rx: mpsc::Receiver<Message<Request, T::Future>>,
         service: T,
         finish: bool,
         failed: Option<ServiceError>,
         handle: Handle,
-        close: Option<Weak<Semaphore>>,
-    }
-
-    impl<T: Service<Request>, Request> PinnedDrop for Worker<T, Request>
-    {
-        fn drop(mut this: Pin<&mut Self>) {
-            this.as_mut().close_semaphore();
-        }
     }
 }
 
@@ -48,22 +40,6 @@ pub(crate) struct Handle {
     inner: Arc<Mutex<Option<ServiceError>>>,
 }
 
-impl<T, Request> Worker<T, Request>
-where
-    T: Service<Request>,
-{
-    /// Closes the buffer's semaphore if it is still open, waking any pending
-    /// tasks.
-    fn close_semaphore(&mut self) {
-        if let Some(close) = self.close.take().as_ref().and_then(Weak::upgrade) {
-            tracing::debug!("buffer closing; waking pending tasks");
-            close.close();
-        } else {
-            tracing::trace!("buffer already closed");
-        }
-    }
-}
-
 impl<T, Request> Worker<T, Request>
 where
     T: Service<Request>,
@@ -71,14 +47,12 @@ where
 {
     pub(crate) fn new(
         service: T,
-        rx: mpsc::UnboundedReceiver<Message<Request, T::Future>>,
-        semaphore: &Arc<Semaphore>,
+        rx: mpsc::Receiver<Message<Request, T::Future>>,
     ) -> (Handle, Worker<T, Request>) {
         let handle = Handle {
             inner: Arc::new(Mutex::new(None)),
         };
 
-        let semaphore = Arc::downgrade(semaphore);
         let worker = Worker {
             current_message: None,
             finish: false,
@@ -86,7 +60,6 @@ where
             rx,
             service,
             handle: handle.clone(),
-            close: Some(semaphore),
         };
 
         (handle, worker)
@@ -221,8 +194,6 @@ where
                                 .as_ref()
                                 .expect("Worker::failed did not set self.failed?")
                                 .clone()));
-                            // Wake any tasks waiting on channel capacity.
-                            self.close_semaphore();
                         }
                     }
                 }


Test patch:
diff --git a/tower/tests/buffer/main.rs b/tower/tests/buffer/main.rs
index b5e7c19..b65de6b 100644
--- a/tower/tests/buffer/main.rs
+++ b/tower/tests/buffer/main.rs
@@ -159,7 +159,7 @@ async fn waits_for_channel_capacity() {
 
     let (service, mut handle) = mock::pair::<&'static str, &'static str>();
 
-    let (service, worker) = Buffer::pair(service, 3);
+    let (service, worker) = Buffer::pair(service, 2);
 
     let mut service = mock::Spawn::new(service);
     let mut worker = task::spawn(worker);
@@ -243,6 +243,14 @@ async fn wakes_pending_waiters_on_close() {
     assert_pending!(worker.poll());
     let mut response = task::spawn(service1.call("hello"));
 
+    assert!(worker.is_woken(), "worker task should be woken by request");
+    assert_pending!(worker.poll());
+
+    // fill the channel so all subsequent requests will wait for capacity
+    let service1 = assert_ready_ok!(task::spawn(service.ready()).poll());
+    assert_pending!(worker.poll());
+    let mut response2 = task::spawn(service1.call("world"));
+
     let mut service1 = service.clone();
     let mut ready1 = task::spawn(service1.ready());
     assert_pending!(worker.poll());
@@ -263,6 +271,13 @@ async fn wakes_pending_waiters_on_close() {
         err
     );
 
+    let err = assert_ready_err!(response2.poll());
+    assert!(
+        err.is::<error::Closed>(),
+        "response should fail with a Closed, got: {:?}",
+        err
+    );
+
     assert!(
         ready1.is_woken(),
         "dropping worker should wake ready task 1"
@@ -301,6 +316,14 @@ async fn wakes_pending_waiters_on_failure() {
     assert_pending!(worker.poll());
     let mut response = task::spawn(service1.call("hello"));
 
+    assert!(worker.is_woken(), "worker task should be woken by request");
+    assert_pending!(worker.poll());
+
+    // fill the channel so all subsequent requests will wait for capacity
+    let service1 = assert_ready_ok!(task::spawn(service.ready()).poll());
+    assert_pending!(worker.poll());
+    let mut response2 = task::spawn(service1.call("world"));
+
     let mut service1 = service.clone();
     let mut ready1 = task::spawn(service1.ready());
     assert_pending!(worker.poll());
@@ -322,6 +345,12 @@ async fn wakes_pending_waiters_on_failure() {
         "response should fail with a ServiceError, got: {:?}",
         err
     );
+    let err = assert_ready_err!(response2.poll());
+    assert!(
+        err.is::<error::ServiceError>(),
+        "response should fail with a ServiceError, got: {:?}",
+        err
+    );
 
     assert!(
         ready1.is_woken(),