fix(ECMAScript): Atomics Async Critical Section in the Main Thread (#976)

Author: komyg

nova_vm/src/ecmascript/builtins/structured_data/atomics_object.rs

@@ -5,10 +5,7 @@
 use std::{
     hint::assert_unchecked,
     ops::ControlFlow,
-    sync::{
-        Arc,
-        atomic::{AtomicBool, Ordering as StdOrdering},
-    },
+    sync::Arc,
     thread::{self, JoinHandle},
     time::Duration,
 };
@@ -1421,6 +1418,7 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
     gc: NoGcScope<'gc, '_>,
 ) -> Value<'gc> {
     let slot = buffer.as_slice(agent).slice_from(byte_index_in_buffer);
+    let data_block = buffer.get_data_block(agent).clone();
     // 14. Let WL be GetWaiterList(block, byteIndexInBuffer).
     // 15. If mode is sync, then
     // a. Let promiseCapability be blocking.
@@ -1429,6 +1427,14 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
     // a. Let promiseCapability be ! NewPromiseCapability(%Promise%).
     // b. Let resultObject be OrdinaryObjectCreate(%Object.prototype%).
     // 17. Perform EnterCriticalSection(WL).
+
+    // SAFETY: buffer is a valid SharedArrayBuffer and cannot be detached. A 0-sized SAB would
+    // have a dangling data block, but Atomics.wait requires `byteIndex` to be within bounds,
+    // so a 0-sized SAB would have been rejected earlier with a RangeError.
+    let waiters = unsafe { data_block.get_or_init_waiters() };
+    let waiter_record = WaiterRecord::new_shared();
+    let mut guard = waiters.lock().unwrap();
+
     // 18. Let elementType be TypedArrayElementType(typedArray).
     // 19. Let w be GetValueFromBuffer(buffer, byteIndexInBuffer, elementType, true, seq-cst).
     let v_not_equal_to_w = if IS_I64 {
@@ -1447,6 +1453,8 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
     // 20. If v ≠ w, then
     if v_not_equal_to_w {
         // a. Perform LeaveCriticalSection(WL).
+        drop(guard);
+
         // b. If mode is sync, return "not-equal".
         if !IS_ASYNC {
             return BUILTIN_STRING_MEMORY.not_equal.into();
@@ -1464,6 +1472,7 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
         //    timeouts. Asynchronous immediate timeouts have special handling
         //    in order to fail fast and avoid unnecessary Promise jobs.
         // b. Perform LeaveCriticalSection(WL).
+        drop(guard);
         // c. Perform ! CreateDataPropertyOrThrow(resultObject, "async", false).
         // d. Perform ! CreateDataPropertyOrThrow(resultObject, "value", "timed-out").
         let result_object =
@@ -1484,32 +1493,10 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
     //         [[Result]]: "ok"
     // }.
     // 28. Perform AddWaiter(WL, waiterRecord).
+    guard.push_to_list(byte_index_in_buffer, waiter_record.clone());
     // 29. If mode is sync, then
     if !IS_ASYNC {
-        let data_block = buffer.get_data_block(agent);
-        // SAFETY: buffer is a valid SharedArrayBuffer and cannot be detached. A 0-sized SAB would
-        // have a dangling data block, but Atomics.wait requires `byteIndex` to be within bounds,
-        // so a 0-sized SAB would have been rejected earlier with a RangeError.
-        let waiters = unsafe { data_block.get_or_init_waiters() };
-        let waiter_record = WaiterRecord::new_shared();
-        let mut guard = waiters.lock().unwrap();
-
-        // Re-read value under critical section to avoid TOCTOU race.
-        let slot = data_block.as_racy_slice().slice_from(byte_index_in_buffer);
-        let v_changed = if IS_I64 {
-            let slot = unsafe { slot.as_aligned::<u64>().unwrap_unchecked() };
-            v as u64 != slot.load(Ordering::SeqCst)
-        } else {
-            let slot = unsafe { slot.as_aligned::<u32>().unwrap_unchecked() };
-            v as i32 as u32 != slot.load(Ordering::SeqCst)
-        };
-        if v_changed {
-            return BUILTIN_STRING_MEMORY.not_equal.into();
-        }
-
         // a. Perform SuspendThisAgent(WL, waiterRecord).
-        guard.push_to_list(byte_index_in_buffer, waiter_record.clone());
-
         if t == u64::MAX {
             waiter_record.wait(guard);
         } else {
@@ -1532,17 +1519,21 @@ fn do_wait_critical<'gc, const IS_ASYNC: bool, const IS_I64: bool>(
         let promise = Global::new(agent, promise_capability.promise.unbind());
         // 30. Else if timeoutTime is finite, then
         // a. Perform EnqueueAtomicsWaitAsyncTimeoutJob(WL, waiterRecord).
-        let buffer = buffer.get_data_block(agent).clone();
+
+        let data_block_clone = data_block.clone();
         enqueue_atomics_wait_async_job::<IS_I64>(
             agent,
-            buffer,
+            data_block_clone,
             byte_index_in_buffer,
-            v,
+            waiter_record.clone(),
             t,
             promise,
             gc,
         );
+
         // 31. Perform LeaveCriticalSection(WL).
+        drop(guard);
+
         // 33. Perform ! CreateDataPropertyOrThrow(resultObject, "async", true).
         // 34. Perform ! CreateDataPropertyOrThrow(resultObject, "value", promiseCapability.[[Promise]]).
         let result_object =
@@ -1672,7 +1663,7 @@ impl WaitAsyncJob {
         // c. Perform LeaveCriticalSection(WL).
         let promise_capability = PromiseCapability::from_promise(promise, true);
         let result = match result {
-            WaitResult::Ok | WaitResult::NotEqual => BUILTIN_STRING_MEMORY.ok.into(),
+            WaitResult::Ok => BUILTIN_STRING_MEMORY.ok.into(),
             WaitResult::TimedOut => BUILTIN_STRING_MEMORY.timed_out.into(),
         };
         unwrap_try(promise_capability.try_resolve(agent, result, gc));
@@ -1688,43 +1679,21 @@ impl WaitAsyncJob {
 /// unused.
 fn enqueue_atomics_wait_async_job<const IS_I64: bool>(
     agent: &mut Agent,
-    buffer: SharedDataBlock,
+    data_block: SharedDataBlock,
     byte_index_in_buffer: usize,
-    v: i64,
+    waiter_record: Arc<WaiterRecord>,
     t: u64,
     promise: Global<Promise>,
     gc: NoGcScope,
 ) {
     // 1. Let timeoutJob be a new Job Abstract Closure with no parameters that
     //    captures WL and waiterRecord and performs the following steps when
     //    called:
-    let signal = Arc::new(AtomicBool::new(false));
-    let s = signal.clone();
     let handle = thread::spawn(move || {
         // SAFETY: buffer is a cloned SharedDataBlock; non-dangling.
-        let waiters = unsafe { buffer.get_or_init_waiters() };
-        let waiter_record = WaiterRecord::new_shared();
+        let waiters = unsafe { data_block.get_or_init_waiters() };
         let mut guard = waiters.lock().unwrap();
 
-        // Re-check the value under the critical section.
-        let slot = buffer.as_racy_slice().slice_from(byte_index_in_buffer);
-        let v_not_equal = if IS_I64 {
-            let slot = unsafe { slot.as_aligned::<u64>().unwrap_unchecked() };
-            v as u64 != slot.load(Ordering::SeqCst)
-        } else {
-            let slot = unsafe { slot.as_aligned::<u32>().unwrap_unchecked() };
-            v as i32 as u32 != slot.load(Ordering::SeqCst)
-        };
-
-        // Signal the main thread that we have the lock and are about to sleep.
-        s.store(true, StdOrdering::Release);
-
-        if v_not_equal {
-            return WaitResult::NotEqual;
-        }
-
-        guard.push_to_list(byte_index_in_buffer, waiter_record.clone());
-
         if t == u64::MAX {
             waiter_record.wait(guard);
         } else {
@@ -1735,6 +1704,8 @@ fn enqueue_atomics_wait_async_job<const IS_I64: bool>(
                 guard.remove_from_list(byte_index_in_buffer, waiter_record);
 
                 // 31. Perform LeaveCriticalSection(WL).
+                drop(guard);
+
                 // 32. If mode is sync, return waiterRecord.[[Result]].
                 return WaitResult::TimedOut;
             }
@@ -1749,9 +1720,6 @@ fn enqueue_atomics_wait_async_job<const IS_I64: bool>(
             _has_timeout: t != u64::MAX,
         }))),
     };
-    while !signal.load(StdOrdering::Acquire) {
-        // Wait until the thread has started up and is about to go to sleep.
-    }
     // 2. Let now be the time value (UTC) identifying the current time.
     // 3. Let currentRealm be the current Realm Record.
     // 4. Perform HostEnqueueTimeoutJob(timeoutJob, currentRealm, 𝔽(waiterRecord.[[TimeoutTime]]) - now).

nova_vm/src/ecmascript/types/spec/data_block.rs

@@ -480,7 +480,6 @@ impl WaiterRecord {
 pub(crate) enum WaitResult {
     Ok,
     TimedOut,
-    NotEqual,
 }
 
 #[cfg(feature = "shared-array-buffer")]