wake up all futex threads at the same time

On macOS, the futex operations return the number of outstanding waiters. For this count to be correct, we have to remove all eligible threads from the internal futex list before any of the wakeup callbacks are run. This has the nice side-effect of making the runtime removal operation linear in the number of threads instead of quadratic.

Author: joboet

src/concurrency/sync.rs

@@ -128,7 +128,7 @@ struct Condvar {
 /// The futex state.
 #[derive(Default, Debug)]
 struct Futex {
-    waiters: VecDeque<FutexWaiter>,
+    waiters: Vec<FutexWaiter>,
     /// Tracks the happens-before relationship
     /// between a futex-wake and a futex-wait
     /// during a non-spurious wake event.
@@ -748,7 +748,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let mut futex = futex_ref.0.borrow_mut();
         let waiters = &mut futex.waiters;
         assert!(waiters.iter().all(|waiter| waiter.thread != thread), "thread is already waiting");
-        waiters.push_back(FutexWaiter { thread, bitset });
+        waiters.push(FutexWaiter { thread, bitset });
         drop(futex);
 
         this.block_thread(
@@ -782,9 +782,14 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         );
     }
 
-    /// Wake up the first thread in the queue that matches any of the bits in the bitset.
-    /// Returns whether anything was woken.
-    fn futex_wake(&mut self, futex_ref: &FutexRef, bitset: u32) -> InterpResult<'tcx, bool> {
+    /// Wake up `count` of the threads in the queue that match any of the bits
+    /// in the bitset. Returns how many threads were woken.
+    fn futex_wake(
+        &mut self,
+        futex_ref: &FutexRef,
+        bitset: u32,
+        count: usize,
+    ) -> InterpResult<'tcx, usize> {
         let this = self.eval_context_mut();
         let mut futex = futex_ref.0.borrow_mut();
         let data_race = &this.machine.data_race;
@@ -794,13 +799,18 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             data_race.release_clock(&this.machine.threads, |clock| futex.clock.clone_from(clock));
         }
 
-        // Wake up the first thread in the queue that matches any of the bits in the bitset.
-        let Some(i) = futex.waiters.iter().position(|w| w.bitset & bitset != 0) else {
-            return interp_ok(false);
-        };
-        let waiter = futex.waiters.remove(i).unwrap();
+        // Remove `count` of the threads in the queue that match any of the bits in the bitset.
+        // We collect all of them before unblocking because the unblock callback may access the
+        // futex state to retrieve the remaining number of waiters on macOS.
+        let waiters: Vec<_> =
+            futex.waiters.extract_if(.., |w| w.bitset & bitset != 0).take(count).collect();
         drop(futex);
-        this.unblock_thread(waiter.thread, BlockReason::Futex)?;
-        interp_ok(true)
+
+        let woken = waiters.len();
+        for waiter in waiters {
+            this.unblock_thread(waiter.thread, BlockReason::Futex)?;
+        }
+
+        interp_ok(woken)
     }
 }

src/lib.rs

@@ -16,6 +16,7 @@
 #![feature(derive_coerce_pointee)]
 #![feature(arbitrary_self_types)]
 #![feature(unsigned_is_multiple_of)]
+#![feature(extract_if)]
 // Configure clippy and other lints
 #![allow(
     clippy::collapsible_else_if,

src/shims/unix/linux_like/sync.rs

@@ -215,16 +215,8 @@ pub fn futex<'tcx>(
             // will see the latest value on addr which could be changed by our caller
             // before doing the syscall.
             ecx.atomic_fence(AtomicFenceOrd::SeqCst)?;
-            let mut n = 0;
-            #[expect(clippy::arithmetic_side_effects)]
-            for _ in 0..val {
-                if ecx.futex_wake(&futex_ref, bitset)? {
-                    n += 1;
-                } else {
-                    break;
-                }
-            }
-            ecx.write_scalar(Scalar::from_target_isize(n, ecx), dest)?;
+            let woken = ecx.futex_wake(&futex_ref, bitset, val.try_into().unwrap())?;
+            ecx.write_scalar(Scalar::from_target_isize(woken.try_into().unwrap(), ecx), dest)?;
         }
         op => throw_unsup_format!("Miri does not support `futex` syscall with op={}", op),
     }

src/shims/unix/macos/sync.rs

@@ -257,13 +257,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
         // See the Linux futex implementation for why this fence exists.
         this.atomic_fence(AtomicFenceOrd::SeqCst)?;
-
-        if all {
-            while this.futex_wake(&futex_ref, u32::MAX)? {}
-        } else {
-            this.futex_wake(&futex_ref, u32::MAX)?;
-        }
-
+        this.futex_wake(&futex_ref, u32::MAX, if all { usize::MAX } else { 1 })?;
         this.write_scalar(Scalar::from_i32(0), dest)?;
         interp_ok(())
     }

src/shims/windows/sync.rs

@@ -255,7 +255,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         };
         let futex_ref = futex_ref.futex.clone();
 
-        this.futex_wake(&futex_ref, u32::MAX)?;
+        this.futex_wake(&futex_ref, u32::MAX, 1)?;
 
         interp_ok(())
     }
@@ -275,7 +275,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         };
         let futex_ref = futex_ref.futex.clone();
 
-        while this.futex_wake(&futex_ref, u32::MAX)? {}
+        this.futex_wake(&futex_ref, u32::MAX, usize::MAX)?;
 
         interp_ok(())
     }

tests/pass-dep/concurrency/apple-futex.rs

@@ -149,6 +149,64 @@ fn wait_wake() {
     t.join().unwrap();
 }
 
+fn wait_wake_multiple() {
+    let val = 0i32;
+    let futex = &val;
+
+    thread::scope(|s| {
+        // Spawn some threads and make them wait on the futex.
+        for i in 0..4 {
+            s.spawn(move || unsafe {
+                assert_eq!(
+                    libc::os_sync_wait_on_address(
+                        ptr::from_ref(futex).cast_mut().cast(),
+                        0,
+                        size_of::<i32>(),
+                        libc::OS_SYNC_WAIT_ON_ADDRESS_NONE,
+                    ),
+                    // The last two threads will be woken at the same time,
+                    // but for the first two threads the remaining number
+                    // of waiters should be strictly decreasing.
+                    if i < 2 { 3 - i } else { 0 },
+                );
+            });
+
+            thread::sleep(Duration::from_millis(200));
+        }
+
+        // Wake the threads up again.
+        unsafe {
+            assert_eq!(
+                libc::os_sync_wake_by_address_any(
+                    ptr::from_ref(futex).cast_mut().cast(),
+                    size_of::<i32>(),
+                    libc::OS_SYNC_WAKE_BY_ADDRESS_NONE,
+                ),
+                0
+            );
+
+            assert_eq!(
+                libc::os_sync_wake_by_address_any(
+                    ptr::from_ref(futex).cast_mut().cast(),
+                    size_of::<i32>(),
+                    libc::OS_SYNC_WAKE_BY_ADDRESS_NONE,
+                ),
+                0
+            );
+
+            // Wake both remaining threads at the same time.
+            assert_eq!(
+                libc::os_sync_wake_by_address_all(
+                    ptr::from_ref(futex).cast_mut().cast(),
+                    size_of::<i32>(),
+                    libc::OS_SYNC_WAKE_BY_ADDRESS_NONE,
+                ),
+                0
+            );
+        }
+    })
+}
+
 fn param_mismatch() {
     let futex = 0;
     thread::scope(|s| {
@@ -209,5 +267,6 @@ fn main() {
     wait_timeout();
     wait_absolute_timeout();
     wait_wake();
+    wait_wake_multiple();
     param_mismatch();
 }