Add epoch-based interruption for cooperative async timeslicing.

This PR introduces a new way of performing cooperative timeslicing that
is intended to replace the "fuel" mechanism. The tradeoff is that this
mechanism interrupts with less precision: not at deterministic points
where fuel runs out, but rather when the Engine enters a new epoch. The
generated code instrumentation is substantially faster, however, because
it does not need to do as much work as when tracking fuel; it only loads
the global "epoch counter" and does a compare-and-branch at backedges
and function prologues.

This change has been measured as ~twice as fast as fuel-based
timeslicing for some workloads, especially control-flow-intensive
workloads such as the SpiderMonkey JS interpreter on Wasm/WASI.

The intended interface is that the embedder of the `Engine` performs an
`engine.increment_epoch()` call periodically, e.g. once per millisecond.
An async invocation of a Wasm guest on a `Store` can specify a number of
epoch-ticks that are allowed before an async yield back to the
executor's event loop. (The initial amount and automatic "refills" are
configured on the `Store`, just as for fuel.) This call does only
signal-safe work (it increments an `AtomicU64`) so could be invoked from
a periodic signal, or from a thread that wakes up once per period.

crates/runtime/src/instance.rs

@@ -21,6 +21,7 @@ use std::convert::TryFrom;
 use std::hash::Hash;
 use std::ops::Range;
 use std::ptr::NonNull;
+use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
 use std::{mem, ptr, slice};
 use wasmtime_environ::{
@@ -203,6 +204,11 @@ impl Instance {
         unsafe { self.vmctx_plus_offset(self.offsets.vmctx_interrupts()) }
     }
 
+    /// Return a pointer to the global epoch counter used by this instance.
+    pub fn epoch_ptr(&self) -> *mut *const AtomicU64 {
+        unsafe { self.vmctx_plus_offset(self.offsets.vmctx_epoch_ptr()) }
+    }
+
     /// Return a pointer to the `VMExternRefActivationsTable`.
     pub fn externref_activations_table(&self) -> *mut *mut VMExternRefActivationsTable {
         unsafe { self.vmctx_plus_offset(self.offsets.vmctx_externref_activations_table()) }

crates/runtime/src/instance/allocator.rs

@@ -463,6 +463,7 @@ fn initialize_instance(
 unsafe fn initialize_vmcontext(instance: &mut Instance, req: InstanceAllocationRequest) {
     if let Some(store) = req.store.as_raw() {
         *instance.interrupts() = (*store).vminterrupts();
+        *instance.epoch_ptr() = (*store).epoch_ptr();
         *instance.externref_activations_table() = (*store).externref_activations_table().0;
         instance.set_store(store);
     }

crates/runtime/src/instance/allocator/pooling/uffd.rs

@@ -438,6 +438,7 @@ mod test {
         Imports, InstanceAllocationRequest, InstanceLimits, ModuleLimits,
         PoolingAllocationStrategy, Store, StorePtr, VMSharedSignatureIndex,
     };
+    use std::sync::atomic::AtomicU64;
     use std::sync::Arc;
     use wasmtime_environ::{Memory, MemoryPlan, MemoryStyle, Module, PrimaryMap, Tunables};
 
@@ -546,6 +547,12 @@ mod test {
                 fn out_of_gas(&mut self) -> Result<(), anyhow::Error> {
                     Ok(())
                 }
+                fn epoch_ptr(&self) -> *const AtomicU64 {
+                    std::ptr::null()
+                }
+                fn new_epoch(&mut self) -> Result<u64, anyhow::Error> {
+                    Ok(0)
+                }
             }
             struct MockModuleInfo;
             impl crate::ModuleInfoLookup for MockModuleInfo {

crates/runtime/src/lib.rs

@@ -20,6 +20,8 @@
     )
 )]
 
+use std::sync::atomic::AtomicU64;
+
 use anyhow::Error;
 
 mod export;
@@ -84,6 +86,11 @@ pub unsafe trait Store {
     /// in the `VMContext`.
     fn vminterrupts(&self) -> *mut VMInterrupts;
 
+    /// Returns a pointer to the global epoch counter.
+    ///
+    /// Used to configure the `VMContext` on initialization.
+    fn epoch_ptr(&self) -> *const AtomicU64;
+
     /// Returns the externref management structures necessary for this store.
     ///
     /// The first element returned is the table in which externrefs are stored
@@ -119,4 +126,8 @@ pub unsafe trait Store {
     /// is returned that's raised as a trap. Otherwise wasm execution will
     /// continue as normal.
     fn out_of_gas(&mut self) -> Result<(), Error>;
+    /// Callback invoked whenever an instance observes a new epoch
+    /// number. Cannot fail; cooperative epoch-based yielding is
+    /// completely semantically transparent. Returns the new deadline.
+    fn new_epoch(&mut self) -> Result<u64, Error>;
 }

crates/runtime/src/libcalls.rs

@@ -557,3 +557,11 @@ pub unsafe extern "C" fn wasmtime_out_of_gas(vmctx: *mut VMContext) {
         Err(err) => crate::traphandlers::raise_user_trap(err),
     }
 }
+
+/// Hook for when an instance observes that the epoch has changed.
+pub unsafe extern "C" fn wasmtime_new_epoch(vmctx: *mut VMContext) -> u64 {
+    match (*(*vmctx).instance().store()).new_epoch() {
+        Ok(new_deadline) => new_deadline,
+        Err(err) => crate::traphandlers::raise_user_trap(err),
+    }
+}

crates/runtime/src/vmcontext.rs

@@ -631,6 +631,7 @@ impl VMBuiltinFunctionsArray {
         ptrs[BuiltinFunctionIndex::memory_atomic_wait64().index() as usize] =
             wasmtime_memory_atomic_wait64 as usize;
         ptrs[BuiltinFunctionIndex::out_of_gas().index() as usize] = wasmtime_out_of_gas as usize;
+        ptrs[BuiltinFunctionIndex::new_epoch().index() as usize] = wasmtime_new_epoch as usize;
 
         if cfg!(debug_assertions) {
             for i in 0..ptrs.len() {
@@ -694,12 +695,18 @@ pub struct VMInterrupts {
     /// turning positive a wasm trap will be generated. This field is only
     /// modified if wasm is configured to consume fuel.
     pub fuel_consumed: UnsafeCell<i64>,
+
+    /// Deadline epoch for interruption: if epoch-based interruption
+    /// is enabled and the global (per engine) epoch counter is
+    /// observed to reach or exceed this value, the guest code will
+    /// yield if running asynchronously.
+    pub epoch_deadline: UnsafeCell<u64>,
 }
 
-// The `VMInterrupts` type is a pod-type with no destructor, and we only access
-// `stack_limit` from other threads, so add in these trait impls which are
-// otherwise not available due to the `fuel_consumed` variable in
-// `VMInterrupts`.
+// The `VMInterrupts` type is a pod-type with no destructor, and we
+// only access `stack_limit` from other threads, so add in these trait
+// impls which are otherwise not available due to the `fuel_consumed`
+// and `epoch_deadline` variables in `VMInterrupts`.
 //
 // Note that users of `fuel_consumed` understand that the unsafety encompasses
 // ensuring that it's only mutated/accessed from one thread dynamically.
@@ -719,6 +726,7 @@ impl Default for VMInterrupts {
         VMInterrupts {
             stack_limit: AtomicUsize::new(usize::max_value()),
             fuel_consumed: UnsafeCell::new(0),
+            epoch_deadline: UnsafeCell::new(0),
         }
     }
 }