Merge v3-t82: chat exits on stdin EOF, host shutdown in ms (T82)

host/tests/it/chat.rs

@@ -33,6 +33,13 @@ async fn request(port: u16, path: &str, body: Option<&str>) -> anyhow::Result<(u
 }
 #[tokio::test(flavor = "current_thread")]
 async fn seeded_chat() -> anyhow::Result<()> {
+    seeded_chat_case(false).await
+}
+#[tokio::test(flavor = "current_thread")]
+async fn shutdown_with_open_connection() -> anyhow::Result<()> {
+    seeded_chat_case(true).await
+}
+async fn seeded_chat_case(open_connection: bool) -> anyhow::Result<()> {
     let source = Path::new(env!("CARGO_MANIFEST_DIR")).join("..");
     let target = crate::common::target_dir();
     crate::common::cwd_init();
@@ -254,7 +261,25 @@ async fn seeded_chat() -> anyhow::Result<()> {
         )
     })
     .await;
+    let connection = if open_connection {
+        let mut stream = tokio::net::TcpStream::connect(("127.0.0.1", port)).await?;
+        // Leave the request incomplete so the server must wait for more input.
+        stream
+            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n")
+            .await?;
+        Some(stream)
+    } else {
+        None
+    };
+    let started = std::time::Instant::now();
     services.shutdown().await;
+    let elapsed = started.elapsed();
+    drop(connection);
     result??;
+    println!("Chat shutdown (open_connection={open_connection}): {elapsed:?}");
+    anyhow::ensure!(
+        elapsed < Duration::from_secs(3),
+        "Chat delayed shutdown: {elapsed:?}"
+    );
     Ok(())
 }

userland/chat/src/main.rs

@@ -1,4 +1,5 @@
 //! A bootstrap HTTP chat server. Only the host harness API supplies chat data.
+//! Stdin EOF stops the server, including waits for connections and HTTP.
 use clap::Parser;
 use std::{io, net::SocketAddr};
 use wasi::{
@@ -6,7 +7,10 @@ use wasi::{
         outgoing_handler,
         types::{Fields, Method, OutgoingBody, OutgoingRequest, Scheme},
     },
-    io::streams::{InputStream, OutputStream},
+    io::{
+        poll::{Pollable, poll},
+        streams::{InputStream, OutputStream, StreamError},
+    },
     sockets::{
         instance_network::instance_network, network::*, tcp_create_socket::create_tcp_socket,
     },
@@ -19,13 +23,37 @@ struct Options {
 fn error(e: impl std::fmt::Debug) -> io::Error {
     io::Error::other(format!("{e:?}"))
 }
-fn write(stream: &OutputStream, bytes: &[u8]) -> io::Result<()> {
-    for chunk in bytes.chunks(4096) {
-        stream.blocking_write_and_flush(chunk).map_err(error)?;
+fn wait(stdin: &InputStream, event: &Pollable) -> io::Result<()> {
+    let stop = stdin.subscribe();
+    loop {
+        let ready = poll(&[&stop, event]);
+        match stdin.read(65536) {
+            Err(StreamError::Closed) => std::process::exit(0),
+            Err(e) => return Err(error(e)),
+            Ok(_) => {}
+        }
+        if ready.contains(&1) {
+            return Ok(());
+        }
+    }
+}
+fn write(stdin: &InputStream, stream: &OutputStream, mut bytes: &[u8]) -> io::Result<()> {
+    while !bytes.is_empty() {
+        wait(stdin, &stream.subscribe())?;
+        let count = (stream.check_write().map_err(error)? as usize).min(bytes.len());
+        stream.write(&bytes[..count]).map_err(error)?;
+        bytes = &bytes[count..];
+    }
+    stream.flush().map_err(error)?;
+    loop {
+        wait(stdin, &stream.subscribe())?;
+        if stream.check_write().map_err(error)? > 0 {
+            return Ok(());
+        }
     }
-    Ok(())
 }
 fn proxy(
+    stdin: &InputStream,
     method: &str,
     path: &str,
     content_type: Option<&str>,
@@ -61,16 +89,10 @@ fn proxy(
     let future = outgoing_handler::handle(request, None).map_err(error)?;
     if !body.is_empty() {
         let stream = outgoing.write().map_err(error)?;
-        let mut rest = body;
-        while !rest.is_empty() {
-            stream.subscribe().block();
-            let count = (stream.check_write().map_err(error)? as usize).min(rest.len());
-            stream.write(&rest[..count]).map_err(error)?;
-            rest = &rest[count..];
-        }
+        write(stdin, &stream, body)?;
     }
     OutgoingBody::finish(outgoing, None).map_err(error)?;
-    future.subscribe().block();
+    wait(stdin, &future.subscribe())?;
     let response = future
         .get()
         .ok_or_else(|| error("missing HTTP response"))?
@@ -87,7 +109,8 @@ fn proxy(
     let stream = incoming.stream().map_err(error)?;
     let mut bytes = Vec::new();
     loop {
-        match stream.blocking_read(65536) {
+        wait(stdin, &stream.subscribe())?;
+        match stream.read(65536) {
             Ok(chunk) => bytes.extend(chunk),
             Err(wasi::io::streams::StreamError::Closed) => break,
             Err(e) => return Err(error(e)),
@@ -95,13 +118,14 @@ fn proxy(
     }
     Ok((status, content_type, bytes))
 }
-fn serve(input: &InputStream, output: &OutputStream) -> io::Result<()> {
+fn serve(stdin: &InputStream, input: &InputStream, output: &OutputStream) -> io::Result<()> {
     let mut bytes = Vec::new();
     let end = loop {
         if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") {
             break index + 4;
         }
-        bytes.extend(input.blocking_read(65536).map_err(error)?);
+        wait(stdin, &input.subscribe())?;
+        bytes.extend(input.read(65536).map_err(error)?);
     };
     let headers = std::str::from_utf8(&bytes[..end])
         .map_err(error)?
@@ -120,7 +144,8 @@ fn serve(input: &InputStream, output: &OutputStream) -> io::Result<()> {
         .parse::<usize>()
         .map_err(error)?;
     while bytes.len() < end + length {
-        bytes.extend(input.blocking_read(65536).map_err(error)?);
+        wait(stdin, &input.subscribe())?;
+        bytes.extend(input.read(65536).map_err(error)?);
     }
     let (status, kind, body) = if method == "GET" && path == "/" {
         (
@@ -130,6 +155,7 @@ fn serve(input: &InputStream, output: &OutputStream) -> io::Result<()> {
         )
     } else if let Some(path) = path.strip_prefix("/api/") {
         proxy(
+            stdin,
             method,
             path,
             header("content-type"),
@@ -139,11 +165,12 @@ fn serve(input: &InputStream, output: &OutputStream) -> io::Result<()> {
     } else {
         (404, "text/plain".into(), b"Not found".to_vec())
     };
-    write(output,format!("HTTP/1.1 {status} Response\r\nContent-Type: {kind}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",body.len()).as_bytes())?;
-    write(output, &body)
+    write(stdin, output,format!("HTTP/1.1 {status} Response\r\nContent-Type: {kind}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",body.len()).as_bytes())?;
+    write(stdin, output, &body)
 }
 fn main() -> io::Result<()> {
     let address = Options::parse().listen;
+    let stdin = wasi::cli::stdin::get_stdin();
     let (family, address) = match address {
         SocketAddr::V4(a) => (
             IpAddressFamily::Ipv4,
@@ -166,17 +193,17 @@ fn main() -> io::Result<()> {
     socket
         .start_bind(&instance_network(), address)
         .map_err(error)?;
-    socket.subscribe().block();
+    wait(&stdin, &socket.subscribe())?;
     socket.finish_bind().map_err(error)?;
     socket.start_listen().map_err(error)?;
-    socket.subscribe().block();
+    wait(&stdin, &socket.subscribe())?;
     socket.finish_listen().map_err(error)?;
     loop {
         match socket.accept() {
             Ok((_connection, input, output)) => {
-                let _ = serve(&input, &output);
+                let _ = serve(&stdin, &input, &output);
             }
-            Err(ErrorCode::WouldBlock) => socket.subscribe().block(),
+            Err(ErrorCode::WouldBlock) => wait(&stdin, &socket.subscribe())?,
             Err(e) => return Err(error(e)),
         }
     }