Repository: tower-rs/tower
Pull request #671 (2022-06-17 19:13:37)

backport non-breaking changes from `master` to v0.4.x
This branch backports all of the non-breaking changes made on `master`
since `master` was bumped to `tower` 0.5:

* ae745061 ready-cache: Ensure cancelation can be observed (#668)
* a9409115 ready_cache: just use pin_project for `Pending` (#667)
* 5419e425 load_shed: make constructor for Overloaded error public (#661)
* 697b8d0e service: clarify docs around shared resource consumption in poll_ready() (#662)
* 9b6c4862 fix: broken Service link (#659)
* 44880d13 util: Fix `call_all` hang when stream is pending (#656)
* 7edd03af tower: fix docs typo (#646)

olix0r: I think you'll need https://github.com/tower-rs/tower/pull/670 as well?

Patch:
diff --git a/tower-layer/src/lib.rs b/tower-layer/src/lib.rs
index 2187375..48ff79f 100644
--- a/tower-layer/src/lib.rs
+++ b/tower-layer/src/lib.rs
@@ -14,7 +14,7 @@
 //!
 //! A middleware implements the [`Layer`] and [`Service`] trait.
 //!
-//! [`Service`]: https://docs.rs/tower/trait.Service.html
+//! [`Service`]: https://docs.rs/tower/latest/tower/trait.Service.html
 
 mod identity;
 mod layer_fn;
@@ -90,7 +90,7 @@ pub use self::{
 /// is also decoupled from client or server concerns. In other words, the same
 /// log middleware could be used in either a client or a server.
 ///
-/// [`Service`]: https://docs.rs/tower/trait.Service.html
+/// [`Service`]: https://docs.rs/tower/latest/tower/trait.Service.html
 pub trait Layer<S> {
     /// The wrapped service
     type Service;
diff --git a/tower-service/src/lib.rs b/tower-service/src/lib.rs
index c2f531e..a6ef0ec 100644
--- a/tower-service/src/lib.rs
+++ b/tower-service/src/lib.rs
@@ -331,6 +331,12 @@ pub trait Service<Request> {
     /// Once `poll_ready` returns `Poll::Ready(Ok(()))`, a request may be dispatched to the
     /// service using `call`. Until a request is dispatched, repeated calls to
     /// `poll_ready` must return either `Poll::Ready(Ok(()))` or `Poll::Ready(Err(_))`.
+    ///
+    /// Note that `poll_ready` may reserve shared resources that are consumed in a subsequent
+    /// invocation of `call`. Thus, it is critical for implementations to not assume that `call`
+    /// will always be invoked and to ensure that such resources are released if the service is
+    /// dropped before `call` is invoked or the future returned by `call` is dropped before it
+    /// is polled.
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
 
     /// Process the request and return the response asynchronously.
diff --git a/tower/Cargo.toml b/tower/Cargo.toml
index fe8e44d..4b6b368 100644
--- a/tower/Cargo.toml
+++ b/tower/Cargo.toml
@@ -57,7 +57,7 @@ limit = ["__common", "tokio/time", "tokio/sync", "tokio-util", "tracing"]
 load = ["__common", "tokio/time", "tracing"]
 load-shed = ["__common"]
 make = ["futures-util", "pin-project-lite", "tokio/io-std"]
-ready-cache = ["futures-core", "futures-util", "indexmap", "tokio/sync", "tracing"]
+ready-cache = ["futures-core", "futures-util", "indexmap", "tokio/sync", "tracing", "pin-project-lite"]
 reconnect = ["make", "tokio/io-std", "tracing"]
 retry = ["__common", "tokio/time"]
 spawn-ready = ["__common", "futures-util", "tokio/sync", "tokio/rt", "util", "tracing"]
diff --git a/tower/src/lib.rs b/tower/src/lib.rs
index c7f72cf..10230b8 100644
--- a/tower/src/lib.rs
+++ b/tower/src/lib.rs
@@ -141,7 +141,7 @@
 //! Tower.
 //!
 //!
-//! #! Supported Rust Versions
+//! ## Supported Rust Versions
 //!
 //! Tower will keep a rolling MSRV (minimum supported Rust version) policy of **at
 //! least** 6 months. When increasing the MSRV, the new Rust version must have been
diff --git a/tower/src/load_shed/error.rs b/tower/src/load_shed/error.rs
index ef401f0..3a4c025 100644
--- a/tower/src/load_shed/error.rs
+++ b/tower/src/load_shed/error.rs
@@ -7,12 +7,14 @@ use std::fmt;
 /// called.
 ///
 /// [`LoadShed`]: crate::load_shed::LoadShed
+#[derive(Default)]
 pub struct Overloaded {
     _p: (),
 }
 
 impl Overloaded {
-    pub(crate) fn new() -> Self {
+    /// Construct a new overloaded error
+    pub fn new() -> Self {
         Overloaded { _p: () }
     }
 }
diff --git a/tower/src/ready_cache/cache.rs b/tower/src/ready_cache/cache.rs
index f0fc1d0..8070d78 100644
--- a/tower/src/ready_cache/cache.rs
+++ b/tower/src/ready_cache/cache.rs
@@ -2,15 +2,16 @@
 
 use super::error;
 use futures_core::Stream;
-use futures_util::stream::FuturesUnordered;
+use futures_util::{stream::FuturesUnordered, task::AtomicWaker};
 pub use indexmap::Equivalent;
 use indexmap::IndexMap;
 use std::fmt;
 use std::future::Future;
 use std::hash::Hash;
 use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::task::{Context, Poll};
-use tokio::sync::oneshot;
 use tower_service::Service;
 use tracing::{debug, trace};
 
@@ -75,8 +76,18 @@ where
 // Safety: This is safe because we do not use `Pin::new_unchecked`.
 impl<S, K: Eq + Hash, Req> Unpin for ReadyCache<K, S, Req> {}
 
-type CancelRx = oneshot::Receiver<()>;
-type CancelTx = oneshot::Sender<()>;
+#[derive(Debug)]
+struct Cancel {
+    waker: AtomicWaker,
+    canceled: AtomicBool,
+}
+
+#[derive(Debug)]
+struct CancelRx(Arc<Cancel>);
+
+#[derive(Debug)]
+struct CancelTx(Arc<Cancel>);
+
 type CancelPair = (CancelTx, CancelRx);
 
 #[derive(Debug)]
@@ -85,14 +96,16 @@ enum PendingError<K, E> {
     Inner(K, E),
 }
 
-/// A [`Future`] that becomes satisfied when an `S`-typed service is ready.
-///
-/// May fail due to cancelation, i.e. if the service is evicted from the balancer.
-struct Pending<K, S, Req> {
-    key: Option<K>,
-    cancel: Option<CancelRx>,
-    ready: Option<S>,
-    _pd: std::marker::PhantomData<Req>,
+pin_project_lite::pin_project! {
+    /// A [`Future`] that becomes satisfied when an `S`-typed service is ready.
+    ///
+    /// May fail due to cancelation, i.e. if the service is evicted from the balancer.
+    struct Pending<K, S, Req> {
+        key: Option<K>,
+        cancel: Option<CancelRx>,
+        ready: Option<S>,
+        _pd: std::marker::PhantomData<Req>,
+    }
 }
 
 // === ReadyCache ===
@@ -193,7 +206,7 @@ where
     /// must be called to cause the service to be dropped.
     pub fn evict<Q: Hash + Equivalent<K>>(&mut self, key: &Q) -> bool {
         let canceled = if let Some(c) = self.pending_cancel_txs.swap_remove(key) {
-            c.send(()).expect("cancel receiver lost");
+            c.cancel();
             true
         } else {
             false
@@ -224,14 +237,14 @@ where
     ///
     /// [`poll_pending`]: crate::ready_cache::cache::ReadyCache::poll_pending
     pub fn push(&mut self, key: K, svc: S) {
-        let cancel = oneshot::channel();
+        let cancel = cancelable();
         self.push_pending(key, svc, cancel);
     }
 
     fn push_pending(&mut self, key: K, svc: S, (cancel_tx, cancel_rx): CancelPair) {
         if let Some(c) = self.pending_cancel_txs.insert(key.clone(), cancel_tx) {
             // If there is already a service for this key, cancel it.
-            c.send(()).expect("cancel receiver lost");
+            c.cancel();
         }
         self.pending.push(Pending {
             key: Some(key),
@@ -268,21 +281,10 @@ where
                         // recreated after the service is used.
                         self.ready.insert(key, (svc, (cancel_tx, cancel_rx)));
                     } else {
-                        // This should not technically be possible. We must have decided to cancel
-                        // a Service (by sending on the CancelTx), yet that same service then
-                        // returns Ready. Since polling a Pending _first_ polls the CancelRx, that
-                        // _should_ always see our CancelTx send. Yet empirically, that isn't true:
-                        //
-                        //   https://github.com/tower-rs/tower/issues/415
-                        //
-                        // So, we instead detect the endpoint as canceled at this point. That
-                        // should be fine, since the oneshot is only really there to ensure that
-                        // the Pending is polled again anyway.
-                        //
-                        // We assert that this can't happen in debug mode so that hopefully one day
-                        // we can find a test that triggers this reliably.
-                        debug_assert!(cancel_tx.is_some());
-                        debug!("canceled endpoint removed when ready");
+                        assert!(
+                            cancel_tx.is_some(),
+                            "services that become ready must have a pending cancelation"
+                        );
                     }
                 }
                 Poll::Ready(Some(Err(PendingError::Canceled(_)))) => {
@@ -292,13 +294,11 @@ where
                 }
                 Poll::Ready(Some(Err(PendingError::Inner(key, e)))) => {
                     let cancel_tx = self.pending_cancel_txs.swap_remove(&key);
-                    if cancel_tx.is_some() {
-                        return Err(error::Failed(key, e.into())).into();
-                    } else {
-                        // See comment for the same clause under Ready(Some(Ok)).
-                        debug_assert!(cancel_tx.is_some());
-                        debug!("canceled endpoint removed on error");
-                    }
+                    assert!(
+                        cancel_tx.is_some(),
+                        "services that return an error must have a pending cancelation"
+                    );
+                    return Err(error::Failed(key, e.into())).into();
                 }
             }
         }
@@ -398,10 +398,29 @@ where
     }
 }
 
-// === Pending ===
+// === impl Cancel ===
 
-// Safety: No use unsafe access therefore this is safe.
-impl<K, S, Req> Unpin for Pending<K, S, Req> {}
+/// Creates a cancelation sender and receiver.
+///
+/// A `tokio::sync::oneshot` is NOT used, as a `Receiver` is not guaranteed to
+/// observe results as soon as a `Sender` fires. Using an `AtomicBool` allows
+/// the state to be observed as soon as the cancelation is triggered.
+fn cancelable() -> CancelPair {
+    let cx = Arc::new(Cancel {
+        waker: AtomicWaker::new(),
+        canceled: AtomicBool::new(false),
+    });
+    (CancelTx(cx.clone()), CancelRx(cx))
+}
+
+impl CancelTx {
+    fn cancel(self) {
+        self.0.canceled.store(true, Ordering::SeqCst);
+        self.0.waker.wake();
+    }
+}
+
+// === Pending ===
 
 impl<K, S, Req> Future for Pending<K, S, Req>
 where
@@ -409,28 +428,44 @@ where
 {
     type Output = Result<(K, S, CancelRx), PendingError<K, S::Error>>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let mut fut = self.cancel.as_mut().expect("polled after complete");
-        if let Poll::Ready(r) = Pin::new(&mut fut).poll(cx) {
-            assert!(r.is_ok(), "cancel sender lost");
-            let key = self.key.take().expect("polled after complete");
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        // Before checking whether the service is ready, check to see whether
+        // readiness has been canceled.
+        let CancelRx(cancel) = this.cancel.as_mut().expect("polled after complete");
+        if cancel.canceled.load(Ordering::SeqCst) {
+            let key = this.key.take().expect("polled after complete");
             return Err(PendingError::Canceled(key)).into();
         }
 
-        match self
+        match this
             .ready
             .as_mut()
             .expect("polled after ready")
             .poll_ready(cx)
         {
-            Poll::Pending => Poll::Pending,
+            Poll::Pending => {
+                // Before returning Pending, register interest in cancelation so
+                // that this future is polled again if the state changes.
+                let CancelRx(cancel) = this.cancel.as_mut().expect("polled after complete");
+                cancel.waker.register(cx.waker());
+                // Because both the cancel receiver and cancel sender are held
+                // by the `ReadyCache` (i.e., on a single task), then it must
+                // not be possible for the cancelation state to change while
+                // polling a `Pending` service.
+                assert!(
+                    !cancel.canceled.load(Ordering::SeqCst),
+                    "cancelation cannot be notified while polling a pending service"
+                );
+                Poll::Pending
+            }
             Poll::Ready(Ok(())) => {
-                let key = self.key.take().expect("polled after complete");
-                let cancel = self.cancel.take().expect("polled after complete");
-                Ok((key, self.ready.take().expect("polled after ready"), cancel)).into()
+                let key = this.key.take().expect("polled after complete");
+                let cancel = this.cancel.take().expect("polled after complete");
+                Ok((key, this.ready.take().expect("polled after ready"), cancel)).into()
             }
             Poll::Ready(Err(e)) => {
-                let key = self.key.take().expect("polled after compete");
+                let key = this.key.take().expect("polled after compete");
                 Err(PendingError::Inner(key, e)).into()
             }
         }
diff --git a/tower/src/util/call_all/common.rs b/tower/src/util/call_all/common.rs
index 09c458a..8e37afc 100644
--- a/tower/src/util/call_all/common.rs
+++ b/tower/src/util/call_all/common.rs
@@ -99,20 +99,18 @@ where
                 .expect("Using CallAll after extracing inner Service");
             ready!(svc.poll_ready(cx)).map_err(Into::into)?;
 
-            // If it is, gather the next request (if there is one)
-            match this.stream.as_mut().poll_next(cx) {
-                Poll::Ready(r) => match r {
-                    Some(req) => {
-                        this.queue.push(svc.call(req));
-                    }
-                    None => {
-                        // We're all done once any outstanding requests have completed
-                        *this.eof = true;
-                    }
-                },
-                Poll::Pending => {
-                    // TODO: We probably want to "release" the slot we reserved in Svc here.
-                    // It may be a while until we get around to actually using it.
+            // If it is, gather the next request (if there is one), or return `Pending` if the
+            // stream is not ready.
+            // TODO: We probably want to "release" the slot we reserved in Svc if the
+            // stream returns `Pending`. It may be a while until we get around to actually
+            // using it.
+            match ready!(this.stream.as_mut().poll_next(cx)) {
+                Some(req) => {
+                    this.queue.push(svc.call(req));
+                }
+                None => {
+                    // We're all done once any outstanding requests have completed
+                    *this.eof = true;
                 }
             }
         }


Test patch:
diff --git a/tower/tests/ready_cache/main.rs b/tower/tests/ready_cache/main.rs
index f7acbf8..56be6a9 100644
--- a/tower/tests/ready_cache/main.rs
+++ b/tower/tests/ready_cache/main.rs
@@ -2,8 +2,9 @@
 #[path = "../support.rs"]
 mod support;
 
+use std::pin::Pin;
 use tokio_test::{assert_pending, assert_ready, task};
-use tower::ready_cache::ReadyCache;
+use tower::ready_cache::{error, ReadyCache};
 use tower_test::mock;
 
 type Req = &'static str;
@@ -191,3 +192,32 @@ fn duplicate_key_by_index() {
     // _and_ service 0 should now be callable
     assert!(task.enter(|cx, _| cache.check_ready(cx, &0)).unwrap());
 }
+
+// Tests https://github.com/tower-rs/tower/issues/415
+#[tokio::test(flavor = "current_thread")]
+async fn cancelation_observed() {
+    let mut cache = ReadyCache::default();
+    let mut handles = vec![];
+
+    // NOTE This test passes at 129 items, but fails at 130 items (if coop
+    // schedulding interferes with cancelation).
+    for _ in 0..130 {
+        let (svc, mut handle) = tower_test::mock::pair::<(), ()>();
+        handle.allow(1);
+        cache.push("ep0", svc);
+        handles.push(handle);
+    }
+
+    struct Ready(ReadyCache<&'static str, tower_test::mock::Mock<(), ()>, ()>);
+    impl Unpin for Ready {}
+    impl std::future::Future for Ready {
+        type Output = Result<(), error::Failed<&'static str>>;
+        fn poll(
+            self: Pin<&mut Self>,
+            cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Self::Output> {
+            self.get_mut().0.poll_pending(cx)
+        }
+    }
+    Ready(cache).await.unwrap();
+}
diff --git a/tower/tests/util/call_all.rs b/tower/tests/util/call_all.rs
index 6bc0929..02e69ef 100644
--- a/tower/tests/util/call_all.rs
+++ b/tower/tests/util/call_all.rs
@@ -143,3 +143,24 @@ async fn unordered() {
         .unwrap();
     assert!(v.is_none());
 }
+
+#[tokio::test]
+async fn pending() {
+    let _t = support::trace_init();
+
+    let (mock, mut handle) = mock::pair::<_, &'static str>();
+
+    let mut task = task::spawn(());
+
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let ca = mock.call_all(support::IntoStream::new(rx));
+    pin_mut!(ca);
+
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    tx.send("req").unwrap();
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    assert_request_eq!(handle, "req").send_response("res");
+    let res = assert_ready!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+    assert_eq!(res.transpose().unwrap(), Some("res"));
+    assert_pending!(task.enter(|cx, _| ca.as_mut().poll_next(cx)));
+}