feat: implement memory.atomic.notify,wait32,wait64

Added the parking_spot crate, which provides the needed registry for the
operations.

Signed-off-by: Harald Hoyer <harald@profian.com>

crates/runtime/Cargo.toml

@@ -45,6 +45,9 @@ features = [
   "Win32_Security",
 ]
 
+[dev-dependencies]
+once_cell = { workspace = true }
+
 [build-dependencies]
 cc = "1.0"
 

crates/runtime/src/instance.rs

@@ -196,6 +196,17 @@ impl Instance {
         }
     }
 
+    /// Get a locally defined or imported memory.
+    pub(crate) fn get_runtime_memory(&mut self, index: MemoryIndex) -> &mut Memory {
+        if let Some(defined_index) = self.module().defined_memory_index(index) {
+            unsafe { &mut *self.get_defined_memory(defined_index) }
+        } else {
+            let import = self.imported_memory(index);
+            let ctx = unsafe { &mut *import.vmctx };
+            unsafe { &mut *ctx.instance_mut().get_defined_memory(import.index) }
+        }
+    }
+
     /// Return the indexed `VMMemoryDefinition`.
     fn memory(&self, index: DefinedMemoryIndex) -> VMMemoryDefinition {
         unsafe { VMMemoryDefinition::load(self.memory_ptr(index)) }

crates/runtime/src/lib.rs

@@ -83,6 +83,8 @@ mod module_id;
 pub use module_id::{CompiledModuleId, CompiledModuleIdAllocator};
 
 mod cow;
+pub mod parking_spot;
+
 pub use crate::cow::{MemoryImage, MemoryImageSlot, ModuleMemoryImages};
 
 /// Version number of this crate.

crates/runtime/src/libcalls.rs

@@ -56,12 +56,16 @@
 
 use crate::externref::VMExternRef;
 use crate::instance::Instance;
+use crate::parking_spot::{ParkResult, ParkingSpot};
 use crate::table::{Table, TableElementType};
 use crate::vmcontext::{VMCallerCheckedAnyfunc, VMContext};
 use crate::TrapReason;
 use anyhow::Result;
 use std::mem;
 use std::ptr::{self, NonNull};
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
 use wasmtime_environ::{
     DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, TableIndex, TrapCode,
 };
@@ -434,50 +438,89 @@ unsafe fn externref_global_set(vmctx: *mut VMContext, index: u32, externref: *mu
 unsafe fn memory_atomic_notify(
     vmctx: *mut VMContext,
     memory_index: u32,
-    addr: u64,
-    _count: u32,
+    addr_index: u64,
+    count: u32,
 ) -> Result<u32, TrapReason> {
     let memory = MemoryIndex::from_u32(memory_index);
-    let instance = (*vmctx).instance();
-    validate_atomic_addr(instance, memory, addr, 4, 4)?;
-    Err(
-        anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_notify) unsupported",)
-            .into(),
-    )
+    let instance = (*vmctx).instance_mut();
+    let (_, spot) = validate_atomic_addr(instance, memory, addr_index, 4, 4)?;
+
+    if count == 0 {
+        return Ok(0);
+    }
+
+    let unparked_threads = spot.map_or(0, |spot| spot.unpark(addr_index, count));
+
+    Ok(unparked_threads)
 }
 
 // Implementation of `memory.atomic.wait32` for locally defined memories.
 unsafe fn memory_atomic_wait32(
     vmctx: *mut VMContext,
     memory_index: u32,
-    addr: u64,
-    _expected: u32,
-    _timeout: u64,
+    addr_index: u64,
+    expected: u32,
+    timeout: u64,
 ) -> Result<u32, TrapReason> {
+    let timeout = (timeout as i64 >= 0).then(|| Duration::from_nanos(timeout));
+
     let memory = MemoryIndex::from_u32(memory_index);
-    let instance = (*vmctx).instance();
-    validate_atomic_addr(instance, memory, addr, 4, 4)?;
-    Err(
-        anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_wait32) unsupported",)
-            .into(),
-    )
+    let instance = (*vmctx).instance_mut();
+    let (addr, spot) = validate_atomic_addr(instance, memory, addr_index, 4, 4)?;
+
+    // SAFETY: `addr` was validated by `validate_atomic_addr` above.
+    let atomic = unsafe { &*(addr as *const AtomicU32) };
+
+    // We want the sequential consistency of `SeqCst` to ensure that the `load` sees the value that the `notify` will/would see.
+    // All WASM atomic operations are also `SeqCst`.
+    let validate = || atomic.load(Ordering::SeqCst) == expected;
+
+    do_atomic_wait(addr_index, timeout, spot, validate)
 }
 
 // Implementation of `memory.atomic.wait64` for locally defined memories.
 unsafe fn memory_atomic_wait64(
     vmctx: *mut VMContext,
     memory_index: u32,
-    addr: u64,
-    _expected: u64,
-    _timeout: u64,
+    addr_index: u64,
+    expected: u64,
+    timeout: u64,
 ) -> Result<u32, TrapReason> {
+    let timeout = (timeout as i64 >= 0).then(|| Duration::from_nanos(timeout));
+
     let memory = MemoryIndex::from_u32(memory_index);
-    let instance = (*vmctx).instance();
-    validate_atomic_addr(instance, memory, addr, 8, 8)?;
-    Err(
-        anyhow::anyhow!("unimplemented: wasm atomics (fn memory_atomic_wait64) unsupported",)
-            .into(),
-    )
+    let instance = (*vmctx).instance_mut();
+    let (addr, spot) = validate_atomic_addr(instance, memory, addr_index, 8, 8)?;
+
+    // SAFETY: `addr` was validated by `validate_atomic_addr` above.
+    let atomic = unsafe { &*(addr as *const AtomicU64) };
+
+    // We want the sequential consistency of `SeqCst` to ensure that the `load` sees the value that the `notify` will/would see.
+    // All WASM atomic operations are also `SeqCst`.
+    let validate = || atomic.load(Ordering::SeqCst) == expected;
+
+    do_atomic_wait(addr_index, timeout, spot, validate)
+}
+
+#[inline]
+fn do_atomic_wait(
+    addr_index: u64,
+    timeout: Option<Duration>,
+    spot: Option<Arc<ParkingSpot>>,
+    validate: impl FnOnce() -> bool,
+) -> Result<u32, TrapReason> {
+    if let Some(spot) = spot {
+        match spot.park(addr_index, validate, timeout) {
+            ParkResult::Unparked => Ok(0),
+            ParkResult::Invalid => Ok(1),
+            ParkResult::TimedOut => Ok(2),
+        }
+    } else if let Some(timeout) = timeout {
+        std::thread::sleep(timeout);
+        Ok(2)
+    } else {
+        Err(TrapReason::Wasm(TrapCode::AlwaysTrapAdapter))
+    }
 }
 
 macro_rules! ensure {
@@ -493,22 +536,30 @@ macro_rules! ensure {
 /// check is here so we don't segfault from Rust. For other configurations,
 /// these checks are required anyways.
 unsafe fn validate_atomic_addr(
-    instance: &Instance,
+    instance: &mut Instance,
     memory: MemoryIndex,
     addr: u64,
     access_size: u64,
     access_alignment: u64,
-) -> Result<(), TrapCode> {
+) -> Result<(*mut u8, Option<Arc<ParkingSpot>>), TrapCode> {
     debug_assert!(access_alignment.is_power_of_two());
     ensure!(addr % access_alignment == 0, TrapCode::HeapMisaligned);
 
-    let length = u64::try_from(instance.get_memory(memory).current_length()).unwrap();
+    let mem = instance.get_memory(memory);
+
+    let length = u64::try_from(mem.current_length()).unwrap();
     ensure!(
         addr.saturating_add(access_size) < length,
         TrapCode::HeapOutOfBounds
     );
 
-    Ok(())
+    let mem_definition = instance.get_memory(memory);
+
+    let mem = instance.get_runtime_memory(memory);
+
+    let spot = mem.as_shared_memory().map(|mem| mem.parking_spot());
+
+    Ok((mem_definition.base.add(addr as usize), spot))
 }
 
 // Hook for when an instance runs out of fuel.

crates/runtime/src/memory.rs

@@ -3,6 +3,7 @@
 //! `RuntimeLinearMemory` is to WebAssembly linear memories what `Table` is to WebAssembly tables.
 
 use crate::mmap::Mmap;
+use crate::parking_spot::ParkingSpot;
 use crate::vmcontext::VMMemoryDefinition;
 use crate::MemoryImage;
 use crate::MemoryImageSlot;
@@ -462,6 +463,7 @@ impl SharedMemory {
             memory: memory,
             ty,
             def,
+            spot: Arc::new(ParkingSpot::default()),
         }))))
     }
 
@@ -484,12 +486,18 @@ impl SharedMemory {
     pub fn vmmemory_ptr(&self) -> *const VMMemoryDefinition {
         &self.0.read().unwrap().def.0 as *const _
     }
+
+    /// Return a reference to the shared memory's [ParkingSpot].
+    pub fn parking_spot(&self) -> Arc<ParkingSpot> {
+        self.0.read().unwrap().spot.clone()
+    }
 }
 
 struct SharedMemoryInner {
     memory: Box<dyn RuntimeLinearMemory>,
     ty: wasmtime_environ::Memory,
     def: LongTermVMMemoryDefinition,
+    spot: Arc<ParkingSpot>,
 }
 
 /// Shared memory needs some representation of a `VMMemoryDefinition` for