diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs index 76a9c3e20c..143fc5a9cb 100644 --- a/src/client/dispatch.rs +++ b/src/client/dispatch.rs @@ -30,14 +30,16 @@ pub struct TrySendError { pub(crate) fn channel() -> (Sender, Receiver) { let (tx, rx) = mpsc::unbounded_channel(); - let (giver, taker) = want::new(); - let tx = Sender { - #[cfg(feature = "http1")] - buffered_once: false, - giver, - inner: tx, + let (giver, mut taker) = want::new(); + // Seed with one want so the first request can be sent before the + // receiver is polled; every request must consume a want + taker.want(); + let tx = Sender { giver, inner: tx }; + let rx = Receiver { + inner: rx, + taker, + wanting: true, }; - let rx = Receiver { inner: rx, taker }; (tx, rx) } @@ -46,17 +48,11 @@ pub(crate) fn channel() -> (Sender, Receiver) { /// While the inner sender is unbounded, the Giver is used to determine /// if the Receiver is ready for another request. pub(crate) struct Sender { - /// One message is always allowed, even if the Receiver hasn't asked - /// for it yet. This boolean keeps track of whether we've sent one - /// without notice. - #[cfg(feature = "http1")] - buffered_once: bool, - /// The Giver helps watch that the Receiver side has been polled - /// when the queue is empty. This helps us know when a request and - /// response have been fully processed, and a connection is ready - /// for more. + // Controls whether the sender is permitted to dispatch a request + // Beyond the initial request, the Receiver only grants readiness once + // the preceding request/response cycle has fully completed giver: want::Giver, - /// Actually bounded by the Giver, plus `buffered_once`. + // Capacity is enforced via the Giver inner: mpsc::UnboundedSender>, } @@ -89,23 +85,9 @@ impl Sender { self.giver.is_canceled() } - #[cfg(feature = "http1")] - fn can_send(&mut self) -> bool { - if self.giver.give() || !self.buffered_once { - // If the receiver is ready *now*, then of course we can send. - // - // If the receiver isn't ready yet, but we don't have anything - // in the channel yet, then allow one message. - self.buffered_once = true; - true - } else { - false - } - } - #[cfg(feature = "http1")] pub(crate) fn try_send(&mut self, val: T) -> Result, T> { - if !self.can_send() { + if !self.giver.give() { return Err(val); } let (tx, rx) = oneshot::channel(); @@ -117,7 +99,7 @@ impl Sender { #[cfg(feature = "http1")] pub(crate) fn send(&mut self, val: T) -> Result, T> { - if !self.can_send() { + if !self.giver.give() { return Err(val); } let (tx, rx) = oneshot::channel(); @@ -176,24 +158,25 @@ impl Clone for UnboundedSender { pub(crate) struct Receiver { inner: mpsc::UnboundedReceiver>, taker: want::Taker, + // Tracks whether readiness (`want`) has already been granted for the next request + wanting: bool, } impl Receiver { pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll)>> { + if !self.wanting { + self.taker.want(); + self.wanting = true; + } match self.inner.poll_recv(cx) { Poll::Ready(item) => { - // A want signaled after finding the queue empty can land after - // the Sender's `give()` for the message taken here, as can one - // signaled on a coop-budget Pending with a message queued. - // Withdraw it, or the Sender reports this connection as ready - // while it is still serving this message. - self.taker.unwant(); + self.wanting = false; Poll::Ready(item.map(|mut env| env.0.take().expect("envelope not dropped"))) } - Poll::Pending => { - self.taker.want(); - Poll::Pending - } + // Avoid re-signaling `want` here: the sender might have already claimed it + // and enqueued a request, even if polling the channel returns `Pending` + // due to Tokio's coop budget exhaustion or a concurrent send + Poll::Pending => Poll::Pending, } } diff --git a/tests/h1_ready_in_flight.rs b/tests/h1_ready_in_flight.rs index f2317ba308..19edc104cc 100644 --- a/tests/h1_ready_in_flight.rs +++ b/tests/h1_ready_in_flight.rs @@ -1,16 +1,15 @@ // Test: an HTTP/1 client connection must not report ready while a request is // in flight on it. // -// The dispatch `Receiver` signals want when its queue reports `Pending`, and -// `SendRequest::is_ready` reports that want. The signal can go stale: the -// connection task can find its queue empty and a request land before it -// signals, or tokio's coop budget can make the queue report `Pending` with a -// request already in it. The request is then taken with want still set, so -// `is_ready` stays true until the response is complete. Pools that return a -// connection on `is_ready`, such as hyper-util's legacy client, then give that -// connection to the next request, which waits behind the whole response. +// Re-granting `want` on every `Pending` poll from the dispatch channel can +// falsely mark the connection as ready while a request is already in flight +// (eg, during a concurrent send or when Tokio's coop budget runs out with +// a buffered request) +// Revoking readiness when the request is dequeued is too late, as the caller +// may have already seen `Ready` and will fail on send // -// The coop path is deterministic, so the test drives that one. +// Deterministically trigger coop exhaustion and verify readiness state both +// before and after dequeuing the buffered request use std::future::Future; use std::io; @@ -73,14 +72,20 @@ async fn h1_connection_is_not_ready_while_a_request_is_in_flight() { let _response = sender.send_request(Request::new(Empty::new())); assert!(!sender.is_ready()); - // Out of coop budget, the queue reports Pending with the request in it, so - // the connection signals ready again. A connection task preempted between - // finding its queue empty and signaling leaves the same stale signal. + // With the coop budget exhausted, polling the queue yields `Pending` even + // though a request is buffered + // Readiness must not be re-signaled here poll_fn(|cx| { while let Poll::Ready(restore) = tokio::task::coop::poll_proceed(cx) { restore.made_progress(); } - poll_conn(&mut conn, cx) + let _ = poll_conn(&mut conn, cx); + assert!( + sender.poll_ready(cx).is_pending(), + "dispatcher should not report ready while a request is still buffered" + ); + assert!(!sender.is_ready()); + Poll::Ready(()) }) .await; tokio::task::yield_now().await;