From b82a1ff70ad638d71a4796779763d0683218a48c Mon Sep 17 00:00:00 2001 From: Surma Date: Wed, 11 Nov 2026 12:59:26 +0000 Subject: [PATCH] Stop bootstrap chat on stdin EOF during socket and HTTP waits --- host/tests/it/chat.rs | 25 ++++++++++++++ userland/chat/src/main.rs | 73 +++++++++++++++++++++++++++------------ 2 files changed, 75 insertions(+), 23 deletions(-) diff --git a/host/tests/it/chat.rs b/host/tests/it/chat.rs index 7dddf61..f1047fd 100644 --- a/host/tests/it/chat.rs +++ b/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(()) } diff --git a/userland/chat/src/main.rs b/userland/chat/src/main.rs index 7f6a552..03bce89 100644 --- a/userland/chat/src/main.rs +++ b/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::() .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)), } }