fix(http1): flush bytes buffered by the write re-check before yielding (#4143)

`poll_loop`'s main path always calls `poll_flush` after `poll_write`. The
"wants_write_again" re-check added in #3988 calls `poll_write` a second time
and returns straight out of the loop when it pends, skipping that flush.

That second write can buffer bytes before it pends. When a response body
reaches end-of-stream between the two write polls, `end_body()` buffers the
end of the message and the write then pends on the *next* message
(`poll_msg`). Returning there strands the terminating chunk in the write
buffer: the wake-ups the connection is left waiting on are for reads, so
nothing flushes it. The peer receives the body but never the terminator and
waits until it gives up, at which point the connection reports
`IncompleteMessage` from `mid_message_detect_eof`.

Observed on a server streaming a chunked body fed from another thread, at
roughly one connection in 600k. hyper's own trace shows the divergence:

    healthy:  buf.len=24, buf.len=5, flushed 29 bytes
    stalled:  buf.len=24, flushed 24 bytes, buf.len=5, <nothing>

Flush what the re-check buffered before yielding. Guard the flush on there
being buffered bytes so the call pattern is otherwise unchanged.

Add a test that drives the interleaving deterministically: a body that yields
one data frame, then pends, then ends the stream on the very next poll, all
within a single `poll_loop` iteration.

Author: ricochet

src/proto/h1/conn.rs

@@ -592,6 +592,11 @@ where
         self.io.can_buffer()
     }
 
+    /// Whether bytes are sitting in the write buffer waiting to be flushed.
+    pub(crate) fn has_buffered_write(&self) -> bool {
+        self.io.has_buffered_write()
+    }
+
     pub(crate) fn write_head(&mut self, head: MessageHead<T::Outgoing>, body: Option<BodyLength>) {
         if let Some(encoder) = self.encode_head(head, body) {
             self.state.writing = if !encoder.is_eof() {

src/proto/h1/dispatch.rs

@@ -205,6 +205,15 @@ where
                 // we need to check it again. If it is still pending, it is safe to yield and rely
                 // on wake-up from the connection futures.
                 if self.poll_write(cx)?.is_pending() {
+                    // That write can have buffered bytes before going pending: a body that
+                    // reached end-of-stream between the two write polls buffers the end of the
+                    // message here, and then the write goes pending on the *next* message.
+                    // Yielding without flushing would strand those bytes in the write buffer
+                    // until the peer gives up, since the wake-ups we then rely on are for
+                    // reads. Flush what was just buffered before yielding.
+                    if self.conn.has_buffered_write() {
+                        let _ = self.poll_flush(cx)?;
+                    }
                     return Poll::Ready(Ok(()));
                 }
             }

src/proto/h1/io.rs

@@ -149,6 +149,11 @@ where
         self.write_buf.buffer(buf);
     }
 
+    /// Whether there are bytes waiting in the write buffer to be flushed.
+    pub(crate) fn has_buffered_write(&self) -> bool {
+        self.write_buf.remaining() > 0
+    }
+
     pub(crate) fn can_buffer(&self) -> bool {
         self.flush_pipeline || self.write_buf.can_buffer()
     }

tests/h1_flush_before_yield.rs

@@ -0,0 +1,118 @@
+// Test: `poll_loop` must never yield leaving buffered bytes unflushed.
+//
+// `poll_loop`'s main path always calls `poll_flush` after `poll_write`. The
+// "wants_write_again" re-check added a *second* `poll_write` whose `Pending`
+// returned straight out of the loop, skipping that flush.
+//
+// That second write can buffer bytes before it pends. A response body that
+// reaches end-of-stream between the two write polls has its end-of-message
+// buffered by `end_body()`, and the write then pends on the *next* message
+// (`poll_msg`). Returning there strands the terminating chunk in the write
+// buffer: the wake-ups the connection then waits on are for reads, so nothing
+// flushes it and the peer waits for a response that is already written but
+// never sent.
+//
+// The body below drives exactly that interleaving deterministically: one data
+// frame, then `Pending`, then end-of-stream on the very next poll, both of
+// which happen inside a single `poll_loop` iteration.
+
+use std::convert::Infallible;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::time::Duration;
+
+use bytes::Bytes;
+use hyper::body::{Body, Frame};
+use hyper::server::conn::http1;
+use hyper::service::service_fn;
+use hyper::Response;
+use support::TokioIo;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::{TcpListener, TcpStream};
+use tokio::time::timeout;
+
+mod support;
+
+/// Yields one data frame, then pends once, then ends the stream.
+///
+/// The `Pending` deliberately arranges no wake-up: it stands in for a body
+/// whose readiness changes between hyper's two write polls of the same
+/// `poll_loop` iteration, which is what leaves the end of the message buffered
+/// by the re-check write.
+#[derive(Default)]
+struct PendOnceThenEnd {
+    polls: u8,
+}
+
+impl Body for PendOnceThenEnd {
+    type Data = Bytes;
+    type Error = Infallible;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        self.polls += 1;
+        match self.polls {
+            1 => Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"hello"))))),
+            2 => Poll::Pending,
+            _ => Poll::Ready(None),
+        }
+    }
+}
+
+#[tokio::test]
+async fn h1_server_flushes_end_of_body_buffered_by_write_recheck() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    tokio::spawn(async move {
+        let (socket, _) = listener.accept().await.unwrap();
+        let service = service_fn(|_req| async {
+            Ok::<_, Infallible>(Response::new(PendOnceThenEnd::default()))
+        });
+        let _ = http1::Builder::new()
+            .serve_connection(TokioIo::new(socket), service)
+            .await;
+    });
+
+    let mut client = TcpStream::connect(addr).await.unwrap();
+    client
+        .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
+        .await
+        .unwrap();
+
+    // The whole response must arrive, terminating chunk included. Before the
+    // fix the body's chunk arrives but `0\r\n\r\n` never does, so this read
+    // waits forever.
+    let mut received = Vec::new();
+    let read = timeout(Duration::from_secs(5), async {
+        let mut buf = [0u8; 256];
+        loop {
+            let n = client.read(&mut buf).await.unwrap();
+            if n == 0 {
+                break;
+            }
+            received.extend_from_slice(&buf[..n]);
+            if received.ends_with(b"\r\n0\r\n\r\n") {
+                break;
+            }
+        }
+    })
+    .await;
+
+    assert!(
+        read.is_ok(),
+        "response never completed; got {:?}",
+        String::from_utf8_lossy(&received)
+    );
+    let received = String::from_utf8_lossy(&received);
+    assert!(
+        received.ends_with("\r\n0\r\n\r\n"),
+        "missing terminating chunk; got {received:?}"
+    );
+    assert!(
+        received.contains("hello"),
+        "missing body content; got {received:?}"
+    );
+}