Fix lost wakeups with stdin and wasip3

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

@@ -0,0 +1,25 @@
+use test_programs::p3::wasi;
+
+struct Component;
+
+test_programs::p3::export!(Component);
+
+impl test_programs::p3::exports::wasi::cli::run::Guest for Component {
+    async fn run() -> Result<(), ()> {
+        let (mut stream, result) = wasi::cli::stdin::read_via_stream();
+        let (sresult, buf) = stream.read(Vec::with_capacity(100)).await;
+        assert_eq!(buf, b"hello!".to_vec());
+        assert_eq!(sresult, wit_bindgen::StreamResult::Complete(6));
+
+        let (sresult, buf) = stream.read(Vec::with_capacity(100)).await;
+        assert!(buf.is_empty());
+        assert_eq!(sresult, wit_bindgen::StreamResult::Dropped);
+
+        result.await.unwrap();
+        Ok(())
+    }
+}
+
+fn main() {
+    unreachable!();
+}

crates/wasi/src/cli/worker_thread_stdin.rs

@@ -121,7 +121,8 @@ fn create() -> GlobalStdin {
                 *state.state.lock().unwrap(),
                 StdinState::ReadRequested
             ));
-            *state.state.lock().unwrap() = new_state;
+            let mut lock = state.state.lock().unwrap();
+            *lock = new_state;
             state.read_completed.notify_waiters();
             if done {
                 break;
@@ -205,6 +206,14 @@ impl AsyncRead for WasiStdinAsyncRead {
     ) -> Poll<io::Result<()>> {
         let g = GlobalStdin::get();
 
+        // Everything below is executed under the global stdin lock. It's not
+        // going to block below so that's semantically fine. Optimization-wise
+        // it's probably possible to move this within the loop around just a
+        // small part of reading/writing the state, but that was done
+        // historically and it resulted in lost wakeups with `Notify`, so this
+        // is conservatively hoisted up here.
+        let mut locked = g.state.lock().unwrap();
+
         // Perform everything below in a `loop` to handle the case that a read
         // was stolen by another thread, for example, or perhaps a spurious
         // notification to `Notified`.
@@ -222,7 +231,6 @@ impl AsyncRead for WasiStdinAsyncRead {
 
             // Once we're in the "ready" state then take a look at the global
             // state of stdin.
-            let mut locked = g.state.lock().unwrap();
             match mem::replace(&mut *locked, StdinState::ReadRequested) {
                 // If data is available then drain what we can into `buf`.
                 StdinState::Data(mut data) => {
@@ -260,11 +268,6 @@ impl AsyncRead for WasiStdinAsyncRead {
             }
 
             self.set(WasiStdinAsyncRead::Waiting(g.read_completed.notified()));
-
-            // Intentionally drop the lock after the `notified()` future
-            // creation just above as to work correctly this needs to happen
-            // within the lock.
-            drop(locked);
         }
     }
 }

tests/all/cli_tests.rs

@@ -2780,6 +2780,27 @@ start a print 1234
 
         Ok(())
     }
+
+    #[test]
+    fn p3_cli_read_stdin() -> Result<()> {
+        let mut cmd = get_wasmtime_command()?;
+        let mut child = cmd
+            .arg("-Sp3")
+            .arg("-Wcomponent-model-async")
+            .arg(P3_CLI_READ_STDIN_COMPONENT)
+            .stdin(std::process::Stdio::piped())
+            .stdout(std::process::Stdio::piped())
+            .stderr(std::process::Stdio::piped())
+            .spawn()
+            .unwrap();
+        child.stdin.take().unwrap().write_all(b"hello!").unwrap();
+        let output = child.wait_with_output()?;
+        println!("stdout: {}", String::from_utf8_lossy(&output.stdout));
+        println!("stderr: {}", String::from_utf8_lossy(&output.stderr));
+        assert!(output.status.success());
+
+        Ok(())
+    }
 }
 
 #[test]