Experiment with loom

litebox/Cargo.toml

@@ -20,6 +20,7 @@ buddy_system_allocator = { version = "0.11.0", default-features = false, feature
 # Depend on (currently unreleased) slabmalloc `main`, which contains some fixes on top of `0.11.0`
 slabmalloc = { git = "https://github.com/gz/rust-slabmalloc.git", rev = "19480b2e82704210abafe575fb9699184c1be110" }
 litebox_util_log = { version = "0.1.0", path = "../litebox_util_log" }
+loom = { version = "0.7", optional = true }
 
 [target.'cfg(windows)'.dependencies]
 windows-sys = { version = "0.60.2", features = [
@@ -34,6 +35,7 @@ windows-sys = { version = "0.60.2", features = [
 lock_tracing = ["dep:arrayvec", "spin/mutex"]
 panic_on_unclosed_fd_drop = []
 enforce_singleton_litebox_instance = []
+loom = ["dep:loom"]
 
 [lints]
 workspace = true

litebox/src/event/observer.rs

@@ -5,7 +5,11 @@
 
 use alloc::collections::btree_map::BTreeMap;
 use alloc::sync::{Arc, Weak};
+
+#[cfg(not(feature = "loom"))]
 use core::sync::atomic::{AtomicUsize, Ordering};
+#[cfg(feature = "loom")]
+use loom::sync::atomic::{AtomicUsize, Ordering};
 
 use crate::sync::{Mutex, RawSyncPrimitivesProvider};
 

litebox/src/event/polling.rs

@@ -3,11 +3,14 @@
 
 //! Polling-related functionality
 
-use core::sync::atomic::AtomicBool;
-
 use alloc::sync::{Arc, Weak};
 use thiserror::Error;
 
+#[cfg(not(feature = "loom"))]
+use core::sync::atomic::{AtomicBool, Ordering};
+#[cfg(feature = "loom")]
+use loom::sync::atomic::{AtomicBool, Ordering};
+
 use super::{
     Events,
     observer::{Observer, Subject},
@@ -158,19 +161,99 @@ impl<Platform: RawSyncPrimitivesProvider> PolleeObserver<Platform> {
     }
 
     fn reset(&self) {
-        self.ready
-            .store(false, core::sync::atomic::Ordering::SeqCst);
+        self.ready.store(false, Ordering::SeqCst);
     }
 
     fn is_ready(&self) -> bool {
-        self.ready.load(core::sync::atomic::Ordering::SeqCst)
+        self.ready.load(Ordering::Acquire)
     }
 }
 
 impl<Platform: RawSyncPrimitivesProvider> Observer<Events> for PolleeObserver<Platform> {
     fn on_events(&self, _events: &Events) {
-        self.ready
-            .store(true, core::sync::atomic::Ordering::Release);
+        self.ready.store(true, Ordering::SeqCst);
+        #[cfg(feature = "loom")]
+        loom::sync::atomic::fence(Ordering::SeqCst);
         self.waker.wake();
     }
 }
+
+#[cfg(all(test, feature = "loom"))]
+mod loom_tests {
+    use alloc::boxed::Box;
+
+    use super::{Pollee, TryOpError};
+    use crate::event::{Events, wait::WaitState};
+    use crate::platform::loom_model::{Arc, LoomPlatform, atomic};
+
+    fn model(f: impl Fn() + Send + Sync + 'static) {
+        let mut builder = loom::model::Builder::new();
+        builder.preemption_bound = Some(1);
+        builder.check(f);
+    }
+
+    fn platform() -> &'static LoomPlatform {
+        Box::leak(Box::new(LoomPlatform::new()))
+    }
+
+    #[test]
+    fn wait_does_not_miss_notification() {
+        model(|| {
+            let pollee = Arc::new(Pollee::<LoomPlatform>::new());
+            let ready = Arc::new(loom::sync::Mutex::new(false));
+            let registered = Arc::new(atomic::AtomicBool::new(false));
+
+            let waiter = {
+                let pollee = Arc::clone(&pollee);
+                let ready = Arc::clone(&ready);
+                let registered = Arc::clone(&registered);
+                loom::thread::spawn(move || {
+                    let wait_state = WaitState::new(platform());
+                    wait_state.context().wait_on_events(
+                        false,
+                        Events::IN,
+                        |observer, filter| {
+                            pollee.register_observer(observer, filter);
+                            registered.store(true, atomic::Ordering::SeqCst);
+                            Ok(())
+                        },
+                        || {
+                            if *ready.lock().unwrap() {
+                                Ok(())
+                            } else {
+                                Err(TryOpError::<()>::TryAgain)
+                            }
+                        },
+                    )
+                })
+            };
+
+            let notifier = loom::thread::spawn(move || {
+                while !registered.load(atomic::Ordering::SeqCst) {
+                    loom::thread::yield_now();
+                }
+                *ready.lock().unwrap() = true;
+                pollee.notify_observers(Events::IN);
+            });
+
+            waiter.join().unwrap().unwrap();
+            notifier.join().unwrap();
+        });
+    }
+
+    #[test]
+    fn nonblock_does_not_register_observer() {
+        model(|| {
+            let pollee = Pollee::<LoomPlatform>::new();
+            let wait_state = WaitState::new(platform());
+
+            let result: Result<(), TryOpError<()>> =
+                pollee.wait(&wait_state.context(), true, Events::IN, || {
+                    Err(TryOpError::<()>::TryAgain)
+                });
+
+            assert!(matches!(result, Err(TryOpError::TryAgain)));
+            pollee.notify_observers(Events::IN);
+        });
+    }
+}

litebox/src/event/wait.rs

@@ -25,6 +25,11 @@
 use alloc::sync::Arc;
 use core::{marker::PhantomData, sync::atomic::Ordering};
 
+#[cfg(not(feature = "loom"))]
+use core::sync::atomic::fence;
+#[cfg(feature = "loom")]
+use loom::sync::atomic::fence;
+
 use crate::{
     platform::{
         ImmediatelyWokenUp, Instant as _, RawMutex, ThreadProvider, TimeProvider,
@@ -83,7 +88,7 @@ impl<Platform: RawSyncPrimitivesProvider> WaitState<Platform> {
         Self {
             waker: Waker(Arc::new(WaitStateInner {
                 platform,
-                condvar: <Platform::RawMutex as RawMutex>::INIT,
+                condvar: <Platform::RawMutex as RawMutex>::new(),
             })),
             _phantom: PhantomData,
         }
@@ -110,6 +115,11 @@ impl<Platform: RawSyncPrimitivesProvider> WaitState<Platform> {
         }
     }
 
+    #[cfg(all(test, feature = "loom"))]
+    pub(crate) fn is_running_in_host_for_test(&self) -> bool {
+        self.waker.0.state_for_assert() == ThreadState::RUNNING_IN_HOST
+    }
+
     /// Sets the wait state so that [`ThreadHandle::interrupt`] will interrupt
     /// the guest execution, then calls `f` to see if the guest is still ready
     /// to run.
@@ -132,6 +142,9 @@ impl<Platform: RawSyncPrimitivesProvider> WaitState<Platform> {
         self.waker
             .0
             .set_state(ThreadState::RUNNING_IN_GUEST, Ordering::SeqCst);
+        // loom does not support SeqCst for load/store: see https://github.com/tokio-rs/loom/issues/180
+        #[cfg(feature = "loom")]
+        loom::sync::atomic::fence(Ordering::SeqCst);
         let ready_to_run_guest = f();
         if !ready_to_run_guest {
             self.waker
@@ -182,7 +195,7 @@ impl<Platform: RawSyncPrimitivesProvider> WaitStateInner<Platform> {
             Err(_) => {
                 // Provide a consistent release fence even if we didn't wake up
                 // the thread.
-                core::sync::atomic::fence(Ordering::Release);
+                fence(Ordering::Release);
             }
         }
     }
@@ -192,9 +205,13 @@ impl<Platform: RawSyncPrimitivesProvider> WaitStateInner<Platform> {
     }
 
     fn set_state(&self, new_state: ThreadState, ordering: Ordering) {
+        #[cfg(not(feature = "loom"))]
         self.condvar
             .underlying_atomic()
             .store(new_state.0, ordering);
+        // See test `relaxed_load_does_not_observe_own_relaxed_store`
+        #[cfg(feature = "loom")]
+        let _ = self.condvar.underlying_atomic().swap(new_state.0, ordering);
     }
 }
 
@@ -235,7 +252,7 @@ impl<Platform: RawSyncPrimitivesProvider + ThreadProvider> ThreadHandle<Platform
             Err(_) => {
                 // Provide a consistent release fence even if we didn't wake up
                 // the thread.
-                core::sync::atomic::fence(Ordering::Release);
+                fence(Ordering::Release);
             }
         }
     }
@@ -292,7 +309,7 @@ pub trait CheckForInterrupt {
     ///
     /// This is called by [`WaitContext::wait_until`] each time it is about to
     /// block the thread. If this returns `true`, the wait will return with
-    /// [`WaitError::Interrupted`].
+    /// [`WaitError::Interrupted`], which in turn is converted into the wait's error type.
     fn check_for_interrupt(&self) -> bool;
 }
 
@@ -378,6 +395,9 @@ impl<'a, Platform: RawSyncPrimitivesProvider + TimeProvider> WaitContext<'a, Pla
         self.waker
             .0
             .set_state(ThreadState::WAITING, Ordering::SeqCst);
+        // loom does not support SeqCst for load/store: see https://github.com/tokio-rs/loom/issues/180
+        #[cfg(feature = "loom")]
+        loom::sync::atomic::fence(Ordering::SeqCst);
     }
 
     /// Returns the thread to the running state after a wait.