diff --git a/src/proto/h2/upgrade.rs b/src/proto/h2/upgrade.rs index 217aa27e88..def1fcfc96 100644 --- a/src/proto/h2/upgrade.rs +++ b/src/proto/h2/upgrade.rs @@ -1,11 +1,8 @@ use std::future::Future; use std::io::Cursor; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; use std::task::{Context, Poll}; -use atomic_waker::AtomicWaker; use bytes::{Buf, Bytes}; use futures_channel::{mpsc, oneshot}; use futures_core::{ready, Stream}; @@ -23,15 +20,10 @@ pub(super) fn pair( ) -> (H2Upgraded, UpgradedSendStreamTask) { let (tx, rx) = mpsc::channel(1); let (error_tx, error_rx) = oneshot::channel(); - let close_notify = Arc::new(UpgradedCloseNotify::new()); ( H2Upgraded { - send_stream: UpgradedSendStreamBridge { - tx, - error_rx, - close_notify: close_notify.clone(), - }, + send_stream: UpgradedSendStreamBridge { tx, error_rx }, recv_stream, ping, buf: Bytes::new(), @@ -39,7 +31,7 @@ pub(super) fn pair( UpgradedSendStreamTask { h2_tx: send_stream, rx, - close_notify, + buffered: None, error_tx: Some(error_tx), }, ) @@ -55,46 +47,6 @@ pub(super) struct H2Upgraded { struct UpgradedSendStreamBridge { tx: mpsc::Sender>>, error_rx: oneshot::Receiver, - close_notify: Arc, -} - -impl Drop for UpgradedSendStreamBridge { - fn drop(&mut self) { - self.close_notify.close(); - } -} - -struct UpgradedCloseNotify { - closed: AtomicBool, - task: AtomicWaker, -} - -impl UpgradedCloseNotify { - fn new() -> Self { - Self { - closed: AtomicBool::new(false), - task: AtomicWaker::new(), - } - } - - fn close(&self) { - self.closed.store(true, Ordering::Release); - self.task.wake(); - } - - fn poll_closed(&self, cx: &mut Context<'_>) -> Poll<()> { - if self.closed.load(Ordering::Acquire) { - return Poll::Ready(()); - } - - self.task.register(cx.waker()); - - if self.closed.load(Ordering::Acquire) { - Poll::Ready(()) - } else { - Poll::Pending - } - } } pin_project! { @@ -104,7 +56,7 @@ pin_project! { h2_tx: SendStream>, #[pin] rx: mpsc::Receiver>>, - close_notify: Arc, + buffered: Option>>, error_tx: Option>, } } @@ -123,35 +75,6 @@ where // one of the sides hanging up, so the task doesn't live around // longer than it's meant to. loop { - // we don't have the next chunk of data yet, so just reserve 1 byte to make - // sure there's some capacity available. h2 will handle the capacity management - // for the actual body chunk. - me.h2_tx.reserve_capacity(1); - - let h2_has_capacity = if me.h2_tx.capacity() == 0 { - // poll_capacity oddly needs a loop - loop { - match me.h2_tx.poll_capacity(cx) { - Poll::Ready(Some(Ok(0))) => {} - Poll::Ready(Some(Ok(_))) => break true, - Poll::Ready(Some(Err(e))) => { - return Poll::Ready(Err(crate::Error::new_body_write(e))) - } - Poll::Ready(None) => { - // None means the stream is no longer in a - // streaming state, we either finished it - // somehow, or the remote reset us. - return Poll::Ready(Err(crate::Error::new_body_write( - "send stream capacity unexpectedly closed", - ))); - } - Poll::Pending => break false, - } - } - } else { - true - }; - match me.h2_tx.poll_reset(cx) { Poll::Ready(Ok(reason)) => { trace!("stream received RST_STREAM: {:?}", reason); @@ -165,30 +88,43 @@ where Poll::Pending => (), } - // If h2 has no capacity, don't pull another item from the mpsc - // receiver. That would free a channel slot and let the writer - // enqueue more data without h2 backpressure. - // - // Still allow the task to finish once the upgraded write side is - // gone and the mpsc queue is empty. - if !h2_has_capacity { - // `size_hint` reads the queued message count without popping, - // so an accepted write stays queued until h2 capacity returns. - if me.rx.size_hint().0 == 0 && me.close_notify.poll_closed(cx).is_ready() { - me.h2_tx - .send_data(SendBuf::None, true) - .map_err(crate::Error::new_body_write)?; - return Poll::Ready(Ok(())); + // A write taken from the mpsc receiver waits here for h2 capacity, + // and the next one isn't pulled until it has been handed to h2, so + // the writer still sees h2 backpressure. + if me.buffered.is_some() { + // poll_capacity oddly needs a loop + while me.h2_tx.capacity() == 0 { + match ready!(me.h2_tx.poll_capacity(cx)) { + Some(Ok(0)) => {} + Some(Ok(_)) => break, + Some(Err(e)) => return Poll::Ready(Err(crate::Error::new_body_write(e))), + None => { + // None means the stream is no longer in a + // streaming state, we either finished it + // somehow, or the remote reset us. + return Poll::Ready(Err(crate::Error::new_body_write( + "send stream capacity unexpectedly closed", + ))); + } + } } - return Poll::Pending; + let cursor = me.buffered.take().expect("checked is_some above"); + me.h2_tx + .send_data(SendBuf::Cursor(cursor), false) + .map_err(crate::Error::new_body_write)?; + continue; } match me.rx.as_mut().poll_next(cx) { Poll::Ready(Some(cursor)) => { - me.h2_tx - .send_data(SendBuf::Cursor(cursor), false) - .map_err(crate::Error::new_body_write)?; + // Only reserve capacity once there is something to send. + // Reserving while idle, even a single byte, pins that + // capacity on the connection-level window (#4003). As in + // `PipeToSendStream`, h2 raises the request to the + // buffered length inside `send_data`. + me.h2_tx.reserve_capacity(1); + *me.buffered = Some(cursor); } Poll::Ready(None) => { me.h2_tx @@ -328,7 +264,6 @@ impl Write for H2Upgraded { cx: &mut Context<'_>, ) -> Poll> { self.send_stream.tx.close_channel(); - self.send_stream.close_notify.close(); match Pin::new(&mut self.send_stream.error_rx).poll(cx) { Poll::Ready(Ok(reason)) => Poll::Ready(Err(io_error(reason))), Poll::Ready(Err(_task_dropped)) => Poll::Ready(Ok(())), diff --git a/tests/client.rs b/tests/client.rs index b512260cc5..39fba04dab 100644 --- a/tests/client.rs +++ b/tests/client.rs @@ -3315,6 +3315,118 @@ mod conn { drop(tx_a); let _ = tokio::time::timeout(Duration::from_secs(5), a_handle).await; } + + // https://github.com/hyperium/hyper/issues/4003, for HTTP/2 CONNECT + // + // Like `h2_idle_stream_does_not_pin_connection_window`, but the idle + // stream is the send side of an `Upgraded` tunnel. It must not reserve + // connection-level flow control capacity while it has nothing to write. + #[tokio::test] + async fn h2_idle_upgraded_does_not_pin_connection_window() { + // One byte short of the initial connection-level window. + const STREAM_A_LEN: usize = 65534; + + let (client_io, server_io, _) = setup_duplex_test_server(); + let (stream_a_full_tx, stream_a_full_rx) = oneshot::channel::<()>(); + let (stream_a_done_tx, stream_a_done_rx) = oneshot::channel::<()>(); + let (stream_b_got_tx, stream_b_got_rx) = oneshot::channel::(); + + // Raw h2 server that never calls `release_capacity`, so it never + // sends a connection-level WINDOW_UPDATE. + tokio::spawn(async move { + let mut h2 = h2::server::handshake(server_io).await.unwrap(); + let mut stream_a_full_tx = Some(stream_a_full_tx); + let mut stream_a_done_rx = Some(stream_a_done_rx); + let mut stream_b_got_tx = Some(stream_b_got_tx); + while let Some(result) = h2.accept().await { + let (req, mut respond) = result.unwrap(); + if req.method() == Method::CONNECT { + let full_tx = stream_a_full_tx.take().unwrap(); + let done_rx = stream_a_done_rx.take().unwrap(); + tokio::spawn(async move { + let _send = respond.send_response(Response::new(()), false).unwrap(); + let mut body = req.into_body(); + let mut received = 0usize; + while received < STREAM_A_LEN { + match body.data().await { + Some(Ok(frame)) => received += frame.len(), + _ => return, + } + } + let _ = full_tx.send(()); + // Hold on to the recv stream, dropping it would release + // its capacity and send a WINDOW_UPDATE. + let _ = done_rx.await; + drop(body); + }); + } else { + let got_tx = stream_b_got_tx.take().unwrap(); + tokio::spawn(async move { + let mut body = req.into_body(); + let mut received = 0usize; + if let Some(Ok(frame)) = body.data().await { + received += frame.len(); + } + let _ = got_tx.send(received); + let mut send = respond.send_response(Response::new(()), false).unwrap(); + let _ = send.send_data(Bytes::from_static(b"ok"), true); + }); + } + } + }); + + let io = TokioIo::new(client_io); + let (mut client, conn) = conn::http2::Builder::new(TokioExecutor) + .handshake::<_, Full>(io) + .await + .expect("http handshake"); + tokio::spawn(async move { + let _ = conn.await; + }); + + // Stream A: a tunnel that writes STREAM_A_LEN bytes and then stays + // open, leaving one byte of connection window. + let req_a = Request::connect("localhost") + .body(Full::new(Bytes::new())) + .unwrap(); + let res_a = client.send_request(req_a).await.expect("send_request A"); + assert_eq!(res_a.status(), StatusCode::OK); + let mut upgraded = TokioIo::new(hyper::upgrade::on(res_a).await.unwrap()); + upgraded + .write_all(&vec![b'A'; STREAM_A_LEN]) + .await + .expect("write to tunnel"); + + tokio::time::timeout(Duration::from_secs(5), stream_a_full_rx) + .await + .expect("server should receive all of stream A in time") + .expect("stream_a_full_rx"); + + // Let the tunnel's send task park waiting for more writes. + for _ in 0..16 { + tokio::task::yield_now().await; + } + + // Stream B: one byte of body, which needs the last byte of window. + let req_b = Request::post("http://localhost/b") + .body(Full::new(Bytes::from_static(b"b"))) + .unwrap(); + let b_fut = client.send_request(req_b); + + let received_b = tokio::time::timeout(Duration::from_secs(5), stream_b_got_rx) + .await + .expect("stream B must reach the server even while the tunnel is idle") + .expect("stream_b_got_rx"); + assert_eq!( + received_b, 1, + "stream B should deliver its single body byte" + ); + + let _ = tokio::time::timeout(Duration::from_secs(5), b_fut).await; + + let _ = stream_a_done_tx.send(()); + drop(upgraded); + } } trait FutureHyperExt: TryFuture {