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

Use parking_lot_core, as this established crate does exactly what is
needed here.

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

crates/runtime/Cargo.toml

@@ -26,6 +26,7 @@ anyhow = { workspace = true }
 memfd = "0.6.1"
 paste = "1.0.3"
 encoding_rs = { version = "0.8.31", optional = true }
+parking_lot_core = "0.9.4"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 mach = "0.3.2"

crates/runtime/src/libcalls.rs

@@ -60,8 +60,13 @@ use crate::table::{Table, TableElementType};
 use crate::vmcontext::{VMCallerCheckedAnyfunc, VMContext};
 use crate::TrapReason;
 use anyhow::Result;
+use parking_lot_core::{
+    park, unpark_all, unpark_one, ParkResult, DEFAULT_PARK_TOKEN, DEFAULT_UNPARK_TOKEN,
+};
 use std::mem;
 use std::ptr::{self, NonNull};
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::time::{Duration, Instant};
 use wasmtime_environ::{
     DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, TableIndex, TrapCode,
 };
@@ -435,49 +440,122 @@ unsafe fn memory_atomic_notify(
     vmctx: *mut VMContext,
     memory_index: u32,
     addr: u64,
-    _count: u32,
+    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 addr = validate_atomic_addr(instance, memory, addr, 4, 4)?;
+    if count == 0 {
+        return Ok(0);
+    }
+
+    let unparked_threads = if count == u32::MAX {
+        // SAFETY: `addr` is a valid pointer into the given memory and unique to parking_lot_core.
+        unsafe { unpark_all(addr as _, DEFAULT_UNPARK_TOKEN) }
+    } else {
+        let mut num = 0;
+        for _ in 0..count {
+            // SAFETY: `addr` is a valid pointer into the given memory and unique to parking_lot_core.
+            let num_t = unsafe { unpark_one(addr as _, |_| DEFAULT_UNPARK_TOKEN).unparked_threads };
+
+            if num_t == 0 {
+                break;
+            }
+
+            num += num_t;
+        }
+        num
+    };
+
+    u32::try_from(unparked_threads).map_err(|e| TrapReason::user_with_backtrace(e.into()))
 }
 
 // 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,
+    expected: u32,
+    timeout: u64,
 ) -> Result<u32, TrapReason> {
+    // convert to absolute timestamp as soon as possible
+    let wait_until = (timeout > 0)
+        .then(|| {
+            Instant::now()
+                .checked_add(Duration::from_nanos(timeout))
+                .ok_or_else(|| {
+                    TrapReason::user_with_backtrace(anyhow::anyhow!(
+                        "overflow when adding timeout to current time"
+                    ))
+                })
+        })
+        .transpose()?;
+
     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 addr = validate_atomic_addr(instance, memory, addr, 4, 4)?;
+
+    // SAFETY: `addr` was validated by `validate_atomic_addr` above.
+    let atomic = unsafe { &*(addr as *const AtomicU32) };
+    // SAFETY: `addr` is a valid pointer into the given memory and unique to parking_lot_core.
+    match unsafe {
+        park(
+            addr as usize,
+            || atomic.load(Ordering::SeqCst) == expected,
+            || {},
+            |_, _| {},
+            DEFAULT_PARK_TOKEN,
+            wait_until,
+        )
+    } {
+        ParkResult::Unparked(_) => return Ok(0),
+        ParkResult::Invalid => return Ok(1),
+        ParkResult::TimedOut => return Ok(2),
+    }
 }
 
 // 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,
+    expected: u64,
+    timeout: u64,
 ) -> Result<u32, TrapReason> {
+    // convert to absolute timestamp as soon as possible
+    let wait_until = (timeout > 0)
+        .then(|| {
+            Instant::now()
+                .checked_add(Duration::from_nanos(timeout))
+                .ok_or_else(|| {
+                    TrapReason::user_with_backtrace(anyhow::anyhow!(
+                        "overflow when adding timeout to current time"
+                    ))
+                })
+        })
+        .transpose()?;
+
     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 addr = validate_atomic_addr(instance, memory, addr, 8, 8)?;
+
+    // SAFETY: `addr` was validated by `validate_atomic_addr` above.
+    let atomic = unsafe { &*(addr as *const AtomicU64) };
+    // SAFETY: `addr` is a valid pointer into the given memory and unique to parking_lot_core.
+    match unsafe {
+        park(
+            addr as usize,
+            || atomic.load(Ordering::SeqCst) == expected,
+            || {},
+            |_, _| {},
+            DEFAULT_PARK_TOKEN,
+            wait_until,
+        )
+    } {
+        ParkResult::Unparked(_) => return Ok(0),
+        ParkResult::Invalid => return Ok(1),
+        ParkResult::TimedOut => return Ok(2),
+    }
 }
 
 macro_rules! ensure {
@@ -498,17 +576,19 @@ unsafe fn validate_atomic_addr(
     addr: u64,
     access_size: u64,
     access_alignment: u64,
-) -> Result<(), TrapCode> {
+) -> Result<*mut u8, 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(())
+    Ok(mem.base.add(addr as usize))
 }
 
 // Hook for when an instance runs out of fuel.