ready-cache: Ensure cancelation can be observed

tower/src/ready_cache/cache.rs

@@ -9,8 +9,10 @@ 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 tokio::sync::Notify;
 use tower_service::Service;
 use tracing::{debug, trace};
 
@@ -75,8 +77,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 {
+    notify: Notify,
+    canceled: AtomicBool,
+}
+
+#[derive(Debug)]
+struct CancelRx(Arc<Cancel>);
+
+#[derive(Debug)]
+struct CancelTx(Arc<Cancel>);
+
 type CancelPair = (CancelTx, CancelRx);
 
 #[derive(Debug)]
@@ -193,7 +205,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 +236,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 +280,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 +293,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,6 +397,28 @@ where
     }
 }
 
+// === impl Cancel ===
+
+/// 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 {
+        notify: Notify::new(),
+        canceled: AtomicBool::new(false),
+    });
+    (CancelTx(cx.clone()), CancelRx(cx))
+}
+
+impl CancelTx {
+    fn cancel(self) {
+        self.0.canceled.store(true, Ordering::SeqCst);
+        self.0.notify.notify_waiters();
+    }
+}
+
 // === Pending ===
 
 // Safety: No use unsafe access therefore this is safe.
@@ -410,20 +431,40 @@ 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");
+        // Before checking whether a service ready, check to see whether
+        // readiness has been canceled.
+        let CancelRx(cancel) = self.cancel.as_mut().expect("polled after complete");
+        if cancel.canceled.load(Ordering::SeqCst) {
             let key = self.key.take().expect("polled after complete");
             return Err(PendingError::Canceled(key)).into();
         }
 
+        // If readiness hasn't been canceled, check to see whether the service
+        // is ready.
         match self
             .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) = self.cancel.as_mut().expect("polled after complete");
+                tokio::pin! {
+                    let cancel_notified = cancel.notify.notified();
+                }
+                let crx_ready = cancel_notified.poll(cx);
+                // 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.
+                assert!(
+                    crx_ready.is_pending() && !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");

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();
+}