From 66638807dd412301b7a6fe582d1b22ca0b1dbe41 Mon Sep 17 00:00:00 2001 From: Morax Date: Mon, 10 Aug 2026 14:16:47 +0200 Subject: [PATCH 1/4] fix(http1): synchronize dispatch shutdown with sends Drain a closed HTTP/1 dispatch channel until any concurrent synchronous send has either published or observed closure, so the queued request cannot remain owned by the sender. Closes #4122 --- src/client/dispatch.rs | 71 +++++++++++++++++++++++++++++++++++----- src/proto/h1/dispatch.rs | 3 +- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs index a9375f627a..c8554838be 100644 --- a/src/client/dispatch.rs +++ b/src/client/dispatch.rs @@ -192,16 +192,18 @@ impl Receiver { } #[cfg(feature = "http1")] - pub(crate) fn close(&mut self) { + pub(crate) fn close_and_recv(&mut self) -> Option<(T, Callback)> { self.taker.cancel(); self.inner.close(); - } - - #[cfg(feature = "http1")] - pub(crate) fn try_recv(&mut self) -> Option<(T, Callback)> { - match crate::common::task::now_or_never(self.inner.recv()) { - Some(Some(mut env)) => env.0.take(), - _ => None, + loop { + match crate::common::task::now_or_never(self.inner.recv()) { + Some(Some(mut env)) => return env.0.take(), + Some(None) => return None, + // A send can reserve capacity immediately before close() and + // publish its envelope immediately afterwards. Once closed, + // Pending means that such a synchronous send is still in flight. + None => std::hint::spin_loop(), + } } } } @@ -394,11 +396,64 @@ 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); + + #[cfg(feature = "http1")] + impl Drop for TrackDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + #[cfg(feature = "http1")] + #[test] + fn receiver_shutdown_reclaims_concurrent_send() { + let barrier = Arc::new(Barrier::new(2)); + let worker_barrier = barrier.clone(); + let (work_tx, work_rx) = std_mpsc::channel::<(super::Sender, 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::(); + let dropped = Arc::new(AtomicBool::new(false)); + work_tx.send((tx, TrackDrop(dropped.clone()))).unwrap(); + barrier.wait(); + drop(rx.close_and_recv()); + 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 Future for Receiver { type Output = Option<(T, Callback)>; diff --git a/src/proto/h1/dispatch.rs b/src/proto/h1/dispatch.rs index d98aebcf1f..905859a22c 100644 --- a/src/proto/h1/dispatch.rs +++ b/src/proto/h1/dispatch.rs @@ -728,8 +728,7 @@ cfg_client! { })); Ok(()) } else if !self.rx_closed { - self.rx.close(); - if let Some((req, cb)) = self.rx.try_recv() { + if let Some((req, cb)) = self.rx.close_and_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 From fca16719f5ed2053577d2f462a69ed512c79bc5b Mon Sep 17 00:00:00 2001 From: Morax Date: Thu, 8 Oct 2026 05:25:03 +0200 Subject: [PATCH 2/4] fix(http1): avoid task budgets during dispatch shutdown --- src/client/dispatch.rs | 37 ++++++++++++++++++++++++++++++++----- src/common/task.rs | 36 ------------------------------------ 2 files changed, 32 insertions(+), 41 deletions(-) diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs index c8554838be..3d54fd5b63 100644 --- a/src/client/dispatch.rs +++ b/src/client/dispatch.rs @@ -196,13 +196,13 @@ impl Receiver { self.taker.cancel(); self.inner.close(); loop { - match crate::common::task::now_or_never(self.inner.recv()) { - Some(Some(mut env)) => return env.0.take(), - Some(None) => return None, + match self.inner.try_recv() { + Ok(mut env) => return env.0.take(), + Err(mpsc::error::TryRecvError::Disconnected) => return None, // A send can reserve capacity immediately before close() and // publish its envelope immediately afterwards. Once closed, - // Pending means that such a synchronous send is still in flight. - None => std::hint::spin_loop(), + // Empty means that such a synchronous send is still in flight. + Err(mpsc::error::TryRecvError::Empty) => std::hint::spin_loop(), } } } @@ -416,6 +416,33 @@ mod tests { } } + #[cfg(feature = "http1")] + #[tokio::test] + async fn receiver_shutdown_ignores_exhausted_task_budget() { + use futures_util::FutureExt; + + tokio::spawn(async { + let (mut tx, mut rx) = channel::(); + 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) = rx.close_and_recv().expect("queued request"); + assert_eq!(value, 43); + drop(callback); + assert!(rx.close_and_recv().is_none()); + }) + .await + .unwrap(); + } + #[cfg(feature = "http1")] #[test] fn receiver_shutdown_reclaims_concurrent_send() { diff --git a/src/common/task.rs b/src/common/task.rs index 7e6024557a..41671b1453 100644 --- a/src/common/task.rs +++ b/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 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(fut: F) -> Option { - 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, - } -} From eadd7e6146c15aa74a297c514b6e70befc682e53 Mon Sep 17 00:00:00 2001 From: Morax Date: Thu, 8 Oct 2026 05:36:14 +0200 Subject: [PATCH 3/4] fix(http1): yield while draining dispatch on shutdown --- src/client/dispatch.rs | 38 ++++++++------ src/proto/h1/dispatch.rs | 107 +++++++++++++++++++++++++++++---------- 2 files changed, 103 insertions(+), 42 deletions(-) diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs index 3d54fd5b63..cfe9cf7b21 100644 --- a/src/client/dispatch.rs +++ b/src/client/dispatch.rs @@ -192,19 +192,18 @@ impl Receiver { } #[cfg(feature = "http1")] - pub(crate) fn close_and_recv(&mut self) -> Option<(T, Callback)> { + pub(crate) fn poll_close_and_recv( + &mut self, + cx: &mut Context<'_>, + ) -> Poll)>> { self.taker.cancel(); self.inner.close(); - loop { - match self.inner.try_recv() { - Ok(mut env) => return env.0.take(), - Err(mpsc::error::TryRecvError::Disconnected) => return None, - // A send can reserve capacity immediately before close() and - // publish its envelope immediately afterwards. Once closed, - // Empty means that such a synchronous send is still in flight. - Err(mpsc::error::TryRecvError::Empty) => std::hint::spin_loop(), - } - } + // 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")) + }) } } @@ -418,7 +417,7 @@ mod tests { #[cfg(feature = "http1")] #[tokio::test] - async fn receiver_shutdown_ignores_exhausted_task_budget() { + async fn receiver_shutdown_yields_when_task_budget_is_exhausted() { use futures_util::FutureExt; tokio::spawn(async { @@ -434,10 +433,14 @@ mod tests { while budget_rx.recv().now_or_never().is_some() {} assert!(budget_rx.try_recv().is_ok()); - let (value, callback) = rx.close_and_recv().expect("queued request"); + 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!(rx.close_and_recv().is_none()); + assert!(futures_util::future::poll_fn(|cx| rx.poll_close_and_recv(cx)) + .await + .is_none()); }) .await .unwrap(); @@ -446,6 +449,9 @@ mod tests { #[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)>(); @@ -466,7 +472,9 @@ mod tests { let dropped = Arc::new(AtomicBool::new(false)); work_tx.send((tx, TrackDrop(dropped.clone()))).unwrap(); barrier.wait(); - drop(rx.close_and_recv()); + 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(); diff --git a/src/proto/h1/dispatch.rs b/src/proto/h1/dispatch.rs index 905859a22c..8b55c8c520 100644 --- a/src/proto/h1/dispatch.rs +++ b/src/proto/h1/dispatch.rs @@ -25,6 +25,7 @@ pub(crate) struct Dispatcher { body_tx: SenderDropGuard, body_rx: Pin>>, is_closing: bool, + error: Option, } pub(crate) trait Dispatch { @@ -38,6 +39,13 @@ pub(crate) trait Dispatch { ) -> Poll>>; 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, + ) -> Poll> { + Poll::Ready(self.recv_msg(Err(error.take().expect("pending connection error")))) + } fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll>; 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> { - 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,19 +733,6 @@ cfg_client! { message: None, })); Ok(()) - } else if !self.rx_closed { - if let Some((req, cb)) = self.rx.close_and_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) } @@ -747,6 +740,29 @@ cfg_client! { } } + fn poll_recv_error( + &mut self, + cx: &mut Context<'_>, + error: &mut Option, + ) -> Poll> { + 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> { match &mut self.callback { Some(cb) => match cb.poll_canceled(cx) { @@ -808,6 +824,43 @@ 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() { From a0f2a0a7d018d138e6f57c9fba2e9934e3774992 Mon Sep 17 00:00:00 2001 From: Morax Date: Thu, 8 Oct 2026 05:39:13 +0200 Subject: [PATCH 4/4] style(http1): format dispatch modules with repository checks --- src/client/dispatch.rs | 14 ++++++++------ src/proto/h1/dispatch.rs | 5 ++++- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs index cfe9cf7b21..0e189dc2db 100644 --- a/src/client/dispatch.rs +++ b/src/client/dispatch.rs @@ -201,9 +201,9 @@ impl Receiver { // 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")) - }) + self.inner + .poll_recv(cx) + .map(|item| item.map(|mut env| env.0.take().expect("envelope not dropped"))) } } @@ -438,9 +438,11 @@ mod tests { .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()); + assert!( + futures_util::future::poll_fn(|cx| rx.poll_close_and_recv(cx)) + .await + .is_none() + ); }) .await .unwrap(); diff --git a/src/proto/h1/dispatch.rs b/src/proto/h1/dispatch.rs index 8b55c8c520..0c7462938f 100644 --- a/src/proto/h1/dispatch.rs +++ b/src/proto/h1/dispatch.rs @@ -853,7 +853,10 @@ mod tests { assert!(dispatcher.error.is_some()); dispatcher.await.unwrap(); - let err = response.await.unwrap().expect_err("request must be canceled"); + let err = response + .await + .unwrap() + .expect_err("request must be canceled"); assert!(err.error.is_canceled()); assert!(err.message.is_some()); })