feat: add support for suspending execution by timer, atomic flag or user callback

Author: WhaleKit

crates/tinywasm/src/coro.rs

@@ -4,18 +4,26 @@ use crate::Result;
 // use alloc::boxed::Box;
 pub(crate) use tinywasm_types::{ResumeArgument, YieldedValue};
 
+
+/// explains why did execution suspend, and carries payload if needed
 #[derive(Debug)]
 pub enum SuspendReason {
     /// host function yielded
     /// potentially some host functions might expect resume argument when calling resume
     Yield(YieldedValue),
-    // /// timer ran out (not implemented),
-    // /// host shouldn't provide resume argument when calling resume
-    // SuspendedEpoch,
 
-    // /// async should_suspend flag was set (not implemented)
-    // /// host shouldn't provide resume argument when calling resume
-    // SuspendedFlag,
+    /// time to suspend has come,
+    /// host shouldn't provide resume argument when calling resume
+    #[cfg(feature = "std")]
+    SuspendedEpoch,
+
+    /// user's should-suspend-callback,
+    /// host shouldn't provide resume argument when calling resume
+    SuspendedCallback,
+
+    /// async should_suspend flag was set
+    /// host shouldn't provide resume argument when calling resume
+    SuspendedFlag,
 }
 
 /// result of a function that might pause in the middle and yield

crates/tinywasm/src/func.rs

@@ -37,14 +37,20 @@ pub(self) enum SuspendFuncInner {
     Host(SuspendedHostCoroState),
 }
 
+/// handle to function that was suspended and can be resumed
 #[derive(Debug)]
 pub struct SuspendFunc {
-    pub(self) func: SuspendFuncInner,
-    pub(crate) module_addr: ModuleInstanceAddr,
+    func: SuspendFuncInner,
+    module_addr: ModuleInstanceAddr,
+    store_id: usize,
 }
 
 impl<'a> crate::coro::CoroState<Vec<WasmValue>, &mut Store> for SuspendFunc {
     fn resume(&mut self, store: &mut Store, arg: ResumeArgument) -> Result<FuncHandleResumeOutcome> {
+        if store.id() != self.store_id {
+            return Err(Error::InvalidStore);
+        }
+
         let ctx = FuncContext { store, module_addr: self.module_addr };
         match &mut self.func {
             SuspendFuncInner::Wasm(wasm) => wasm.resume(ctx, arg),
@@ -110,6 +116,7 @@ impl FuncHandle {
                         coro_orig_function: self.addr,
                     }),
                     module_addr: self.module_addr,
+                    store_id: store.id(),
                 }));
             }
             Function::Wasm(wasm_func) => wasm_func,
@@ -145,6 +152,7 @@ impl FuncHandle {
                         result_types: func_ty.results.clone(),
                     }),
                     module_addr: self.module_addr,
+                    store_id: store.id(),
                 }
             }))
     }

crates/tinywasm/src/interpreter/executor.rs

@@ -78,8 +78,13 @@ impl<'store, 'stack> Executor<'store, 'stack> {
                     return Ok(ExecOutcome::Suspended(suspend_reason));
                 }
             };
-            self.after_resume_host_coro(&res);
+            self.stack.values.extend_from_wasmvalues(&res);
             self.suspended_host_coro = None;
+
+            // we don't know how much time we spent in host function
+            if let ControlFlow::Break(ReasonToBreak::Suspended(reason)) = self.check_should_suspend() {
+                return Ok(ExecOutcome::Suspended(reason));
+            }
         }
 
         loop {
@@ -93,9 +98,35 @@ impl<'store, 'stack> Executor<'store, 'stack> {
         }
     }
 
-    fn after_resume_host_coro(&mut self, vals: &[WasmValue]) {
-        self.stack.values.extend_from_wasmvalues(&vals);
-        self.cf.incr_instr_ptr();
+    /// for controlling how long execution spends in wasm
+    /// called when execution loops back, because that might happen indefinite amount of times
+    /// and before and after function calls, because even without loops or infinite recursion, wasm function calls
+    /// can mutliply time spent in execution
+    /// execution may not be suspended in the middle of execution the funcion:
+    /// so only do it as the last thing or first thing in the intsruction execution
+    fn check_should_suspend(&mut self) -> ControlFlow<ReasonToBreak> {
+        if let Some(flag) = &self.store.suspend_cond.suspend_flag {
+            if flag.load(core::sync::atomic::Ordering::Acquire) {
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedFlag).into();
+            }
+        }
+
+        #[cfg(feature = "std")]
+        if let Some(when) = &self.store.suspend_cond.timeout_instant {
+            if crate::std::time::Instant::now() >= *when {
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedEpoch).into();
+            }
+        }
+
+        if let Some(mut cb) = self.store.suspend_cond.suspend_cb.take() {
+            let should_suspend = matches!(cb(&self.store), ControlFlow::Break(()));
+            self.store.suspend_cond.suspend_cb = Some(cb); // put it back
+            if should_suspend {
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedCallback).into();
+            }
+        }
+
+        ControlFlow::Continue(())
     }
 
     #[inline(always)]
@@ -382,35 +413,41 @@ impl<'store, 'stack> Executor<'store, 'stack> {
         self.module.swap_with(self.cf.module_addr(), self.store);
         ControlFlow::Continue(())
     }
+    fn exec_call_host(&mut self, host_func: Rc<HostFunction>, func_ref: u32) -> ControlFlow<ReasonToBreak> {
+        let params = self.stack.values.pop_params(&host_func.ty.params);
+        let res =
+            (host_func.func)(FuncContext { store: self.store, module_addr: self.module.id() }, &params).to_cf()?;
+        match res {
+            PotentialCoroCallResult::Return(res) => {
+                self.stack.values.extend_from_wasmvalues(&res);
+                self.cf.incr_instr_ptr();
+                self.check_should_suspend()?; // who knows how long we've spent in host function
+                return ControlFlow::Continue(());
+            }
+            PotentialCoroCallResult::Suspended(suspend_reason, state) => {
+                self.suspended_host_coro =
+                    Some(SuspendedHostCoroState { coro_state: state, coro_orig_function: func_ref });
+                self.cf.incr_instr_ptr();
+                return ReasonToBreak::Suspended(suspend_reason).into();
+            }
+        }
+    }
     fn exec_call_direct(&mut self, v: u32) -> ControlFlow<ReasonToBreak> {
+        self.check_should_suspend()?; // don't commit to function if we should be stopping now
         let func_ref = self.module.resolve_func_addr(v);
         let func_inst = self.store.get_func(func_ref);
         let wasm_func = match &func_inst.func {
             crate::Function::Wasm(wasm_func) => wasm_func,
             crate::Function::Host(host_func) => {
-                let func = &host_func.clone();
-                let params = self.stack.values.pop_params(&host_func.ty.params);
-                let res =
-                    (func.func)(FuncContext { store: self.store, module_addr: self.module.id() }, &params).to_cf()?;
-                match res {
-                    PotentialCoroCallResult::Return(res) => {
-                        self.stack.values.extend_from_wasmvalues(&res);
-                        self.cf.incr_instr_ptr();
-                        return ControlFlow::Continue(());
-                    }
-                    PotentialCoroCallResult::Suspended(suspend_reason, state) => {
-                        self.suspended_host_coro =
-                            Some(SuspendedHostCoroState { coro_state: state, coro_orig_function: func_ref });
-                        return ReasonToBreak::Suspended(suspend_reason).into();
-                    }
-                }
+                return self.exec_call_host(host_func.clone(), func_ref);
             }
         };
 
         self.exec_call(wasm_func.clone(), func_inst.owner)
     }
     fn exec_call_indirect(&mut self, type_addr: u32, table_addr: u32) -> ControlFlow<ReasonToBreak> {
-        // verify that the table is of the right type, this should be validated by the parser already
+        self.check_should_suspend()?; // check if we should suspend now before commiting to function
+                                      // verify that the table is of the right type, this should be validated by the parser already
         let func_ref = {
             let table = self.store.get_table(self.module.resolve_table_addr(table_addr));
             let table_idx: u32 = self.stack.values.pop::<i32>() as u32;
@@ -436,27 +473,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
                     )
                     .into();
                 }
-
-                let host_func = host_func.clone();
-                let params = self.stack.values.pop_params(&host_func.ty.params);
-                let res =
-                    match (host_func.func)(FuncContext { store: self.store, module_addr: self.module.id() }, &params) {
-                        Ok(res) => res,
-                        Err(e) => return ReasonToBreak::Errored(e).into(),
-                    };
-                match res {
-                    PotentialCoroCallResult::Return(res) => {
-                        self.stack.values.extend_from_wasmvalues(&res);
-                        self.cf.incr_instr_ptr();
-                    }
-                    PotentialCoroCallResult::Suspended(suspend_reason, state) => {
-                        self.suspended_host_coro =
-                            Some(SuspendedHostCoroState { coro_state: state, coro_orig_function: func_ref });
-                        return ReasonToBreak::Suspended(suspend_reason).into();
-                    }
-                }
-
-                return ControlFlow::Continue(());
+                return self.exec_call_host(host_func.clone(), func_ref);
             }
         };
 
@@ -505,20 +522,30 @@ impl<'store, 'stack> Executor<'store, 'stack> {
         });
     }
     fn exec_br(&mut self, to: u32) -> ControlFlow<ReasonToBreak> {
-        if self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks).is_none() {
+        let break_type = if let Some(bl_ty) = self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks) {
+            bl_ty
+        } else {
             return self.exec_return();
-        }
+        };
 
         self.cf.incr_instr_ptr();
+
+        if matches!(break_type, BlockType::Loop) {
+            self.check_should_suspend()?;
+        }
+
         ControlFlow::Continue(())
     }
     fn exec_br_if(&mut self, to: u32) -> ControlFlow<ReasonToBreak> {
-        if self.stack.values.pop::<i32>() != 0
-            && self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks).is_none()
-        {
+        let break_type = self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks);
+        if self.stack.values.pop::<i32>() != 0 && break_type.is_none() {
             return self.exec_return();
         }
         self.cf.incr_instr_ptr();
+
+        if matches!(break_type, Some(BlockType::Loop)) {
+            self.check_should_suspend()?;
+        }
         ControlFlow::Continue(())
     }
     fn exec_brtable(&mut self, default: u32, len: u32) -> ControlFlow<ReasonToBreak> {
@@ -540,11 +567,19 @@ impl<'store, 'stack> Executor<'store, 'stack> {
             _ => return ReasonToBreak::Errored(Error::Other("br_table out of bounds".to_string())).into(),
         };
 
-        if self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks).is_none() {
+        let break_type = self.cf.break_to(to, &mut self.stack.values, &mut self.stack.blocks);
+        let break_type = if let Some(br_ty) = break_type {
+            br_ty
+        } else {
             return self.exec_return();
-        }
+        };
 
         self.cf.incr_instr_ptr();
+
+        if matches!(break_type, BlockType::Loop) {
+            self.check_should_suspend()?;
+        }
+
         ControlFlow::Continue(())
     }
     fn exec_return(&mut self) -> ControlFlow<ReasonToBreak> {
@@ -559,6 +594,8 @@ impl<'store, 'stack> Executor<'store, 'stack> {
         }
 
         self.module.swap_with(self.cf.module_addr(), self.store);
+
+        self.check_should_suspend()?;
         ControlFlow::Continue(())
     }
     fn exec_end_block(&mut self) {

crates/tinywasm/src/interpreter/stack/block_stack.rs

@@ -60,7 +60,7 @@ pub(crate) struct BlockFrame {
     pub(crate) ty: BlockType,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub(crate) enum BlockType {
     Loop,
     If,

crates/tinywasm/src/interpreter/stack/call_stack.rs

@@ -107,12 +107,13 @@ impl CallFrame {
         break_to_relative: u32,
         values: &mut super::ValueStack,
         blocks: &mut super::BlockStack,
-    ) -> Option<()> {
+    ) -> Option<BlockType> {
         let break_to = blocks.get_relative_to(break_to_relative, self.block_ptr)?;
-        
+
+        let block_ty = break_to.ty;
         // instr_ptr points to the label instruction, but the next step
         // will increment it by 1 since we're changing the "current" instr_ptr
-        match break_to.ty {
+        match block_ty {
             BlockType::Loop => {
                 // this is a loop, so we want to jump back to the start of the loop
                 self.instr_ptr = break_to.instr_ptr;
@@ -124,7 +125,7 @@ impl CallFrame {
                 if break_to_relative != 0 {
                     // we also want to trim the label stack to the loop (but not including the loop)
                     blocks.truncate(blocks.len() as u32 - break_to_relative);
-                    return Some(());
+                    return Some(BlockType::Loop);
                 }
             }
 
@@ -141,7 +142,7 @@ impl CallFrame {
             }
         }
 
-        Some(())
+        Some(block_ty)
     }
 
     #[inline]

crates/tinywasm/src/lib.rs

@@ -93,7 +93,7 @@ pub(crate) mod log {
 mod error;
 pub use coro::{CoroState, CoroStateResumeResult, PotentialCoroCallResult, SuspendReason};
 pub use error::*;
-pub use func::{FuncHandle, FuncHandleTyped};
+pub use func::{FuncHandle, FuncHandleTyped, SuspendFunc};
 pub use imports::*;
 pub use instance::ModuleInstance;
 pub use module::Module;