fixes

Author: lthiery

Cargo.toml

@@ -242,8 +242,8 @@ path = "tests/integration.rs"
 required-features = ["full"]
 
 [[test]]
-name = "ready_stream"
-path = "tests/ready_stream.rs"
+name = "ready_on_poll_stream"
+path = "tests/ready_on_poll_stream.rs"
 required-features = ["full"]
 
 [[test]]

src/proto/h1/dispatch.rs

@@ -170,14 +170,15 @@ where
         // benchmarks often use. Perhaps it should be a config option instead.
         for _ in 0..16 {
             let _ = self.poll_read(cx)?;
-            let _ = self.poll_write(cx)?;
-            let conn_ready = self.poll_flush(cx)?.is_ready();
+            let write_ready = self.poll_write(cx)?.is_ready();
+            let flush_ready = self.poll_flush(cx)?.is_ready();
 
             // If we can write more body and the connection is ready, we should
             // write again. If we return `Ready(Ok(())` here, we will yield
-            // without a guaranteed wakeup from the write side of the connection.
+            // without a guaranteed wake-up from the write side of the connection.
             // This would lead to a deadlock if we also don't expect reads.
-            let wants_write_again = self.can_write_again() && conn_ready;
+            let wants_write_again = self.can_write_again() && (write_ready || flush_ready);
+
             // This could happen if reading paused before blocking on IO,
             // such as getting to the end of a framed message, but then
             // writing/flushing set the state back to Init. In that case,
@@ -187,16 +188,30 @@ where
             // Using this instead of task::current() and notify() inside
             // the Conn is noticeably faster in pipelined benchmarks.
             let wants_read_again = self.conn.wants_read_again();
+
             // If we cannot write or read again, we yield and rely on the
-            // wakeup from the connection futures.
+            // wake-up from the connection futures.
             if !(wants_write_again || wants_read_again) {
-                //break;
                 return Poll::Ready(Ok(()));
             }
-        }
 
+            // If we are continuing only because "wants_write_again", check if write is ready.
+            if !wants_read_again && wants_write_again {
+                // If write was ready, just proceed with the loop
+                if write_ready {
+                    continue;
+                }
+                // Write was previously pending, but may have become ready since polling flush, so
+                // we need to check it again. If we simply proceeded, the case of an unbuffered
+                // writer where flush is always ready would cause us to hot loop.
+                if self.poll_write(cx)?.is_pending() {
+                    // write is pending, so it is safe to yield and rely on wake-up from connection
+                    // futures.
+                    return Poll::Ready(Ok(()));
+                }
+            }
+        }
         trace!("poll_loop yielding (self = {:p})", self);
-
         task::yield_now(cx).map(|never| match never {})
     }
 

tests/h1_server/mod.rs

@@ -0,0 +1,5 @@
+mod streams;
+mod test_framework;
+
+pub use streams::StreamReadHalf;
+pub use test_framework::{init_tracing, run_body_test, TestConfig, TestStream};

tests/h1_server/streams.rs

@@ -0,0 +1,87 @@
+use hyper::rt::{Read, ReadBufCursor};
+use pin_project_lite::pin_project;
+use std::io;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+use tokio::sync::mpsc;
+
+// Common read half shared by both stream types
+pin_project! {
+    #[derive(Debug)]
+    pub struct StreamReadHalf {
+        #[pin]
+        read_rx: mpsc::UnboundedReceiver<Vec<u8>>,
+        read_buffer: Vec<u8>,
+    }
+}
+
+impl StreamReadHalf {
+    pub fn new(read_rx: mpsc::UnboundedReceiver<Vec<u8>>) -> Self {
+        Self {
+            read_rx,
+            read_buffer: Vec::new(),
+        }
+    }
+
+    /// Receive data written to this stream by the other end (async)
+    pub async fn recv(&mut self) -> Option<Vec<u8>> {
+        self.read_rx.recv().await
+    }
+}
+
+impl Read for StreamReadHalf {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: ReadBufCursor<'_>,
+    ) -> Poll<io::Result<()>> {
+        let mut this = self.as_mut().project();
+
+        // First, try to satisfy the read request from the internal buffer
+        if !this.read_buffer.is_empty() {
+            let to_read = std::cmp::min(this.read_buffer.len(), buf.remaining());
+            // Copy data from internal buffer to the read buffer
+            buf.put_slice(&this.read_buffer[..to_read]);
+            // Remove the consumed data from the internal buffer
+            this.read_buffer.drain(..to_read);
+            return Poll::Ready(Ok(()));
+        }
+
+        // If internal buffer is empty, try to get data from the channel
+        match this.read_rx.as_mut().get_mut().try_recv() {
+            Ok(data) => {
+                // Copy as much data as we can fit in the buffer
+                let to_read = std::cmp::min(data.len(), buf.remaining());
+                buf.put_slice(&data[..to_read]);
+
+                // Store any remaining data in the internal buffer for next time
+                if to_read < data.len() {
+                    let remaining = &data[to_read..];
+                    this.read_buffer.extend_from_slice(remaining);
+                }
+                Poll::Ready(Ok(()))
+            }
+            Err(mpsc::error::TryRecvError::Empty) => {
+                match ready!(this.read_rx.poll_recv(cx)) {
+                    Some(data) => {
+                        // Copy as much data as we can fit in the buffer
+                        let to_read = std::cmp::min(data.len(), buf.remaining());
+                        buf.put_slice(&data[..to_read]);
+
+                        // Store any remaining data in the internal buffer for next time
+                        if to_read < data.len() {
+                            let remaining = &data[to_read..];
+                            this.read_buffer.extend_from_slice(remaining);
+                        }
+                        Poll::Ready(Ok(()))
+                    }
+                    None => Poll::Ready(Ok(())),
+                }
+            }
+            Err(mpsc::error::TryRecvError::Disconnected) => {
+                // Channel closed, return EOF
+                Poll::Ready(Ok(()))
+            }
+        }
+    }
+}

tests/h1_server/test_framework.rs

@@ -0,0 +1,123 @@
+use http_body_util::StreamBody;
+use hyper::body::Bytes;
+use hyper::body::Frame;
+use hyper::server::conn::http1;
+use hyper::service::service_fn;
+use hyper::{Response, StatusCode};
+use std::convert::Infallible;
+use tracing::{error, info};
+
+pub struct TestConfig {
+    pub total_chunks: usize,
+    pub chunk_size: usize,
+}
+
+impl Default for TestConfig {
+    fn default() -> Self {
+        Self {
+            total_chunks: 16,
+            chunk_size: 64 * 1024,
+        }
+    }
+}
+
+pub fn init_tracing() {
+    use std::sync::Once;
+    static INIT: Once = Once::new();
+    INIT.call_once(|| {
+        tracing_subscriber::fmt()
+            .with_max_level(tracing::Level::INFO)
+            .with_target(true)
+            .with_thread_ids(true)
+            .with_thread_names(true)
+            .init();
+    });
+}
+
+// Trait for streams that can send and receive data directly
+pub trait TestStream: hyper::rt::Read + hyper::rt::Write + Send + Unpin + 'static {
+    fn send(&self, data: &[u8]) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
+    fn recv(&mut self) -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<Vec<u8>>> + Send + '_>>;
+}
+
+pub async fn run_body_test<S>(stream_pair: (S, S), config: TestConfig)
+where
+    S: TestStream,
+{
+    let (server_stream, mut client_stream) = stream_pair;
+
+    let mut http_builder = http1::Builder::new();
+    http_builder.max_buf_size(config.chunk_size);
+    
+    let total_chunks = config.total_chunks;
+    let chunk_size = config.chunk_size;
+    
+    let service = service_fn(move |_| {
+        let total_chunks = total_chunks;
+        let chunk_size = chunk_size;
+        async move {
+            info!(
+                "Creating payload of {} chunks of {} KiB each ({} MiB total)...",
+                total_chunks,
+                chunk_size / 1024,
+                total_chunks * chunk_size / (1024 * 1024)
+            );
+            let bytes = Bytes::from(vec![0; chunk_size]);
+            let data = vec![bytes.clone(); total_chunks];
+            let stream = futures_util::stream::iter(
+                data.into_iter()
+                    .map(|b| Ok::<_, Infallible>(Frame::data(b))),
+            );
+            let body = StreamBody::new(stream);
+            info!("Server: Sending data response...");
+            Ok::<_, hyper::Error>(
+                Response::builder()
+                    .status(StatusCode::OK)
+                    .header("content-type", "application/octet-stream")
+                    .header("content-length", (total_chunks * chunk_size).to_string())
+                    .body(body)
+                    .unwrap(),
+            )
+        }
+    });
+
+    let server_task = tokio::spawn(async move {
+        let conn = http_builder.serve_connection(Box::pin(server_stream), service);
+        let conn_result = conn.await;
+        if let Err(e) = &conn_result {
+            error!("Server connection error: {}", e);
+        }
+        conn_result
+    });
+
+    let get_request = "GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n";
+    client_stream.send(get_request.as_bytes())
+        .map_err(|e| Box::new(std::io::Error::new(std::io::ErrorKind::Other, format!("Failed to send request: {}", e))))
+        .unwrap();
+    
+    info!("Client is reading response...");
+    let mut bytes_received = 0;
+    let mut all_data = Vec::new();
+    while let Some(chunk) = client_stream.recv().await {
+        bytes_received += chunk.len();
+        all_data.extend_from_slice(&chunk);
+    }
+    
+    // Clean up
+    let result = server_task.await.unwrap();
+    result.unwrap();
+    
+    // Parse HTTP response to find body start
+    // HTTP response format: "HTTP/1.1 200 OK\r\n...headers...\r\n\r\n<body>"
+    let body_start = all_data.windows(4)
+        .position(|w| w == b"\r\n\r\n")
+        .map(|pos| pos + 4)
+        .unwrap_or(0);
+    
+    let body_bytes = bytes_received - body_start;
+    assert_eq!(body_bytes, config.total_chunks * config.chunk_size, 
+               "Expected {} body bytes, got {} (total received: {}, headers: {})", 
+               config.total_chunks * config.chunk_size, body_bytes, bytes_received, body_start);
+    info!(bytes_received, body_bytes, "Client done receiving bytes");
+}
+