Add futex wake-path store-buffering reproduction

litebox/Cargo.toml

@@ -34,6 +34,10 @@ windows-sys = { version = "0.60.2", features = [
 lock_tracing = ["dep:arrayvec", "spin/mutex"]
 panic_on_unclosed_fd_drop = []
 enforce_singleton_litebox_instance = []
+# Exposes the futex wake-path ordering stress instrumentation (see
+# `src/ordering_stress.rs`) to dependent crates so they can drive the reproduction
+# over a real platform. Test-only; never enable in production builds.
+futex_ordering_stress = []
 
 [lints]
 workspace = true

litebox/src/event/wait.rs

@@ -174,6 +174,8 @@ impl<Platform: RawSyncPrimitivesProvider> WaitStateInner<Platform> {
                 state => unreachable!("{state:?}"),
             },
         );
+        #[cfg(any(test, feature = "futex_ordering_stress"))]
+        crate::ordering_stress::record_waker_result(v);
         match v.map(ThreadState) {
             Ok(ThreadState::WAITING) => {
                 condvar.wake_one();
@@ -375,6 +377,8 @@ impl<'a, Platform: RawSyncPrimitivesProvider + TimeProvider> WaitContext<'a, Pla
     /// missed.
     fn start_wait(&self) {
         self.waker.0.platform.update_waker(Some(self.waker.clone()));
+        #[cfg(any(test, feature = "futex_ordering_stress"))]
+        crate::ordering_stress::waiter_rendezvous();
         self.waker
             .0
             .set_state(ThreadState::WAITING, Ordering::SeqCst);

litebox/src/lib.rs

@@ -28,6 +28,13 @@ pub mod shim;
 pub mod sync;
 pub mod tls;
 
+// Test-only instrumentation for the futex wake-path store-buffering stress test.
+// Compiled under `cfg(test)` for the in-crate mock test, or the
+// `futex_ordering_stress` feature so a dependent crate can drive the same
+// reproduction over a real platform. Never enable in production builds.
+#[cfg(any(test, feature = "futex_ordering_stress"))]
+pub mod ordering_stress;
+
 // The core [`LiteBox`] object itself, re-exported here publicly, just to keep management of the
 // code cleaner.
 mod litebox;

litebox/src/ordering_stress.rs

@@ -0,0 +1,130 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+//! Test-only instrumentation for the futex wake-path store-buffering stress test.
+//!
+//! This module exposes a small rendezvous protocol that a stress harness uses to align
+//! a waiter and a waker at the exact instant of the two conflicting stores in the futex
+//! wake path (the waiter's `WAITING` store in [`WaitContext::start_wait`] and the waker's
+//! relaxed `done` store in [`FutexManager::wake`]), and to observe whether each side read
+//! the other's stale value.
+//!
+//! It is compiled only under `cfg(test)` (for the in-crate mock-platform test) or the
+//! `futex_ordering_stress` feature (so a dependent crate can drive the same protocol over
+//! a real platform). It must never be enabled in production builds: the hooks add a
+//! rendezvous barrier inside the live wait/wake paths.
+//!
+//! [`WaitContext::start_wait`]: crate::event::wait
+//! [`FutexManager::wake`]: crate::sync::futex::FutexManager::wake
+
+use core::sync::atomic::{AtomicBool, AtomicU32, Ordering};
+
+/// The number of participants (waiter + waker) that must park before the harness releases them.
+const PARTICIPANTS: u32 = 2;
+
+/// Sentinel meaning "not yet recorded this round".
+const UNSET: u32 = u32::MAX;
+
+static ACTIVE: AtomicBool = AtomicBool::new(false);
+static PARKED: AtomicU32 = AtomicU32::new(0);
+static RELEASE: AtomicBool = AtomicBool::new(false);
+static WAITER_REGISTERED: AtomicBool = AtomicBool::new(false);
+static WAITER_DONE: AtomicU32 = AtomicU32::new(UNSET);
+static WAKER_RESULT: AtomicU32 = AtomicU32::new(UNSET);
+
+/// Enables the instrumentation hooks. Call once before a stress run.
+pub fn activate() {
+    ACTIVE.store(true, Ordering::Relaxed);
+}
+
+/// Disables the instrumentation hooks. Call once after a stress run.
+pub fn deactivate() {
+    ACTIVE.store(false, Ordering::Relaxed);
+}
+
+/// Resets all per-round observation state. Call at the top of each iteration before
+/// releasing the two threads.
+pub fn begin_round() {
+    PARKED.store(0, Ordering::Relaxed);
+    RELEASE.store(false, Ordering::Relaxed);
+    WAITER_REGISTERED.store(false, Ordering::Relaxed);
+    WAITER_DONE.store(UNSET, Ordering::Relaxed);
+    WAKER_RESULT.store(UNSET, Ordering::Relaxed);
+}
+
+/// Returns whether the waiter has inserted its entry (so the waker's `wake` will select it).
+#[must_use]
+pub fn waiter_is_registered() -> bool {
+    WAITER_REGISTERED.load(Ordering::Acquire)
+}
+
+/// Spins until both the waiter and the waker have parked at the rendezvous.
+pub fn wait_until_parked() {
+    while PARKED.load(Ordering::Acquire) != PARTICIPANTS {
+        core::hint::spin_loop();
+    }
+}
+
+/// Releases the parked waiter and waker together, so their two stores race.
+pub fn release() {
+    RELEASE.store(true, Ordering::Release);
+}
+
+/// Returns whether this round observed the both-old outcome: the waiter's first `done`
+/// load read `false` and the waker's `fetch_update` read `RUNNING_IN_HOST` (encoded 0).
+#[must_use]
+pub fn observed_both_old() -> bool {
+    WAITER_DONE.load(Ordering::Relaxed) == 0 && WAKER_RESULT.load(Ordering::Relaxed) == 0
+}
+
+/// Hook: the waiter has inserted its entry but has not yet parked.
+pub(crate) fn waiter_registered() {
+    if ACTIVE.load(Ordering::Relaxed) {
+        WAITER_REGISTERED.store(true, Ordering::Release);
+    }
+}
+
+/// Hook: called immediately before the waiter's `WAITING` store.
+pub(crate) fn waiter_rendezvous() {
+    if ACTIVE.load(Ordering::Relaxed) {
+        rendezvous();
+    }
+}
+
+/// Hook: called immediately before the waker's relaxed `done` store.
+pub(crate) fn waker_rendezvous() {
+    if ACTIVE.load(Ordering::Relaxed) {
+        rendezvous();
+    }
+}
+
+/// Hook: records the value of the waiter's first `done` load (first write wins).
+pub(crate) fn record_waiter_done(done: bool) {
+    if ACTIVE.load(Ordering::Relaxed) {
+        let _ = WAITER_DONE.compare_exchange(
+            UNSET,
+            u32::from(done),
+            Ordering::Relaxed,
+            Ordering::Relaxed,
+        );
+    }
+}
+
+/// Hook: records the encoded result of the waker's `fetch_update`. `Ok` results set the
+/// high bit so a failed `Err(RUNNING_IN_HOST)` (encoded 0) is distinguishable.
+pub(crate) fn record_waker_result(result: Result<u32, u32>) {
+    if ACTIVE.load(Ordering::Relaxed) {
+        let encoded = match result {
+            Ok(state) => state | (1 << 31),
+            Err(state) => state,
+        };
+        WAKER_RESULT.store(encoded, Ordering::Relaxed);
+    }
+}
+
+fn rendezvous() {
+    PARKED.fetch_add(1, Ordering::Release);
+    while !RELEASE.load(Ordering::Acquire) {
+        core::hint::spin_loop();
+    }
+}

litebox/src/sync/futex.rs

@@ -111,11 +111,20 @@ impl<Platform: RawSyncPrimitivesProvider + RawPointerProvider + TimeProvider>
         if value != expected_value {
             return Err(FutexError::ImmediatelyWokenBecauseValueMismatch);
         }
+
+        #[cfg(any(test, feature = "futex_ordering_stress"))]
+        crate::ordering_stress::waiter_registered();
+
         // Only return when woken--don't reevaluate the futex word. This
         // ensures that the rate control mechanisms provided by the futex
         // interface are effective.
-        cx.wait_until(|| entry.get().done.load(Ordering::Acquire))
-            .map_err(FutexError::WaitError)
+        cx.wait_until(|| {
+            let done = entry.get().done.load(Ordering::Acquire);
+            #[cfg(any(test, feature = "futex_ordering_stress"))]
+            crate::ordering_stress::record_waiter_done(done);
+            done
+        })
+        .map_err(FutexError::WaitError)
     }
 
     /// Wakes waiters on the given futex word.
@@ -160,6 +169,8 @@ impl<Platform: RawSyncPrimitivesProvider + RawPointerProvider + TimeProvider>
         // Wake the waiters outside the `extract_if` closure to minimize the list's lock hold
         // time.
         for entry in entries {
+            #[cfg(any(test, feature = "futex_ordering_stress"))]
+            crate::ordering_stress::waker_rendezvous();
             entry.done.store(true, Ordering::Relaxed);
             entry.waker.wake();
         }
@@ -186,11 +197,11 @@ mod tests {
 
     use super::*;
     use crate::LiteBox;
-    use crate::event::wait::WaitState;
+    use crate::event::wait::{WaitError, WaitState};
     use crate::platform::mock::MockPlatform;
     use alloc::sync::Arc;
     use core::num::NonZeroU32;
-    use core::sync::atomic::{AtomicU32, Ordering};
+    use core::sync::atomic::{AtomicBool, AtomicU32, Ordering};
     use std::sync::Barrier;
     use std::thread;
     use std::time::Duration;
@@ -358,4 +369,119 @@ mod tests {
 
         assert!((1..=3).contains(&woken));
     }
+
+    /// Reproduces the store-buffering hazard in the original futex wake path: the waker's
+    /// relaxed `done` store and the waiter's `SeqCst` `WAITING` store can each read the
+    /// other's stale value, so the waker's `fetch_update` sees `RUNNING_IN_HOST` and skips
+    /// the wake while the waiter blocks and times out. Ignored because it is probabilistic;
+    /// run with `LITEBOX_FUTEX_STRESS_ITERS` to control the iteration count.
+    #[test]
+    #[ignore = "probabilistic weak-memory stress test"]
+    fn stress_registered_waiter_does_not_miss_wake() {
+        let platform = MockPlatform::new();
+        let _litebox = LiteBox::new(platform);
+        let futex_manager = Arc::new(FutexManager::new());
+        let futex_word = Arc::new(AtomicU32::new(0));
+        let iterations = std::env::var("LITEBOX_FUTEX_STRESS_ITERS")
+            .ok()
+            .and_then(|value| value.parse().ok())
+            .unwrap_or(1_000_000);
+        crate::ordering_stress::activate();
+        let iteration_start = Arc::new(Barrier::new(3));
+        let iteration_finish = Arc::new(Barrier::new(3));
+        let waiter_result = Arc::new(AtomicU32::new(u32::MAX));
+        let was_selected = Arc::new(AtomicBool::new(false));
+        let waiter = {
+            let futex_manager = Arc::clone(&futex_manager);
+            let futex_word = Arc::clone(&futex_word);
+            let iteration_start = Arc::clone(&iteration_start);
+            let iteration_finish = Arc::clone(&iteration_finish);
+            let waiter_result = Arc::clone(&waiter_result);
+            thread::spawn(move || {
+                for _ in 0..iterations {
+                    iteration_start.wait();
+                    let futex_addr = <MockPlatform as crate::platform::RawPointerProvider>::RawMutPointer::from_usize(
+                        futex_word.as_ptr() as usize,
+                    );
+                    let result = futex_manager.wait(
+                        &WaitState::new(platform)
+                            .context()
+                            .with_timeout(Duration::from_millis(20)),
+                        futex_addr,
+                        0,
+                        None,
+                    );
+                    waiter_result.store(
+                        u32::from(matches!(
+                            result,
+                            Err(FutexError::WaitError(WaitError::TimedOut))
+                        )),
+                        Ordering::Relaxed,
+                    );
+                    iteration_finish.wait();
+                }
+            })
+        };
+        let waker = {
+            let futex_manager = Arc::clone(&futex_manager);
+            let futex_word = Arc::clone(&futex_word);
+            let iteration_start = Arc::clone(&iteration_start);
+            let iteration_finish = Arc::clone(&iteration_finish);
+            let was_selected = Arc::clone(&was_selected);
+            thread::spawn(move || {
+                for _ in 0..iterations {
+                    iteration_start.wait();
+                    while !crate::ordering_stress::waiter_is_registered() {
+                        core::hint::spin_loop();
+                    }
+                    let futex_addr = <MockPlatform as crate::platform::RawPointerProvider>::RawMutPointer::from_usize(
+                        futex_word.as_ptr() as usize,
+                    );
+                    was_selected.store(
+                        futex_manager
+                            .wake(futex_addr, NonZeroU32::new(1).unwrap(), None)
+                            .unwrap()
+                            == 1,
+                        Ordering::Relaxed,
+                    );
+                    iteration_finish.wait();
+                }
+            })
+        };
+        let mut selected = 0;
+        let mut both_old = 0;
+        let mut lost_wakeups = 0;
+
+        for _ in 0..iterations {
+            futex_word.store(0, Ordering::Relaxed);
+            crate::ordering_stress::begin_round();
+            waiter_result.store(u32::MAX, Ordering::Relaxed);
+            was_selected.store(false, Ordering::Relaxed);
+            iteration_start.wait();
+            crate::ordering_stress::wait_until_parked();
+            crate::ordering_stress::release();
+            iteration_finish.wait();
+            if was_selected.load(Ordering::Relaxed) {
+                selected += 1;
+                if crate::ordering_stress::observed_both_old() {
+                    both_old += 1;
+                }
+                if waiter_result.load(Ordering::Relaxed) == 1 {
+                    lost_wakeups += 1;
+                }
+            }
+        }
+
+        waiter.join().unwrap();
+        waker.join().unwrap();
+        crate::ordering_stress::deactivate();
+        std::eprintln!(
+            "iterations={iterations} selected={selected} both_old={both_old} lost_wakeups={lost_wakeups}"
+        );
+        assert_eq!(
+            both_old, 0,
+            "the original futex wake path saw both old values"
+        );
+        assert_eq!(lost_wakeups, 0, "a selected futex waiter missed its wake");
+    }
 }

litebox_platform_linux_userland/Cargo.toml

@@ -16,6 +16,15 @@ syscalls = { version = "0.6", default-features = false }
 zerocopy = { version = "0.8", default-features = false }
 seccompiler = { version = "0.5.0" }
 
+[dev-dependencies]
+# Re-declare litebox with the ordering-stress instrumentation enabled so the
+# real-platform futex reproduction test can drive the rendezvous hooks. The
+# feature is active only for this crate's own test builds, never for normal
+# builds or when this crate is used as a dependency.
+litebox = { path = "../litebox/", version = "0.1.0", features = [
+    "futex_ordering_stress",
+] }
+
 [features]
 default = ["linux_syscall"]
 linux_syscall = []