fix(http1): synchronize dispatch shutdown with sends

src/client/dispatch.rs

@@ -192,17 +192,18 @@ impl<T, U> Receiver<T, U> {
     }
 
     #[cfg(feature = "http1")]
-    pub(crate) fn close(&mut self) {
+    pub(crate) fn poll_close_and_recv(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<(T, Callback<T, U>)>> {
         self.taker.cancel();
         self.inner.close();
-    }
-
-    #[cfg(feature = "http1")]
-    pub(crate) fn try_recv(&mut self) -> Option<(T, Callback<T, U>)> {
-        match crate::common::task::now_or_never(self.inner.recv()) {
-            Some(Some(mut env)) => env.0.take(),
-            _ => None,
-        }
+        // A send may have reserved capacity before close() without publishing
+        // its envelope yet. Let it wake this task rather than spinning, including
+        // when Tokio's cooperative task budget is exhausted.
+        self.inner
+            .poll_recv(cx)
+            .map(|item| item.map(|mut env| env.0.take().expect("envelope not dropped")))
     }
 }
 
@@ -394,11 +395,102 @@ mod tests {
     use std::pin::Pin;
     use std::task::{Context, Poll};
 
+    #[cfg(feature = "http1")]
+    use std::sync::atomic::{AtomicBool, Ordering};
+    #[cfg(feature = "http1")]
+    use std::sync::{mpsc as std_mpsc, Arc, Barrier};
+
     use super::{channel, Callback, Receiver};
 
     #[derive(Debug)]
     struct Custom(#[allow(dead_code)] i32);
 
+    #[cfg(feature = "http1")]
+    struct TrackDrop(Arc<AtomicBool>);
+
+    #[cfg(feature = "http1")]
+    impl Drop for TrackDrop {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    #[cfg(feature = "http1")]
+    #[tokio::test]
+    async fn receiver_shutdown_yields_when_task_budget_is_exhausted() {
+        use futures_util::FutureExt;
+
+        tokio::spawn(async {
+            let (mut tx, mut rx) = channel::<i32, ()>();
+            let _promise = tx.try_send(43).unwrap();
+
+            // Exhaust this task's budget without yielding, leaving a queued
+            // message to distinguish budget exhaustion from an empty channel.
+            let (budget_tx, mut budget_rx) = tokio::sync::mpsc::unbounded_channel();
+            for _ in 0..1024 {
+                budget_tx.send(()).unwrap();
+            }
+            while budget_rx.recv().now_or_never().is_some() {}
+            assert!(budget_rx.try_recv().is_ok());
+
+            let (value, callback) = futures_util::future::poll_fn(|cx| rx.poll_close_and_recv(cx))
+                .await
+                .expect("queued request");
+            assert_eq!(value, 43);
+            drop(callback);
+            assert!(
+                futures_util::future::poll_fn(|cx| rx.poll_close_and_recv(cx))
+                    .await
+                    .is_none()
+            );
+        })
+        .await
+        .unwrap();
+    }
+
+    #[cfg(feature = "http1")]
+    #[test]
+    fn receiver_shutdown_reclaims_concurrent_send() {
+        let rt = tokio::runtime::Builder::new_current_thread()
+            .build()
+            .unwrap();
+        let barrier = Arc::new(Barrier::new(2));
+        let worker_barrier = barrier.clone();
+        let (work_tx, work_rx) = std_mpsc::channel::<(super::Sender<TrackDrop, ()>, TrackDrop)>();
+        let (done_tx, done_rx) = std_mpsc::channel();
+        let worker = std::thread::spawn(move || {
+            while let Ok((mut tx, val)) = work_rx.recv() {
+                worker_barrier.wait();
+                drop(tx.send(val));
+                worker_barrier.wait();
+                done_tx.send(tx).unwrap();
+            }
+        });
+
+        // Tokio's unbounded channel reserves capacity before publishing the
+        // envelope. Repeat enough times to exercise shutdown in that window.
+        for _ in 0..10_000 {
+            let (tx, mut rx) = channel::<TrackDrop, ()>();
+            let dropped = Arc::new(AtomicBool::new(false));
+            work_tx.send((tx, TrackDrop(dropped.clone()))).unwrap();
+            barrier.wait();
+            drop(rt.block_on(futures_util::future::poll_fn(|cx| {
+                rx.poll_close_and_recv(cx)
+            })));
+            drop(rx);
+            barrier.wait();
+            let tx = done_rx.recv().unwrap();
+            assert!(
+                dropped.load(Ordering::SeqCst),
+                "value remained queued after the receiver was dropped"
+            );
+            drop(tx);
+        }
+
+        drop(work_tx);
+        worker.join().unwrap();
+    }
+
     impl<T, U> Future for Receiver<T, U> {
         type Output = Option<(T, Callback<T, U>)>;
 

src/common/task.rs

@@ -1,6 +1,4 @@
 use std::task::{Context, Poll};
-#[cfg(feature = "client")]
-use std::task::{RawWaker, RawWakerVTable, Waker};
 
 /// A function to help "yield" a future, such that it is re-scheduled immediately.
 ///
@@ -9,37 +7,3 @@ pub(crate) fn yield_now(cx: &mut Context<'_>) -> Poll<std::convert::Infallible>
     cx.waker().wake_by_ref();
     Poll::Pending
 }
-
-// TODO: replace with `std::task::Waker::noop()` once MSRV >= 1.85
-#[cfg(feature = "client")]
-fn noop_waker() -> Waker {
-    const NOOP_RAW_WAKER: RawWaker = RawWaker::new(std::ptr::null(), &NOOP_VTABLE);
-    const NOOP_VTABLE: RawWakerVTable = RawWakerVTable::new(
-        // `clone` returns the same noop waker again
-        |_: *const ()| NOOP_RAW_WAKER,
-        // `wake`, `wake_by_ref`, and `drop` do nothing
-        |_: *const ()| {},
-        |_: *const ()| {},
-        |_: *const ()| {},
-    );
-
-    // SAFETY: all functions in the vtable are safe to call, and Waker's safety does not require
-    // them to actually do anything.
-    unsafe { Waker::from_raw(NOOP_RAW_WAKER) }
-}
-
-/// Poll the future once and return `Some` if it is ready, else `None`.
-///
-/// If the future wasn't ready, it future likely can't be driven to completion any more: the polling
-/// uses a no-op waker, so knowledge of what the pending future was waiting for is lost.
-#[cfg(feature = "client")]
-pub(crate) fn now_or_never<F: std::future::Future>(fut: F) -> Option<F::Output> {
-    let waker = noop_waker();
-    let mut cx = Context::from_waker(&waker);
-    // TODO: replace with std::pin::pin! once MSRV >= 1.68
-    tokio::pin!(fut);
-    match fut.poll(&mut cx) {
-        Poll::Ready(res) => Some(res),
-        Poll::Pending => None,
-    }
-}

src/proto/h1/dispatch.rs

@@ -25,6 +25,7 @@ pub(crate) struct Dispatcher<D, Bs: Body, I, T> {
     body_tx: SenderDropGuard,
     body_rx: Pin<Box<Option<Bs>>>,
     is_closing: bool,
+    error: Option<crate::Error>,
 }
 
 pub(crate) trait Dispatch {
@@ -38,6 +39,13 @@ pub(crate) trait Dispatch {
     ) -> Poll<Option<Result<(Self::PollItem, Self::PollBody), Self::PollError>>>;
     fn recv_msg(&mut self, msg: crate::Result<(Self::RecvItem, IncomingBody)>)
         -> crate::Result<()>;
+    fn poll_recv_error(
+        &mut self,
+        _cx: &mut Context<'_>,
+        error: &mut Option<crate::Error>,
+    ) -> Poll<crate::Result<()>> {
+        Poll::Ready(self.recv_msg(Err(error.take().expect("pending connection error"))))
+    }
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), ()>>;
     fn should_poll(&self) -> bool;
 }
@@ -84,6 +92,7 @@ where
             body_tx: SenderDropGuard::none(),
             body_rx: Box::pin(None),
             is_closing: false,
+            error: None,
         }
     }
 
@@ -125,19 +134,21 @@ where
         cx: &mut Context<'_>,
         should_shutdown: bool,
     ) -> Poll<crate::Result<Dispatched>> {
-        Poll::Ready(ready!(self.poll_inner(cx, should_shutdown)).or_else(|e| {
+        if self.error.is_none() {
+            match ready!(self.poll_inner(cx, should_shutdown)) {
+                Ok(dispatched) => return Poll::Ready(Ok(dispatched)),
+                Err(error) => self.error = Some(error),
+            }
             // Be sure to alert a streaming body of the failure with a
             // more specific error than the drop guard would provide.
             if let Some(mut body) = self.body_tx.take() {
                 body.send_error(crate::Error::new_body("connection error"));
             }
-            // An error means we're shutting down either way.
-            // We just try to give the error to the user,
-            // and close the connection with an Ok. If we
-            // cannot give it to the user, then return the Err.
-            self.dispatch.recv_msg(Err(e))?;
-            Ok(Dispatched::Shutdown)
-        }))
+        }
+        // Keep the error while the client drains an in-flight send. Pending may
+        // also mean that Tokio needs this task to yield its cooperative budget.
+        ready!(self.dispatch.poll_recv_error(cx, &mut self.error))?;
+        Poll::Ready(Ok(Dispatched::Shutdown))
     }
 
     fn poll_inner(
@@ -333,12 +344,7 @@ where
             }
             Some(Err(err)) => {
                 debug!("read_head error: {}", err);
-                self.dispatch.recv_msg(Err(err))?;
-                // if here, the dispatcher gave the user the error
-                // somewhere else. we still need to shutdown, but
-                // not as a second error.
-                self.close();
-                Poll::Ready(Ok(()))
+                Poll::Ready(Err(err))
             }
             None => {
                 // read eof, the write side will have been closed too unless
@@ -727,27 +733,36 @@ cfg_client! {
                             message: None,
                         }));
                         Ok(())
-                    } else if !self.rx_closed {
-                        self.rx.close();
-                        if let Some((req, cb)) = self.rx.try_recv() {
-                            trace!("canceling queued request with connection error: {}", err);
-                            // in this case, the message was never even started, so it's safe to tell
-                            // the user that the request was completely canceled
-                            cb.send(Err(TrySendError {
-                                error: crate::Error::new_canceled().with(err),
-                                message: Some(req),
-                            }));
-                            Ok(())
-                        } else {
-                            Err(err)
-                        }
                     } else {
                         Err(err)
                     }
                 }
             }
         }
 
+        fn poll_recv_error(
+            &mut self,
+            cx: &mut Context<'_>,
+            error: &mut Option<crate::Error>,
+        ) -> Poll<crate::Result<()>> {
+            if self.callback.is_none() && !self.rx_closed {
+                let queued = ready!(self.rx.poll_close_and_recv(cx));
+                self.rx_closed = true;
+                if let Some((req, cb)) = queued {
+                    let err = error.take().expect("pending connection error");
+                    trace!("canceling queued request with connection error: {}", err);
+                    // The request was never started, so it can be returned to
+                    // the caller along with the original connection error.
+                    cb.send(Err(TrySendError {
+                        error: crate::Error::new_canceled().with(err),
+                        message: Some(req),
+                    }));
+                    return Poll::Ready(Ok(()));
+                }
+            }
+            Poll::Ready(self.recv_msg(Err(error.take().expect("pending connection error"))))
+        }
+
         fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), ()>> {
             match &mut self.callback {
                 Some(cb) => match cb.poll_canceled(cx) {
@@ -809,6 +824,46 @@ mod tests {
         });
     }
 
+    #[tokio::test]
+    async fn client_shutdown_preserves_error_across_budget_yield() {
+        use futures_util::{future::poll_fn, FutureExt};
+
+        tokio::spawn(async {
+            let io = tokio_test::io::Builder::new().build();
+            let (mut tx, rx) = crate::client::dispatch::channel();
+            let conn = Conn::<_, bytes::Bytes, ClientTransaction>::new(Compat::new(io));
+            let mut dispatcher = Dispatcher::new(Client::new(rx), conn);
+            // Start at error delivery so mock IO does not consume the exhausted
+            // budget before the dispatch receiver is polled.
+            dispatcher.error = Some(crate::Error::new_unexpected_message());
+            let response = tx
+                .try_send(crate::Request::new(IncomingBody::empty()))
+                .unwrap();
+
+            let (budget_tx, mut budget_rx) = tokio::sync::mpsc::unbounded_channel();
+            for _ in 0..1024 {
+                budget_tx.send(()).unwrap();
+            }
+            while budget_rx.recv().now_or_never().is_some() {}
+            assert!(budget_rx.try_recv().is_ok());
+
+            assert!(poll_fn(|cx| Pin::new(&mut dispatcher).poll(cx))
+                .now_or_never()
+                .is_none());
+            assert!(dispatcher.error.is_some());
+
+            dispatcher.await.unwrap();
+            let err = response
+                .await
+                .unwrap()
+                .expect_err("request must be canceled");
+            assert!(err.error.is_canceled());
+            assert!(err.message.is_some());
+        })
+        .await
+        .unwrap();
+    }
+
     #[cfg(not(miri))]
     #[tokio::test]
     async fn client_flushing_is_not_ready_for_next_request() {