Add a TryLock primitive

Author: alexcrichton

crates/guest-rust/src/rt/async_support.rs

@@ -68,6 +68,7 @@ mod futures_stream;
 mod inter_task_wakeup;
 mod stream_support;
 mod subtask;
+mod try_lock;
 #[cfg(feature = "inter-task-wakeup")]
 mod unit_stream;
 mod waitable;

crates/guest-rust/src/rt/async_support/future_support.rs

@@ -111,23 +111,20 @@
 //! alive until that write completes but otherwise shouldn't hinder anything
 //! else.
 
-use crate::rt::Cleanup;
-use crate::rt::async_support::ReturnCode;
+use crate::rt::async_support::try_lock::TryLock;
 use crate::rt::async_support::waitable::{WaitableOp, WaitableOperation};
+use crate::rt::async_support::ReturnCode;
+use crate::rt::Cleanup;
 use alloc::sync::Arc;
 use alloc::task::Wake;
 use core::alloc::Layout;
-use core::cell::UnsafeCell;
 use core::fmt;
 use core::future::{Future, IntoFuture};
 use core::marker;
 use core::mem::{self, ManuallyDrop};
 use core::pin::Pin;
 use core::ptr;
-use core::sync::atomic::{
-    AtomicBool, AtomicU32,
-    Ordering::{Acquire, Relaxed, Release},
-};
+use core::sync::atomic::{AtomicU32, Ordering::Relaxed};
 use core::task::{Context, Poll, Waker};
 
 /// Helper trait which encapsulates the various operations which can happen
@@ -506,8 +503,7 @@ impl<O: FutureOps> RawFutureWriter<O> {
         O: 'static,
     {
         return Arc::new(DeferredWrite {
-            write: UnsafeCell::new(self.write(value)),
-            in_use: AtomicBool::new(false),
+            write: TryLock::new(self.write(value)),
         })
         .wake();
 
@@ -525,8 +521,7 @@ impl<O: FutureOps> RawFutureWriter<O> {
         /// doesn't require the `async-spawn` feature and instead works with the
         /// `wasip3_task` C ABI structures (which spawn doesn't support).
         struct DeferredWrite<O: FutureOps> {
-            write: UnsafeCell<RawFutureWrite<O>>,
-            in_use: AtomicBool,
+            write: TryLock<RawFutureWrite<O>>,
         }
 
         // SAFETY: Needed to satisfy `Waker::from` but otherwise should be ok
@@ -536,10 +531,6 @@ impl<O: FutureOps> RawFutureWriter<O> {
 
         impl<O: FutureOps + 'static> Wake for DeferredWrite<O> {
             fn wake(self: Arc<Self>) {
-                if self.in_use.swap(true, Acquire) {
-                    panic!("recursive wake detected");
-                }
-
                 // When a `wake` signal comes in that should happen in two
                 // locations:
                 //
@@ -561,16 +552,15 @@ impl<O: FutureOps> RawFutureWriter<O> {
                 let poll = {
                     let waker = Waker::from(self.clone());
                     let mut cx = Context::from_waker(&waker);
-                    unsafe { Pin::new_unchecked(&mut *self.write.get()).poll(&mut cx) }
+                    let mut write = self.write.try_lock().unwrap();
+                    unsafe { Pin::new_unchecked(&mut *write).poll(&mut cx) }
                 };
                 if poll.is_ready() {
                     assert_eq!(Arc::strong_count(&self), 1);
                 } else {
                     assert!(Arc::strong_count(&self) > 1);
                 }
                 assert_eq!(Arc::weak_count(&self), 0);
-
-                self.in_use.store(false, Release);
             }
         }
     }

crates/guest-rust/src/rt/async_support/inter_task_wakeup.rs

@@ -1,12 +1,8 @@
 use super::FutureState;
+use crate::rt::async_support::try_lock::TryLock;
 use crate::rt::async_support::{BLOCKED, COMPLETED};
 use crate::{RawStreamReader, RawStreamWriter, StreamOps, UnitStreamOps};
-use core::cell::UnsafeCell;
 use core::ptr;
-use core::sync::atomic::{
-    AtomicBool,
-    Ordering::{Acquire, Release},
-};
 
 #[derive(Default)]
 pub struct State {
@@ -31,7 +27,9 @@ impl FutureState<'_> {
             assert!(!self.inter_task_wakeup.stream_reading);
             let (writer, reader) = UnitStreamOps::new();
             self.inter_task_wakeup.stream = Some(reader);
-            self.waker.inter_task_stream.set(writer);
+            let mut waker_stream = self.waker.inter_task_stream.lock.try_lock().unwrap();
+            assert!(waker_stream.is_none());
+            *waker_stream = Some(writer);
         }
 
         // If there's not already a pending read then schedule a new read here.
@@ -83,37 +81,18 @@ impl State {
 
 #[derive(Default)]
 pub struct WakerState {
-    in_use: AtomicBool,
-    lock: UnsafeCell<Option<RawStreamWriter<UnitStreamOps>>>,
+    lock: TryLock<Option<RawStreamWriter<UnitStreamOps>>>,
 }
 
-unsafe impl Send for WakerState {}
-unsafe impl Sync for WakerState {}
-
 impl WakerState {
-    fn set(&self, writer: RawStreamWriter<UnitStreamOps>) {
-        self.with(|waker_stream| {
-            assert!(waker_stream.is_none());
-            *waker_stream = Some(writer);
-        })
-    }
-
-    fn with<R>(&self, f: impl FnOnce(&mut Option<RawStreamWriter<UnitStreamOps>>) -> R) -> R {
-        assert!(!self.in_use.swap(true, Acquire));
-        let ret = unsafe { f(&mut *self.lock.get()) };
-        self.in_use.store(false, Release);
-        ret
-    }
-
     pub fn wake(&self) {
         // Here the wakeup stream should already have been filled in by the
         // original future itself. The stream should also have an active read
         // while the future is sleeping. This means that this write should
         // succeed immediately.
-        self.with(|inter_task_stream| {
-            let stream = inter_task_stream.as_mut().unwrap();
-            let rc = unsafe { UnitStreamOps.start_write(stream.handle(), ptr::null_mut(), 1) };
-            assert_eq!(rc, COMPLETED | (1 << 4));
-        })
+        let mut inter_task_stream = self.lock.try_lock().unwrap();
+        let stream = inter_task_stream.as_mut().unwrap();
+        let rc = unsafe { UnitStreamOps.start_write(stream.handle(), ptr::null_mut(), 1) };
+        assert_eq!(rc, COMPLETED | (1 << 4));
     }
 }

crates/guest-rust/src/rt/async_support/try_lock.rs

@@ -0,0 +1,85 @@
+use core::cell::UnsafeCell;
+use core::ops::{Deref, DerefMut};
+use core::sync::atomic::{
+    AtomicBool,
+    Ordering::{Acquire, Release},
+};
+
+/// Small helper type to wrap `T` in a lock-like primitive which only supports
+/// the `try_lock` operation.
+///
+/// This is useful on wasm right now where threads aren't actually a thing so
+/// there shouldn't ever be contention, but the Rust type system still requires
+/// Send/Sync.
+#[derive(Default)]
+pub struct TryLock<T> {
+    locked: AtomicBool,
+    data: UnsafeCell<T>,
+}
+
+unsafe impl<T: Send> Send for TryLock<T> {}
+unsafe impl<T: Send> Sync for TryLock<T> {}
+
+impl<T> TryLock<T> {
+    pub fn new(data: T) -> Self {
+        TryLock {
+            locked: AtomicBool::new(false),
+            data: UnsafeCell::new(data),
+        }
+    }
+
+    pub fn try_lock(&self) -> Option<TryLockGuard<'_, T>> {
+        if self.locked.swap(true, Acquire) {
+            None
+        } else {
+            Some(TryLockGuard { lock: self })
+        }
+    }
+}
+
+pub struct TryLockGuard<'a, T> {
+    lock: &'a TryLock<T>,
+}
+
+impl<T> Deref for TryLockGuard<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.lock.data.get() }
+    }
+}
+
+impl<T> DerefMut for TryLockGuard<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { &mut *self.lock.data.get() }
+    }
+}
+
+impl<T> Drop for TryLockGuard<'_, T> {
+    fn drop(&mut self) {
+        self.lock.locked.store(false, Release);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn smoke() {
+        let lock = TryLock::new(1);
+        assert!(lock.try_lock().is_some());
+        assert!(lock.try_lock().is_some());
+
+        let mut guard = lock.try_lock().unwrap();
+        assert_eq!(*guard, 1);
+        *guard = 2;
+        assert_eq!(*guard, 2);
+        assert!(lock.try_lock().is_none());
+        drop(guard);
+        assert!(lock.try_lock().is_some());
+
+        let guard = lock.try_lock().unwrap();
+        assert_eq!(*guard, 2);
+    }
+}