From 2880dab8f81249138b66528294df09eff45a935b Mon Sep 17 00:00:00 2001
From: Dan Gohman <dev@sunfishcode.online>
Date: Mon, 22 Mar 2021 11:23:45 -0700
Subject: [PATCH] Add a `sleep` function to the `WasiSched` trait.

---
 .../cap-std-sync/src/sched/unix.rs            |  3 +++
 .../cap-std-sync/src/sched/windows.rs         |  3 +++
 crates/wasi-common/src/sched.rs               |  1 +
 crates/wasi-common/src/snapshots/preview_0.rs | 24 +++++++++++++++++++
 crates/wasi-common/src/snapshots/preview_1.rs |  3 +--
 5 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/crates/wasi-common/cap-std-sync/src/sched/unix.rs b/crates/wasi-common/cap-std-sync/src/sched/unix.rs
index 15bbb29138..ab8be60a0a 100644
--- a/crates/wasi-common/cap-std-sync/src/sched/unix.rs
+++ b/crates/wasi-common/cap-std-sync/src/sched/unix.rs
@@ -108,6 +108,9 @@ impl WasiSched for SyncSched {
         std::thread::yield_now();
         Ok(())
     }
+    fn sleep(&self, duration: Duration) {
+        std::thread::sleep(duration)
+    }
 }
 
 fn wasi_file_raw_fd(f: &dyn WasiFile) -> Option<RawFd> {
diff --git a/crates/wasi-common/cap-std-sync/src/sched/windows.rs b/crates/wasi-common/cap-std-sync/src/sched/windows.rs
index 03eb41514c..93545c4040 100644
--- a/crates/wasi-common/cap-std-sync/src/sched/windows.rs
+++ b/crates/wasi-common/cap-std-sync/src/sched/windows.rs
@@ -128,6 +128,9 @@ impl WasiSched for SyncSched {
         thread::yield_now();
         Ok(())
     }
+    fn sleep(&self, duration: Duration) {
+        std::thread::sleep(duration)
+    }
 }
 
 fn wasi_file_raw_handle(f: &dyn WasiFile) -> Option<RawHandle> {
diff --git a/crates/wasi-common/src/sched.rs b/crates/wasi-common/src/sched.rs
index a925cd4e94..707459b4a6 100644
--- a/crates/wasi-common/src/sched.rs
+++ b/crates/wasi-common/src/sched.rs
@@ -10,6 +10,7 @@ use subscription::{MonotonicClockSubscription, RwSubscription, Subscription, Sub
 pub trait WasiSched {
     fn poll_oneoff(&self, poll: &Poll) -> Result<(), Error>;
     fn sched_yield(&self) -> Result<(), Error>;
+    fn sleep(&self, duration: Duration);
 }
 
 pub struct Userdata(u64);
diff --git a/crates/wasi-common/src/snapshots/preview_0.rs b/crates/wasi-common/src/snapshots/preview_0.rs
index 65f4823a3e..d027997122 100644
--- a/crates/wasi-common/src/snapshots/preview_0.rs
+++ b/crates/wasi-common/src/snapshots/preview_0.rs
@@ -738,6 +738,30 @@ impl<'a> wasi_unstable::WasiUnstable for WasiCtx {
             return Err(Error::invalid_argument().context("nsubscriptions must be nonzero"));
         }
 
+        // Special-case a `poll_oneoff` which is just sleeping on a single
+        // relative timer event, such as what WASI libc uses to implement sleep
+        // functions. This supports all clock IDs, because POSIX says that
+        // `clock_settime` doesn't effect relative sleeps.
+        if nsubscriptions == 1 {
+            let sub = subs.read()?;
+            if let types::SubscriptionU::Clock(clocksub) = sub.u {
+                if !clocksub
+                    .flags
+                    .contains(types::Subclockflags::SUBSCRIPTION_CLOCK_ABSTIME)
+                {
+                    self.sched.sleep(Duration::from_nanos(clocksub.timeout));
+
+                    events.write(types::Event {
+                        userdata: sub.userdata,
+                        error: types::Errno::Success,
+                        type_: types::Eventtype::Clock,
+                        fd_readwrite: fd_readwrite_empty(),
+                    })?;
+                    return Ok(1);
+                }
+            }
+        }
+
         let table = self.table();
         let mut poll = Poll::new();
 
diff --git a/crates/wasi-common/src/snapshots/preview_1.rs b/crates/wasi-common/src/snapshots/preview_1.rs
index 70d422c947..ca91b84cc7 100644
--- a/crates/wasi-common/src/snapshots/preview_1.rs
+++ b/crates/wasi-common/src/snapshots/preview_1.rs
@@ -16,7 +16,6 @@ use std::cell::{Ref, RefMut};
 use std::convert::{TryFrom, TryInto};
 use std::io::{IoSlice, IoSliceMut};
 use std::ops::{Deref, DerefMut};
-use std::thread;
 use tracing::debug;
 use wiggle::GuestPtr;
 
@@ -921,7 +920,7 @@ impl<'a> wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiCtx {
                     .flags
                     .contains(types::Subclockflags::SUBSCRIPTION_CLOCK_ABSTIME)
                 {
-                    thread::sleep(Duration::from_nanos(clocksub.timeout));
+                    self.sched.sleep(Duration::from_nanos(clocksub.timeout));
 
                     events.write(types::Event {
                         userdata: sub.userdata,