From 9b0927293675d29854051538aa10987e76b4b3f4 Mon Sep 17 00:00:00 2001
From: Pat Hickey <pat@moreproductive.org>
Date: Fri, 7 May 2021 12:19:51 -0700
Subject: [PATCH] poll_oneoff: bound tests for time, rather than instant
 completion

---
 crates/wasi-common/tokio/tests/poll_oneoff.rs | 87 ++++++++++++++++++-
 1 file changed, 85 insertions(+), 2 deletions(-)

diff --git a/crates/wasi-common/tokio/tests/poll_oneoff.rs b/crates/wasi-common/tokio/tests/poll_oneoff.rs
index 10ac82b9be..818f7e86c0 100644
--- a/crates/wasi-common/tokio/tests/poll_oneoff.rs
+++ b/crates/wasi-common/tokio/tests/poll_oneoff.rs
@@ -1,13 +1,17 @@
 use anyhow::{Context, Error};
+use cap_std::time::Duration;
+use std::collections::HashMap;
 use wasi_common::{
     file::{FdFlags, OFlags},
     sched::{Poll, RwEventFlags, SubscriptionResult, Userdata},
-    WasiDir,
+    WasiDir, WasiFile,
 };
-use wasi_tokio::{sched::poll_oneoff, Dir};
+use wasi_tokio::{clocks_ctx, sched::poll_oneoff, Dir};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn empty_file_readable() -> Result<(), Error> {
+    let clocks = clocks_ctx();
+
     let workspace = unsafe { cap_tempfile::tempdir().expect("create tempdir") };
     workspace.create_dir("d").context("create dir")?;
     let d = workspace.open_dir("d").context("open dir")?;
@@ -30,6 +34,17 @@ async fn empty_file_readable() -> Result<(), Error> {
 
     let mut poll = Poll::new();
     poll.subscribe_read(&mut *f, Userdata::from(123));
+    // Timeout bounds time in poll_oneoff
+    poll.subscribe_monotonic_clock(
+        &*clocks.monotonic,
+        clocks
+            .monotonic
+            .now(clocks.monotonic.resolution())
+            .checked_add(Duration::from_millis(1))
+            .unwrap(),
+        clocks.monotonic.resolution(),
+        Userdata::from(0),
+    );
     poll_oneoff(&mut poll).await?;
 
     let events = poll.results();
@@ -48,6 +63,8 @@ async fn empty_file_readable() -> Result<(), Error> {
 
 #[tokio::test(flavor = "multi_thread")]
 async fn empty_file_writable() -> Result<(), Error> {
+    let clocks = clocks_ctx();
+
     let workspace = unsafe { cap_tempfile::tempdir().expect("create tempdir") };
     workspace.create_dir("d").context("create dir")?;
     let d = workspace.open_dir("d").context("open dir")?;
@@ -60,6 +77,17 @@ async fn empty_file_writable() -> Result<(), Error> {
 
     let mut poll = Poll::new();
     poll.subscribe_write(&mut *writable_f, Userdata::from(123));
+    // Timeout bounds time in poll_oneoff
+    poll.subscribe_monotonic_clock(
+        &*clocks.monotonic,
+        clocks
+            .monotonic
+            .now(clocks.monotonic.resolution())
+            .checked_add(Duration::from_millis(1))
+            .unwrap(),
+        clocks.monotonic.resolution(),
+        Userdata::from(0),
+    );
     poll_oneoff(&mut poll).await?;
 
     let events = poll.results();
@@ -75,3 +103,58 @@ async fn empty_file_writable() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn stdio_readable() -> Result<(), Error> {
+    let clocks = clocks_ctx();
+
+    let deadline = clocks
+        .monotonic
+        .now(clocks.monotonic.resolution())
+        .checked_add(Duration::from_millis(1))
+        .unwrap();
+
+    let mut waiting_on: HashMap<u64, Box<dyn WasiFile>> = vec![
+        (
+            1,
+            Box::new(wasi_tokio::stdio::stdout()) as Box<dyn WasiFile>,
+        ),
+        (2, Box::new(wasi_tokio::stdio::stderr())),
+    ]
+    .into_iter()
+    .collect();
+
+    while !waiting_on.is_empty() {
+        let mut poll = Poll::new();
+
+        for (ix, file) in waiting_on.iter_mut() {
+            poll.subscribe_write(&mut **file, Userdata::from(*ix));
+        }
+        // Timeout bounds time in poll_oneoff
+        poll.subscribe_monotonic_clock(
+            &*clocks.monotonic,
+            deadline,
+            clocks.monotonic.resolution(),
+            Userdata::from(999),
+        );
+        poll_oneoff(&mut poll).await?;
+        let events = poll.results();
+
+        for e in events {
+            match e {
+                (SubscriptionResult::Write(Ok(_)), ud) => {
+                    let _ = waiting_on.remove(&u64::from(ud));
+                }
+                (SubscriptionResult::Write(Err(_)), ud) => {
+                    panic!("error on ix {}", u64::from(ud))
+                }
+                (SubscriptionResult::Read { .. }, _) => unreachable!(),
+                (SubscriptionResult::MonotonicClock { .. }, _) => {
+                    panic!("timed out before stdin and stdout ready for reading")
+                }
+            }
+        }
+    }
+
+    Ok(())
+}