feat: implement memory.atomic.notify,wait32,wait64 (#5255)

* 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 <[email protected]>

* fix: change trap message for HeapMisaligned

The threads spec test wants "unaligned atomic"
instead of "misaligned memory access".

Signed-off-by: Harald Hoyer <[email protected]>

* tests: add test for atomic wait on non-shared memory

Signed-off-by: Harald Hoyer <[email protected]>

* tests: add tests/spec_testsuite/proposals/threads

without pooling and reference types.
Also "shared_memory" is added to the "spectest" interface.

Signed-off-by: Harald Hoyer <[email protected]>

* tests: add atomics_notify.wast

checking that notify with 0 waiters returns 0 on shared and non-shared
memory.

Signed-off-by: Harald Hoyer <[email protected]>

* tests: add tests for atomic wait on shared memory

- return 2 - timeout for 0
- return 2 - timeout for 1000ns
- return 1 - invalid value

Signed-off-by: Harald Hoyer <[email protected]>

* fixup! feat: implement memory.atomic.notify,wait32,wait64

Signed-off-by: Harald Hoyer <[email protected]>

* fixup! feat: implement memory.atomic.notify,wait32,wait64

Signed-off-by: Harald Hoyer <[email protected]>

Signed-off-by: Harald Hoyer <[email protected]>

Author: haraldh

Cargo.lock

@@ -63,6 +63,7 @@ wasmtime-wast = { workspace = true, features = ['component-model'] }
 wasmtime-component-util = { workspace = true }
 component-macro-test = { path = "crates/misc/component-macro-test" }
 component-test-util = { workspace = true }
+bstr = "0.2.17"
 
 [target.'cfg(windows)'.dev-dependencies]
 windows-sys = { workspace = true, features = ["Win32_System_Memory"] }

Cargo.toml

@@ -12,6 +12,7 @@ use std::process::Command;
 
 fn main() -> anyhow::Result<()> {
     println!("cargo:rerun-if-changed=build.rs");
+
     let out_dir = PathBuf::from(
         env::var_os("OUT_DIR").expect("The OUT_DIR environment variable must be set"),
     );
@@ -43,6 +44,7 @@ fn main() -> anyhow::Result<()> {
                     "tests/spec_testsuite/proposals/multi-memory",
                     strategy,
                 )?;
+                test_directory_module(out, "tests/spec_testsuite/proposals/threads", strategy)?;
             } else {
                 println!(
                     "cargo:warning=The spec testsuite is disabled. To enable, run `git submodule \
@@ -62,7 +64,6 @@ fn main() -> anyhow::Result<()> {
     drop(Command::new("rustfmt").arg(&output).status());
     Ok(())
 }
-
 fn test_directory_module(
     out: &mut String,
     path: impl AsRef<Path>,
@@ -91,7 +92,7 @@ fn test_directory(
                 return None;
             }
             // Ignore files starting with `.`, which could be editor temporary files
-            if p.file_stem()?.to_str()?.starts_with(".") {
+            if p.file_stem()?.to_str()?.starts_with('.') {
                 return None;
             }
             Some(p)
@@ -116,8 +117,7 @@ fn extract_name(path: impl AsRef<Path>) -> String {
         .expect("filename should have a stem")
         .to_str()
         .expect("filename should be representable as a string")
-        .replace("-", "_")
-        .replace("/", "_")
+        .replace(['-', '/'], "_")
 }
 
 fn with_test_module<T>(
@@ -163,7 +163,7 @@ fn write_testsuite_tests(
         "    crate::wast::run_wast(r#\"{}\"#, crate::wast::Strategy::{}, {}).unwrap();",
         path.display(),
         strategy,
-        pooling
+        pooling,
     )?;
     writeln!(out, "}}")?;
     writeln!(out)?;

build.rs

@@ -84,6 +84,10 @@ pub enum Trap {
     /// When wasm code is configured to consume fuel and it runs out of fuel
     /// then this trap will be raised.
     OutOfFuel,
+
+    /// Used to indicate that a trap was raised by atomic wait operations on non shared memory.
+    AtomicWaitNonSharedMemory,
+    // if adding a variant here be sure to update the `check!` macro below
 }
 
 impl fmt::Display for Trap {
@@ -93,7 +97,7 @@ impl fmt::Display for Trap {
         let desc = match self {
             StackOverflow => "call stack exhausted",
             MemoryOutOfBounds => "out of bounds memory access",
-            HeapMisaligned => "misaligned memory access",
+            HeapMisaligned => "unaligned atomic",
             TableOutOfBounds => "undefined element: out of bounds table access",
             IndirectCallToNull => "uninitialized element",
             BadSignature => "indirect call type mismatch",
@@ -104,6 +108,7 @@ impl fmt::Display for Trap {
             Interrupt => "interrupt",
             AlwaysTrapAdapter => "degenerate component adapter called",
             OutOfFuel => "all fuel consumed by WebAssembly",
+            AtomicWaitNonSharedMemory => "atomic wait on non-shared memory",
         };
         write!(f, "wasm trap: {desc}")
     }
@@ -217,6 +222,7 @@ pub fn lookup_trap_code(section: &[u8], offset: usize) -> Option<Trap> {
         Interrupt
         AlwaysTrapAdapter
         OutOfFuel
+        AtomicWaitNonSharedMemory
     }
 
     if cfg!(debug_assertions) {

crates/environ/src/trap_encoding.rs

@@ -506,7 +506,7 @@ pub fn spectest(mut fuzz_config: generators::Config, test: generators::SpecTest)
     fuzz_config.set_spectest_compliant();
     log::debug!("running {:?}", test.file);
     let mut wast_context = WastContext::new(fuzz_config.to_store());
-    wast_context.register_spectest().unwrap();
+    wast_context.register_spectest(false).unwrap();
     wast_context
         .run_buffer(test.file, test.contents.as_bytes())
         .unwrap();

crates/fuzzing/src/oracles.rs

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

crates/runtime/Cargo.toml

@@ -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/instance.rs

@@ -40,6 +40,7 @@ mod instance;
 mod memory;
 mod mmap;
 mod mmap_vec;
+mod parking_spot;
 mod table;
 mod traphandlers;
 mod vmcontext;

crates/runtime/src/lib.rs

@@ -55,13 +55,13 @@
 //! ```
 
 use crate::externref::VMExternRef;
-use crate::instance::Instance;
 use crate::table::{Table, TableElementType};
 use crate::vmcontext::{VMCallerCheckedAnyfunc, VMContext};
-use crate::TrapReason;
+use crate::{SharedMemory, TrapReason};
 use anyhow::Result;
 use std::mem;
 use std::ptr::{self, NonNull};
+use std::time::{Duration, Instant};
 use wasmtime_environ::{
     DataIndex, ElemIndex, FuncIndex, GlobalIndex, MemoryIndex, TableIndex, Trap,
 };
@@ -434,81 +434,81 @@ 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();
+    instance
+        .get_memory(memory)
+        .validate_addr(addr_index, 4, 4)?;
+
+    let shared_mem = instance.get_runtime_memory(memory).as_shared_memory();
+
+    if count == 0 {
+        return Ok(0);
+    }
+
+    let unparked_threads = shared_mem.map_or(0, |shared_mem| {
+        // SAFETY: checked `addr_index` above
+        unsafe { shared_mem.unchecked_atomic_notify(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> {
+    // convert timeout to Instant, before any wait happens on locking
+    let timeout = (timeout as i64 >= 0).then(|| Instant::now() + 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 = instance
+        .get_memory(memory)
+        .validate_addr(addr_index, 4, 4)?;
+
+    let shared_mem: SharedMemory = instance
+        .get_runtime_memory(memory)
+        .as_shared_memory()
+        .ok_or(Trap::AtomicWaitNonSharedMemory)?;
+
+    // SAFETY: checked `addr_index` above
+    let res = unsafe { shared_mem.unchecked_atomic_wait32(addr_index, addr, expected, timeout) };
+    Ok(res)
 }
 
 // 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 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(),
-    )
-}
-
-macro_rules! ensure {
-    ($cond:expr, $trap:expr) => {
-        if !($cond) {
-            return Err($trap);
-        }
-    };
-}
+    // convert timeout to Instant, before any wait happens on locking
+    let timeout = (timeout as i64 >= 0).then(|| Instant::now() + Duration::from_nanos(timeout));
 
-/// In the configurations where bounds checks were elided in JIT code (because
-/// we are using static memories with virtual memory guard pages) this manual
-/// 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,
-    memory: MemoryIndex,
-    addr: u64,
-    access_size: u64,
-    access_alignment: u64,
-) -> Result<(), Trap> {
-    debug_assert!(access_alignment.is_power_of_two());
-    ensure!(addr % access_alignment == 0, Trap::HeapMisaligned);
-
-    let length = u64::try_from(instance.get_memory(memory).current_length()).unwrap();
-    ensure!(
-        addr.saturating_add(access_size) < length,
-        Trap::MemoryOutOfBounds
-    );
-
-    Ok(())
+    let memory = MemoryIndex::from_u32(memory_index);
+    let instance = (*vmctx).instance_mut();
+    let addr = instance
+        .get_memory(memory)
+        .validate_addr(addr_index, 8, 8)?;
+
+    let shared_mem: SharedMemory = instance
+        .get_runtime_memory(memory)
+        .as_shared_memory()
+        .ok_or(Trap::AtomicWaitNonSharedMemory)?;
+
+    // SAFETY: checked `addr_index` above
+    let res = unsafe { shared_mem.unchecked_atomic_wait64(addr_index, addr, expected, timeout) };
+    Ok(res)
 }
 
 // Hook for when an instance runs out of fuel.