store futexes in per-allocation data rather than globally

Author: RalfJung

src/concurrency/sync.rs

@@ -1,6 +1,8 @@
+use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::collections::hash_map::Entry;
 use std::ops::Not;
+use std::rc::Rc;
 use std::time::Duration;
 
 use rustc_abi::Size;
@@ -121,6 +123,15 @@ struct Futex {
     clock: VClock,
 }
 
+#[derive(Default, Clone)]
+pub struct FutexRef(Rc<RefCell<Futex>>);
+
+impl VisitProvenance for FutexRef {
+    fn visit_provenance(&self, _visit: &mut VisitWith<'_>) {
+        // No provenance
+    }
+}
+
 /// A thread waiting on a futex.
 #[derive(Debug)]
 struct FutexWaiter {
@@ -137,9 +148,6 @@ pub struct SynchronizationObjects {
     rwlocks: IndexVec<RwLockId, RwLock>,
     condvars: IndexVec<CondvarId, Condvar>,
     pub(super) init_onces: IndexVec<InitOnceId, InitOnce>,
-
-    /// Futex info for the futex at the given address.
-    futexes: FxHashMap<u64, Futex>,
 }
 
 // Private extension trait for local helper methods
@@ -184,7 +192,7 @@ impl SynchronizationObjects {
 }
 
 impl<'tcx> AllocExtra<'tcx> {
-    pub fn get_sync<T: 'static>(&self, offset: Size) -> Option<&T> {
+    fn get_sync<T: 'static>(&self, offset: Size) -> Option<&T> {
         self.sync.get(&offset).and_then(|s| s.downcast_ref::<T>())
     }
 }
@@ -286,7 +294,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         // is truly the only place where the data could be stored.
         this.check_ptr_access(ptr, Size::from_bytes(1), CheckInAllocMsg::InboundsTest)?;
 
-        let (alloc, offset, _) = this.ptr_get_alloc_id(ptr, 0)?;
+        let (alloc, offset, _) = this.ptr_get_alloc_id(ptr, 0).unwrap();
         let (alloc_extra, machine) = this.get_alloc_extra_mut(alloc)?;
         // Due to borrow checker reasons, we have to do the lookup twice.
         if alloc_extra.get_sync::<T>(offset).is_none() {
@@ -296,6 +304,31 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         interp_ok(alloc_extra.get_sync::<T>(offset).unwrap())
     }
 
+    /// If the pointer is inbounds of an allocation, get the synchronization primitive associated
+    /// with the given pointer, or initialize a new one.
+    ///
+    /// Otherwise, return `None`.
+    fn try_get_sync_or_init<'a, T: 'static>(
+        &'a mut self,
+        ptr: Pointer,
+        new: impl FnOnce(&'a mut MiriMachine<'tcx>) -> InterpResult<'tcx, T>,
+    ) -> Option<&'a T>
+    where
+        'tcx: 'a,
+    {
+        let this = self.eval_context_mut();
+        if this.ptr_try_get_alloc_id(ptr, 0).ok().is_some_and(|(alloc_id, offset, ..)| {
+            let (size, _align, kind) = this.get_alloc_info(alloc_id);
+            // FIXME: we should also check mutability
+            kind == AllocKind::LiveData && offset < size
+        }) {
+            // This cannot fail now.
+            Some(this.get_sync_or_init(ptr, new).unwrap())
+        } else {
+            None
+        }
+    }
+
     #[inline]
     /// Get the id of the thread that currently owns this lock.
     fn mutex_get_owner(&mut self, id: MutexId) -> ThreadId {
@@ -690,7 +723,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     /// On a timeout, `retval_timeout` is written to `dest` and `errno_timeout` is set as the last error.
     fn futex_wait(
         &mut self,
-        addr: u64,
+        futex_ref: FutexRef,
         bitset: u32,
         timeout: Option<(TimeoutClock, TimeoutAnchor, Duration)>,
         retval_succ: Scalar,
@@ -700,23 +733,25 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     ) {
         let this = self.eval_context_mut();
         let thread = this.active_thread();
-        let futex = &mut this.machine.sync.futexes.entry(addr).or_default();
+        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 });
+        drop(futex);
+
         this.block_thread(
-            BlockReason::Futex { addr },
+            BlockReason::Futex,
             timeout,
             callback!(
                 @capture<'tcx> {
-                    addr: u64,
+                    futex_ref: FutexRef,
                     retval_succ: Scalar,
                     retval_timeout: Scalar,
                     dest: MPlaceTy<'tcx>,
                     errno_timeout: IoError,
                 }
                 @unblock = |this| {
-                    let futex = this.machine.sync.futexes.get(&addr).unwrap();
+                    let futex = futex_ref.0.borrow();
                     // Acquire the clock of the futex.
                     if let Some(data_race) = &this.machine.data_race {
                         data_race.acquire_clock(&futex.clock, &this.machine.threads);
@@ -728,7 +763,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 @timeout = |this| {
                     // Remove the waiter from the futex.
                     let thread = this.active_thread();
-                    let futex = this.machine.sync.futexes.get_mut(&addr).unwrap();
+                    let mut futex = futex_ref.0.borrow_mut();
                     futex.waiters.retain(|waiter| waiter.thread != thread);
                     // Set errno and write return value.
                     this.set_last_error(errno_timeout)?;
@@ -739,12 +774,11 @@ 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, addr: u64, bitset: u32) -> InterpResult<'tcx, bool> {
+    fn futex_wake(&mut self, futex_ref: &FutexRef, bitset: u32) -> InterpResult<'tcx, bool> {
         let this = self.eval_context_mut();
-        let Some(futex) = this.machine.sync.futexes.get_mut(&addr) else {
-            return interp_ok(false);
-        };
+        let mut futex = futex_ref.0.borrow_mut();
         let data_race = &this.machine.data_race;
 
         // Each futex-wake happens-before the end of the futex wait
@@ -757,7 +791,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             return interp_ok(false);
         };
         let waiter = futex.waiters.remove(i).unwrap();
-        this.unblock_thread(waiter.thread, BlockReason::Futex { addr })?;
+        drop(futex);
+        this.unblock_thread(waiter.thread, BlockReason::Futex)?;
         interp_ok(true)
     }
 }

src/concurrency/thread.rs

@@ -147,7 +147,7 @@ pub enum BlockReason {
     /// Blocked on a reader-writer lock.
     RwLock(RwLockId),
     /// Blocked on a Futex variable.
-    Futex { addr: u64 },
+    Futex,
     /// Blocked on an InitOnce.
     InitOnce(InitOnceId),
     /// Blocked on epoll.

src/shims/unix/linux/sync.rs

@@ -1,6 +1,11 @@
+use crate::concurrency::sync::FutexRef;
 use crate::helpers::check_min_arg_count;
 use crate::*;
 
+struct LinuxFutex {
+    futex: FutexRef,
+}
+
 /// Implementation of the SYS_futex syscall.
 /// `args` is the arguments *including* the syscall number.
 pub fn futex<'tcx>(
@@ -27,7 +32,6 @@ pub fn futex<'tcx>(
 
     // This is a vararg function so we have to bring our own type for this pointer.
     let addr = this.ptr_to_mplace(addr, this.machine.layouts.i32);
-    let addr_usize = addr.ptr().addr().bytes();
 
     let futex_private = this.eval_libc_i32("FUTEX_PRIVATE_FLAG");
     let futex_wait = this.eval_libc_i32("FUTEX_WAIT");
@@ -50,6 +54,15 @@ pub fn futex<'tcx>(
         op if op & !futex_realtime == futex_wait || op & !futex_realtime == futex_wait_bitset => {
             let wait_bitset = op & !futex_realtime == futex_wait_bitset;
 
+            // Storing this inside the allocation means that if an address gets reused by a new allocation,
+            // we'll use an independent futex queue for this... that seems acceptable.
+            let futex_ref = this
+                .get_sync_or_init(addr.ptr(), |_| {
+                    interp_ok(LinuxFutex { futex: Default::default() })
+                })?
+                .futex
+                .clone();
+
             let (timeout, bitset) = if wait_bitset {
                 let [_, _, _, _, timeout, uaddr2, bitset] =
                     check_min_arg_count("`syscall(SYS_futex, FUTEX_WAIT_BITSET, ...)`", args)?;
@@ -63,8 +76,7 @@ pub fn futex<'tcx>(
             };
 
             if bitset == 0 {
-                this.set_last_error_and_return(LibcError("EINVAL"), dest)?;
-                return interp_ok(());
+                return this.set_last_error_and_return(LibcError("EINVAL"), dest);
             }
 
             let timeout = this.deref_pointer_as(timeout, this.libc_ty_layout("timespec"))?;
@@ -144,7 +156,7 @@ pub fn futex<'tcx>(
             if val == futex_val {
                 // The value still matches, so we block the thread and make it wait for FUTEX_WAKE.
                 this.futex_wait(
-                    addr_usize,
+                    futex_ref,
                     bitset,
                     timeout,
                     Scalar::from_target_isize(0, this), // retval_succ
@@ -165,6 +177,15 @@ pub fn futex<'tcx>(
         // FUTEX_WAKE_BITSET: (int *addr, int op = FUTEX_WAKE, int val, const timespect *_unused, int *_unused, unsigned int bitset)
         // Same as FUTEX_WAKE, but allows you to specify a bitset to select which threads to wake up.
         op if op == futex_wake || op == futex_wake_bitset => {
+            let Some(futex_ref) = this.try_get_sync_or_init(addr.ptr(), |_| {
+                interp_ok(LinuxFutex { futex: Default::default() })
+            }) else {
+                // No AllocId, or no live allocation at that AllocId.
+                // Return an error code. (That seems nicer than silently doing something non-intuitive.)
+                return this.set_last_error_and_return(LibcError("EFAULT"), dest);
+            };
+            let futex_ref = futex_ref.futex.clone();
+
             let bitset = if op == futex_wake_bitset {
                 let [_, _, _, _, timeout, uaddr2, bitset] =
                     check_min_arg_count("`syscall(SYS_futex, FUTEX_WAKE_BITSET, ...)`", args)?;
@@ -184,7 +205,7 @@ pub fn futex<'tcx>(
             let mut n = 0;
             #[expect(clippy::arithmetic_side_effects)]
             for _ in 0..val {
-                if this.futex_wake(addr_usize, bitset)? {
+                if this.futex_wake(&futex_ref, bitset)? {
                     n += 1;
                 } else {
                     break;

src/shims/windows/sync.rs

@@ -3,13 +3,18 @@ use std::time::Duration;
 use rustc_abi::Size;
 
 use crate::concurrency::init_once::InitOnceStatus;
+use crate::concurrency::sync::FutexRef;
 use crate::*;
 
 #[derive(Copy, Clone)]
 struct WindowsInitOnce {
     id: InitOnceId,
 }
 
+struct WindowsFutex {
+    futex: FutexRef,
+}
+
 impl<'tcx> EvalContextExtPriv<'tcx> for crate::MiriInterpCx<'tcx> {}
 trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
     // Windows sync primitives are pointer sized.
@@ -168,7 +173,10 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let size = this.read_target_usize(size_op)?;
         let timeout_ms = this.read_scalar(timeout_op)?.to_u32()?;
 
-        let addr = ptr.addr().bytes();
+        let futex_ref = this
+            .get_sync_or_init(ptr, |_| interp_ok(WindowsFutex { futex: Default::default() }))?
+            .futex
+            .clone();
 
         if size > 8 || !size.is_power_of_two() {
             let invalid_param = this.eval_windows("c", "ERROR_INVALID_PARAMETER");
@@ -196,7 +204,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         if futex_val == compare_val {
             // If the values are the same, we have to block.
             this.futex_wait(
-                addr,
+                futex_ref,
                 u32::MAX, // bitset
                 timeout,
                 Scalar::from_i32(1), // retval_succ
@@ -219,8 +227,14 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         // See the Linux futex implementation for why this fence exists.
         this.atomic_fence(AtomicFenceOrd::SeqCst)?;
 
-        let addr = ptr.addr().bytes();
-        this.futex_wake(addr, u32::MAX)?;
+        let Some(futex_ref) = this
+            .try_get_sync_or_init(ptr, |_| interp_ok(WindowsFutex { futex: Default::default() }))
+        else {
+            // Seems like this cannot return an error, so we just wake nobody.
+            return interp_ok(());
+        };
+        let futex_ref = futex_ref.futex.clone();
+        this.futex_wake(&futex_ref, u32::MAX)?;
 
         interp_ok(())
     }
@@ -232,8 +246,14 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         // See the Linux futex implementation for why this fence exists.
         this.atomic_fence(AtomicFenceOrd::SeqCst)?;
 
-        let addr = ptr.addr().bytes();
-        while this.futex_wake(addr, u32::MAX)? {}
+        let Some(futex_ref) = this
+            .try_get_sync_or_init(ptr, |_| interp_ok(WindowsFutex { futex: Default::default() }))
+        else {
+            // Seems like this cannot return an error, so we just wake nobody.
+            return interp_ok(());
+        };
+        let futex_ref = futex_ref.futex.clone();
+        while this.futex_wake(&futex_ref, u32::MAX)? {}
 
         interp_ok(())
     }

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

@@ -41,9 +41,10 @@ fn wake_dangling() {
     let ptr: *const i32 = &*futex;
     drop(futex);
 
-    // Wake 1 waiter. Expect zero waiters woken up, as nobody is waiting.
+    // Expect error since this is now "unmapped" memory.
     unsafe {
-        assert_eq!(libc::syscall(libc::SYS_futex, ptr, libc::FUTEX_WAKE, 1), 0);
+        assert_eq!(libc::syscall(libc::SYS_futex, ptr, libc::FUTEX_WAKE, 1), -1);
+        assert_eq!(io::Error::last_os_error().raw_os_error().unwrap(), libc::EFAULT);
     }
 }