wasip3: Add a test cancelling TCP reads (bytecodealliance#11623)

* wasip3: Add a test cancelling TCP reads

This commit adds a test where a stream of data is being sent over TCP
and the read side reads as much as it can and then cancels any read that
shows up pending. It then waits for readiness via a 0-length read and
retries again. The test is intended to assert that the number of
consecutive zero-length reads reported is capped at some small number
(e.g. this shouldn't infinite loop).

* Fix lint

Author: alexcrichton

Cargo.lock

@@ -325,8 +325,8 @@ io-lifetimes = { version = "2.0.3", default-features = false }
 io-extras = "0.18.1"
 rustix = "1.0.3"
 # wit-bindgen:
-wit-bindgen = { version = "0.45.0", default-features = false }
-wit-bindgen-rust-macro = { version = "0.45.0", default-features = false }
+wit-bindgen = { version = "0.45.1", default-features = false }
+wit-bindgen-rust-macro = { version = "0.45.1", default-features = false }
 
 # wasm-tools family:
 wasmparser = { version = "0.238.1", default-features = false, features = ['simd'] }

Cargo.toml

@@ -1,8 +1,11 @@
 use futures::join;
+use std::pin::pin;
+use std::task::{Context, Poll, Waker};
 use test_programs::p3::wasi::sockets::types::{
     IpAddress, IpAddressFamily, IpSocketAddress, TcpSocket,
 };
 use test_programs::p3::wit_stream;
+use wit_bindgen::StreamResult;
 
 struct Component;
 
@@ -109,6 +112,76 @@ async fn test_tcp_shutdown_should_not_lose_data(family: IpAddressFamily) {
     .await;
 }
 
+/// Model a situation where there's a continuous stream of data coming into the
+/// guest from one side and the other side is reading in chunks but also
+/// cancelling reads occasionally. Should receive the complete stream of data
+/// into the result.
+async fn test_tcp_read_cancellation(family: IpAddressFamily) {
+    // Send 2M of data in 256-byte chunks.
+    const CHUNKS: usize = (2 << 20) / 256;
+    let mut data = [0; 256];
+    for (i, slot) in data.iter_mut().enumerate() {
+        *slot = i as u8;
+    }
+
+    setup(family, |server, client| async move {
+        // Minimize the local send buffer:
+        client.set_send_buffer_size(1024).unwrap();
+
+        let (mut client_tx, client_rx) = wit_stream::new();
+        join!(
+            async {
+                client.send(client_rx).await.unwrap();
+            },
+            async {
+                for _ in 0..CHUNKS {
+                    let ret = client_tx.write_all(data.to_vec()).await;
+                    assert!(ret.is_empty());
+                }
+                drop(client_tx);
+            },
+            async {
+                let mut buf = Vec::with_capacity(1024);
+                let (mut server_rx, server_fut) = server.receive();
+                let mut i = 0_usize;
+                let mut consecutive_zero_length_reads = 0;
+                loop {
+                    assert!(buf.is_empty());
+                    let (status, b) = {
+                        let mut fut = pin!(server_rx.read(buf));
+                        let mut cx = Context::from_waker(Waker::noop());
+                        match fut.as_mut().poll(&mut cx) {
+                            Poll::Ready(pair) => pair,
+                            Poll::Pending => fut.cancel(),
+                        }
+                    };
+                    buf = b;
+                    match status {
+                        StreamResult::Complete(n) => {
+                            assert_eq!(buf.len(), n);
+                            for slot in buf.iter_mut() {
+                                assert_eq!(*slot, i as u8);
+                                i = i.wrapping_add(1);
+                            }
+                            buf.truncate(0);
+                            consecutive_zero_length_reads = 0;
+                        }
+                        StreamResult::Dropped => break,
+                        StreamResult::Cancelled => {
+                            assert!(consecutive_zero_length_reads < 10);
+                            consecutive_zero_length_reads += 1;
+                            server_rx.read(Vec::new()).await;
+                        }
+                    }
+                }
+                assert_eq!(i, CHUNKS * 256);
+                server_fut.await.unwrap();
+            },
+        );
+    })
+    .await;
+}
+
 impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
     async fn run() -> Result<(), ()> {
         test_tcp_input_stream_should_be_closed_by_remote_shutdown(IpAddressFamily::Ipv4).await;
@@ -122,6 +195,8 @@ impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
 
         test_tcp_shutdown_should_not_lose_data(IpAddressFamily::Ipv4).await;
         test_tcp_shutdown_should_not_lose_data(IpAddressFamily::Ipv6).await;
+
+        test_tcp_read_cancellation(IpAddressFamily::Ipv4).await;
         Ok(())
     }
 }

crates/test-programs/src/bin/p3_sockets_tcp_streams.rs

@@ -4753,6 +4753,12 @@ criteria = "safe-to-deploy"
 delta = "0.43.0 -> 0.45.0"
 notes = "The Bytecode Alliance is the author of this crate"
 
+[[audits.wit-bindgen]]
+who = "Alex Crichton <[email protected]>"
+criteria = "safe-to-deploy"
+delta = "0.45.0 -> 0.45.1"
+notes = "The Bytecode Alliance is the author of this crate"
+
 [[audits.wit-bindgen-core]]
 who = "Joel Dice <[email protected]>"
 criteria = "safe-to-run"
@@ -4764,6 +4770,12 @@ criteria = "safe-to-deploy"
 delta = "0.43.0 -> 0.45.0"
 notes = "The Bytecode Alliance is the author of this crate"
 
+[[audits.wit-bindgen-core]]
+who = "Alex Crichton <[email protected]>"
+criteria = "safe-to-deploy"
+delta = "0.45.0 -> 0.45.1"
+notes = "The Bytecode Alliance is the author of this crate"
+
 [[audits.wit-bindgen-rt]]
 who = "Joel Dice <[email protected]>"
 criteria = "safe-to-run"
@@ -4780,6 +4792,12 @@ criteria = "safe-to-deploy"
 delta = "0.43.0 -> 0.45.0"
 notes = "The Bytecode Alliance is the author of this crate"
 
+[[audits.wit-bindgen-rust]]
+who = "Alex Crichton <[email protected]>"
+criteria = "safe-to-deploy"
+delta = "0.45.0 -> 0.45.1"
+notes = "The Bytecode Alliance is the author of this crate"
+
 [[audits.wit-bindgen-rust-macro]]
 who = "Joel Dice <[email protected]>"
 criteria = "safe-to-run"
@@ -4791,6 +4809,12 @@ criteria = "safe-to-deploy"
 delta = "0.43.0 -> 0.45.0"
 notes = "The Bytecode Alliance is the author of this crate"
 
+[[audits.wit-bindgen-rust-macro]]
+who = "Alex Crichton <[email protected]>"
+criteria = "safe-to-deploy"
+delta = "0.45.0 -> 0.45.1"
+notes = "The Bytecode Alliance is the author of this crate"
+
 [[audits.wit-component]]
 who = "Alex Crichton <[email protected]>"
 criteria = "safe-to-deploy"