Repository: bytecodealliance/wasmtime
Pull request #8968 (2024-07-17 09:41:21)

Immediately return StreamError::Closed when the TCP stream is closed.
This saves the caller from having to perform one last `poll`+`read` sequence.

Before this change, it would first return an empty list of bytes, which is the wasi-io equivalent of EWOULDBLOCK.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->


alexcrichton: Thanks! Would it be possible to add a test for this too?
badeend: Yes!

The test includes a `std::thread::sleep`. Under Linux this shouldn't be necessary as communication over loopback devices happens synchronously. But I've thrown it in anyway just in case Windows or MacOS works differently.

Patch:
diff --git a/crates/wasi/src/tcp.rs b/crates/wasi/src/tcp.rs
index a050b97f1..d755a316d 100644
--- a/crates/wasi/src/tcp.rs
+++ b/crates/wasi/src/tcp.rs
@@ -694,7 +694,7 @@ impl HostInputStream for TcpReadStream {
             // A 0-byte read indicates that the stream has closed.
             Ok(0) => {
                 self.closed = true;
-                0
+                return Err(StreamError::Closed);
             }
             Ok(n) => n,
 


Test patch:
diff --git a/crates/test-programs/src/bin/preview2_tcp_streams.rs b/crates/test-programs/src/bin/preview2_tcp_streams.rs
new file mode 100644
index 000000000..a9a01b5ec
--- /dev/null
+++ b/crates/test-programs/src/bin/preview2_tcp_streams.rs
@@ -0,0 +1,48 @@
+use std::time::Duration;
+
+use test_programs::wasi::io::streams::StreamError;
+use test_programs::wasi::sockets::network::{IpAddress, IpAddressFamily, IpSocketAddress, Network};
+use test_programs::wasi::sockets::tcp::{ShutdownType, TcpSocket};
+
+/// InputStream::read should return `StreamError::Closed` after the connection has been shut down by the server.
+fn test_tcp_read_from_closed_input_stream(net: &Network, family: IpAddressFamily) {
+    // Set up server & client sockets:
+    let bind_address = IpSocketAddress::new(IpAddress::new_loopback(family), 0);
+    let listener = TcpSocket::new(family).unwrap();
+    listener.blocking_bind(&net, bind_address).unwrap();
+    listener.blocking_listen().unwrap();
+    let bound_address = listener.local_address().unwrap();
+    let client = TcpSocket::new(family).unwrap();
+    let (connected_input, connected_output) = client.blocking_connect(net, bound_address).unwrap();
+    let (accepted, accepted_input, accepted_output) = listener.blocking_accept().unwrap();
+
+    // Shut down the connection from the server side and give the kernel a bit
+    // of time to propagate the shutdown signal from the server socket to the
+    // client socket.
+    accepted.shutdown(ShutdownType::Both).unwrap();
+    drop(accepted_input);
+    drop(accepted_output);
+    drop(accepted);
+    std::thread::sleep(Duration::from_millis(50));
+
+    // And now the actual test:
+
+    // The input stream should immediately signal StreamError::Closed.
+    // Notably, it should _not_ return an empty list (the wasi-io equivalent of EWOULDBLOCK)
+    // See: https://github.com/bytecodealliance/wasmtime/pull/8968
+    assert!(matches!(connected_input.read(10), Err(StreamError::Closed))); // If this randomly fails, try tweaking the timeout above.
+
+    // Stream should still be closed, even when requesting 0 bytes:
+    assert!(matches!(connected_input.read(0), Err(StreamError::Closed)));
+
+    drop(connected_input);
+    drop(connected_output);
+    drop(client);
+}
+
+fn main() {
+    let net = Network::default();
+
+    test_tcp_read_from_closed_input_stream(&net, IpAddressFamily::Ipv4);
+    test_tcp_read_from_closed_input_stream(&net, IpAddressFamily::Ipv6);
+}
diff --git a/crates/wasi/tests/all/async_.rs b/crates/wasi/tests/all/async_.rs
index 84073e31d..7d2542372 100644
--- a/crates/wasi/tests/all/async_.rs
+++ b/crates/wasi/tests/all/async_.rs
@@ -324,6 +324,10 @@ async fn preview2_tcp_states() {
     run(PREVIEW2_TCP_STATES_COMPONENT, false).await.unwrap()
 }
 #[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn preview2_tcp_streams() {
+    run(PREVIEW2_TCP_STREAMS_COMPONENT, false).await.unwrap()
+}
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
 async fn preview2_tcp_bind() {
     run(PREVIEW2_TCP_BIND_COMPONENT, false).await.unwrap()
 }
diff --git a/crates/wasi/tests/all/sync.rs b/crates/wasi/tests/all/sync.rs
index aa02d3be7..5629d5d3f 100644
--- a/crates/wasi/tests/all/sync.rs
+++ b/crates/wasi/tests/all/sync.rs
@@ -270,6 +270,10 @@ fn preview2_tcp_states() {
     run(PREVIEW2_TCP_STATES_COMPONENT, false).unwrap()
 }
 #[test_log::test]
+fn preview2_tcp_streams() {
+    run(PREVIEW2_TCP_STREAMS_COMPONENT, false).unwrap()
+}
+#[test_log::test]
 fn preview2_tcp_bind() {
     run(PREVIEW2_TCP_BIND_COMPONENT, false).unwrap()
 }