Change around synchronization primitives

Instead of the lock around the reactor:
- use dashmaps on the future tasks
- use mutex for the poll tasks
- use integrated tests for testing functionality

Author: StaticallyTypedAnxiety

Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "wasm-runtime"
+name = "tiny-wasm-runtime"
 version = "0.1.0"
 edition = "2021"
 

src/engine.rs

@@ -2,7 +2,6 @@ use crate::Timer;
 use crate::{bindings::wasi::io::poll::Pollable, poll_tasks::PollTasks};
 use crate::{io::timer::TIMERS, poll_tasks::EventWithWaker};
 use dashmap::DashMap;
-use futures::pin_mut;
 use lazy_static::lazy_static;
 use std::{
     future::Future,
@@ -13,9 +12,10 @@ use std::{
     },
     task::{Context, Wake, Waker},
 };
+use uuid::Uuid;
 lazy_static! {
     /// The global reactor for this runtime
-    pub static ref REACTOR: Mutex<Reactor<'static>> = Mutex::new(Reactor::default());
+    pub static ref REACTOR: Reactor<'static> = Reactor::default();
 
 }
 /// The async engine instance
@@ -35,19 +35,19 @@ impl<'a> Task<'a> {
 }
 #[derive(Default)]
 pub struct Reactor<'a> {
-    events: PollTasks,
-    future_tasks: Vec<Task<'a>>, //right now the engine holds the tasks but depending
+    events: Mutex<PollTasks>,
+    future_tasks: DashMap<Uuid, Mutex<Task<'a>>>, //right now the engine holds the tasks but depending
     timers: DashMap<String, EventWithWaker<Timer>>,
 }
 
 impl<'a> Reactor<'a> {
     //adds event to the queue
-    pub fn register(&mut self, event_name: String, pollable: EventWithWaker<Arc<Pollable>>) {
-        self.events.push(event_name, pollable);
+    pub fn register(&self, event_name: String, pollable: EventWithWaker<Arc<Pollable>>) {
+        self.events.lock().unwrap().push(event_name, pollable);
     }
     //checks if descriptor has been added to the polling queue
     pub fn is_pollable(&self, key: &str) -> bool {
-        self.events.contains(key)
+        self.events.lock().unwrap().contains(key)
     }
 
     //checks if timer is pollable
@@ -56,17 +56,17 @@ impl<'a> Reactor<'a> {
     }
 
     //polls event queue to see if any of the events are readycar
-    pub fn wait_for_io(&mut self) {
-        self.events.wait_for_pollables();
+    pub fn wait_for_io(&self) {
+        self.events.lock().unwrap().wait_for_pollables();
     }
 
     //checks if event is ready
-    pub fn check_ready(&mut self, event_name: &str) -> bool {
-        self.events.check_if_ready(event_name)
+    pub fn check_ready(&self, event_name: &str) -> bool {
+        self.events.lock().unwrap().check_if_ready(event_name)
     }
 
     pub fn is_empty(&self) -> bool {
-        self.events.is_empty() && self.future_tasks.is_empty()
+        self.events.lock().unwrap().is_empty() && self.future_tasks.is_empty()
     }
 
     pub(crate) fn update_timers(&self) {
@@ -85,47 +85,58 @@ impl<'a> Reactor<'a> {
             .unwrap_or_default()
     }
 
-    pub(crate) fn register_timer(&mut self, timer_name: String, event: EventWithWaker<Timer>) {
+    pub(crate) fn register_timer(&self, timer_name: String, event: EventWithWaker<Timer>) {
         self.timers.insert(timer_name, event);
     }
-}
 
-impl WasmRuntimeAsyncEngine {
-    /// function to execute futures
-    pub fn block_on<K, F: Future<Output = K> + Send, Fun: FnOnce() -> F>(async_closure: Fun) {
-        let future = async_closure();
-        pin_mut!(future);
+    pub(crate) fn push_task<K, F: Future<Output = K> + Send + 'static>(&self, future: F) -> Uuid {
         let task = Task::new(Box::pin(async move {
             let _result = future.await;
         }));
-        let mut future_tasks = Vec::new();
-        future_tasks.push(task);
+        let id = Uuid::new_v4();
+        self.future_tasks.insert(id, Mutex::new(task));
+        id
+    }
+}
+
+impl WasmRuntimeAsyncEngine {
+    /// function to execute futures
+    pub fn block_on<K, F: Future<Output = K> + Send + 'static>(future: F) {
+        let reactor = &REACTOR;
+        reactor.push_task(future);
         loop {
-            let mut reactor = REACTOR.lock().unwrap();
             reactor.update_timers();
             reactor.wait_for_io();
-            reactor.future_tasks.retain_mut(|task_info| {
-                if task_info.waker.should_wake() {
-                    task_info.waker.reset();
-                    let waker: Waker = task_info.waker.clone().into();
+
+            reactor.future_tasks.retain(|_, task_info| {
+                let waker = task_info.get_mut().unwrap().waker.clone();
+                if waker.should_wake() {
+                    waker.reset();
+                    let waker: Waker = waker.into();
                     let mut context = Context::from_waker(&waker);
-                    if task_info.task.as_mut().poll(&mut context).is_ready() {
+
+                    if task_info
+                        .get_mut()
+                        .unwrap()
+                        .task
+                        .as_mut()
+                        .poll(&mut context)
+                        .is_ready()
+                    {
                         return false;
                     }
                 }
                 true
             });
+
             if TIMERS.is_empty() && reactor.is_empty() {
                 break;
             }
         }
     }
 
     pub fn spawn<K, F: Future<Output = ()> + Send + 'static>(future: F) {
-        let task = Task::new(Box::pin(async move {
-            let _result = future.await;
-        }));
-        REACTOR.lock().unwrap().future_tasks.push(task);
+        REACTOR.push_task(future);
     }
 }
 
@@ -188,14 +199,15 @@ mod test {
     #[test]
     fn test_enqueue() {
         let count_future = CountFuture { max: 3, min: 0 };
-        let mut reactor = Reactor::default();
-        reactor.future_tasks.push(Task::new(Box::pin(async move {
+        let reactor = Reactor::default();
+        let id = reactor.push_task(async move {
             count_future.await;
-        })));
-        let task = reactor.future_tasks.first_mut().unwrap();
-        let fut_waker = task.waker.clone();
+        });
+        let mut future_task = reactor.future_tasks.get_mut(&id).unwrap();
+        let task = future_task.value_mut();
+        let fut_waker = task.lock().unwrap().waker.clone();
         let waker: Waker = fut_waker.into();
-        let count_future = &mut task.task;
+        let count_future = &mut task.lock().unwrap().task;
         let mut context = Context::from_waker(&waker);
         futures::pin_mut!(count_future);
         let _ = count_future.as_mut().poll(&mut context);
@@ -205,6 +217,6 @@ mod test {
     fn test_block_on() {
         let count_future = CountFuture { max: 3, min: 0 };
 
-        WasmRuntimeAsyncEngine::block_on(|| async move { assert_eq!(count_future.await, 3) });
+        WasmRuntimeAsyncEngine::block_on(async move { assert_eq!(count_future.await, 3) });
     }
 }

src/io/net.rs

@@ -157,16 +157,16 @@ impl<'a> Future for ConnectionFuture<'a> {
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Self::Output> {
         let this = self.get_mut();
-        if !REACTOR.lock().unwrap().is_pollable(&this.async_key) {
+        if !REACTOR.is_pollable(&this.async_key) {
             this.stream.start_connect(this.address, this.port)?;
-            REACTOR.lock().unwrap().register(
+            REACTOR.register(
                 this.async_key.clone(),
                 (this.stream.pollable.clone(), cx.waker().clone()),
             );
         }
 
         //A PLACE TO CHECK IF THE REACTOR UPDATED THIS KEY
-        if REACTOR.lock().unwrap().check_ready(&this.async_key) {
+        if REACTOR.check_ready(&this.async_key) {
             this.stream.finish_connecting()?;
             Poll::Ready(Ok(()))
         } else {

src/io/timer.rs

@@ -34,7 +34,6 @@ impl Timer {
     }
 
     pub async fn timeout<K, F: Future<Output = K>>(
-        &self,
         fut: F,
         deadline: std::time::Duration,
     ) -> std::io::Result<K> {
@@ -75,7 +74,7 @@ impl Future for TimeFuture {
         cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<Self::Output> {
         let this = self.get_mut();
-        let mut reactor = REACTOR.lock().unwrap();
+        let reactor = &REACTOR;
         if reactor.is_timer_pollable(&this.timer_key) {
             let has_elapsed = reactor.timer_has_elapsed(&this.timer_key);
             if has_elapsed {

tests/engine.rs

@@ -0,0 +1,45 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tiny_wasm_runtime::{Timer, WasmRuntimeAsyncEngine};
+
+pub async fn test_timers_with_assertions() {
+    let task_a_done: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
+    let task_b_done = Arc::new(AtomicBool::new(false));
+
+    let task_a_done_clone = task_a_done.clone();
+    let task_b_done_clone = task_b_done.clone();
+    println!("spaw goes here");
+    WasmRuntimeAsyncEngine::spawn::<(), _>(async move {
+        println!("sleep happens here");
+        Timer::sleep(std::time::Duration::from_secs(1)).await;
+        task_a_done_clone.store(true, Ordering::SeqCst);
+    });
+
+    println!("spaw goes here 2");
+
+    WasmRuntimeAsyncEngine::spawn::<(), _>(async move {
+        Timer::sleep(std::time::Duration::from_millis(500)).await;
+        task_b_done_clone.store(true, Ordering::SeqCst);
+    });
+
+    // Do main timeout
+    let slow_future = async {
+        Timer::sleep(std::time::Duration::from_secs(2)).await;
+        "completed"
+    };
+
+    let res = Timer::timeout(slow_future, std::time::Duration::from_secs(1)).await;
+
+    assert!(res.is_err(), "Expected the main future to time out");
+    assert!(
+        task_a_done.load(Ordering::SeqCst) || task_b_done.load(Ordering::SeqCst),
+        "At least one background task should have completed by now"
+    );
+}
+
+#[test]
+fn test_full_engine_runtime() {
+    WasmRuntimeAsyncEngine::block_on(async {
+        test_timers_with_assertions().await;
+    });
+}