addr_monitor 0.3.0

//! Contains primitives for simple blocking until a memory location changes.
//!
//! On devices that support hardware monitors, these functions use them to wait in between
//! atomic loads. On unsupported platforms, they are just spin loops.
//!
//! Use these only when you expect to be waiting a long time. Otherwise, a simple spin loop
//! is a better idea. When there is an OS, parking and unparking the thread is usually more
//! power-efficient than blocking it with a hardware monitor like this. However, that can
//! be expensive and difficult to manage.

use core::sync::atomic::Ordering;

#[cfg(target_arch = "x86_64")]
use crate::hardware::{HardwareMonitor, MonitorControl, MonitorSupport, MonitorSupportBackend};

/// Conditionally waits for the atomic `value` to equal `desired`.
///
/// On devices that support hardware monitors, this function uses them to wait in between
/// atomic loads. On unsupported platforms, this is just a spin loop.
pub fn cond_wait_eq<T>(
    value: &T,
    desired: T::NonAtomic,
    control: MonitorControl,
    support: MonitorSupport,
) -> T::NonAtomic
where
    T: atomic::Atomic,
    T::NonAtomic: Eq,
{
    macro_rules! fallback {
        ($value:ident, $desired:ident) => {
            loop {
                let current = value.load(Ordering::Relaxed);

                if desired == current {
                    return current;
                }

                ::core::hint::spin_loop();
            }
        };
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        fallback!(value, desired);
    }

    #[cfg(target_arch = "x86_64")]
    {
        if matches!(support.inner, MonitorSupportBackend::None) {
            fallback!(value, desired);
        }

        let monitor = HardwareMonitor::new(value, support);

        loop {
            let current = value.load(Ordering::Relaxed);

            if desired == current {
                return current;
            }

            monitor.wait(control, support);
        }
    }
}

/// Conditionally waits for the atomic `value` to not equal `desired`.
///
/// On devices that support hardware monitors, this function uses them to wait in between
/// atomic loads. On unsupported platforms, this is just a spin loop.
pub fn cond_wait_ne<T>(
    value: &T,
    desired: T::NonAtomic,
    control: MonitorControl,
    support: MonitorSupport,
) where
    T: atomic::Atomic,
    T::NonAtomic: Eq,
{
    macro_rules! fallback {
        ($value:ident, $desired:ident) => {
            loop {
                let current = value.load(Ordering::Relaxed);

                if desired != current {
                    return;
                }

                ::core::hint::spin_loop();
            }
        };
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        fallback!(value, desired);
    }

    #[cfg(target_arch = "x86_64")]
    {
        if matches!(support.inner, MonitorSupportBackend::None) {
            fallback!(value, desired);
        }

        let monitor = HardwareMonitor::new(value, support);

        loop {
            let current = value.load(Ordering::Relaxed);

            if desired != current {
                return;
            }

            monitor.wait(control, support);
        }
    }
}

/// Blocks the current thread until the atomic value changes, returning the new value.
///
/// On devices that support hardware monitors, this function uses them to wait in between
/// atomic loads. On unsupported platforms, this is just a spin loop.
pub fn block_on_write<T>(
    value: &T,
    control: MonitorControl,
    support: MonitorSupport,
) -> T::NonAtomic
where
    T: atomic::Atomic,
    T::NonAtomic: Eq,
{
    macro_rules! fallback {
        ($value:ident) => {
            let current = value.load(Ordering::Relaxed);

            loop {
                let new = value.load(Ordering::Relaxed);

                if new != current {
                    return new;
                }

                ::core::hint::spin_loop();
            }
        };
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        fallback!(value);
    }

    #[cfg(target_arch = "x86_64")]
    {
        if matches!(support.inner, MonitorSupportBackend::None) {
            fallback!(value);
        }

        let monitor = HardwareMonitor::new(value, support);

        let current = value.load(Ordering::Relaxed);

        monitor.wait(MonitorControl::Power, support);

        let monitor = HardwareMonitor::new(value, support);

        loop {
            let new = value.load(Ordering::Relaxed);

            if new != current {
                return new;
            }

            monitor.wait(control, support);
        }
    }
}

/// Writes the given `new` value to `value` once it isn't racing.
///
/// Just an efficient CAS loop using hardware monitors instead of spin loops.
pub fn swap_on_write<T>(
    value: &T,
    new: T::NonAtomic,
    control: MonitorControl,
    support: MonitorSupport,
) -> T::NonAtomic
where
    T: atomic::Atomic,
    T::NonAtomic: Eq + Copy,
{
    macro_rules! fallback {
        ($value:ident, $new:ident) => {
            let mut current = value.load(Ordering::Acquire);

            loop {
                match $value.cas(current, $new, Ordering::AcqRel, Ordering::Acquire) {
                    Ok(prev) => return prev,
                    Err(actual) => current = actual,
                }

                ::core::hint::spin_loop();
            }
        };
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        fallback!(value, new);
    }

    #[cfg(target_arch = "x86_64")]
    {
        if matches!(support.inner, MonitorSupportBackend::None) {
            fallback!(value, new);
        }

        let monitor = HardwareMonitor::new(value, support);

        let mut current = value.load(Ordering::Acquire);

        monitor.wait(control, support);

        let monitor = HardwareMonitor::new(value, support);

        loop {
            match value.cas(current, new, Ordering::AcqRel, Ordering::Acquire) {
                Ok(prev) => return prev,
                Err(actual) => current = actual,
            }

            monitor.wait(control, support);
        }
    }
}

mod atomic {
    use core::sync::atomic::{AtomicPtr, Ordering};

    pub trait Atomic {
        type NonAtomic;

        fn load(&self, ordering: Ordering) -> Self::NonAtomic;
        fn cas(
            &self,
            current: Self::NonAtomic,
            new: Self::NonAtomic,
            success: Ordering,
            failure: Ordering,
        ) -> Result<Self::NonAtomic, Self::NonAtomic>;
    }

    macro_rules! impl_atomic {
        ($($atomic:ident=$non_atomic:ty=$bits:literal),*$(,)?) => {
            $(
                #[cfg(target_has_atomic = $bits)]
                impl Atomic for ::core::sync::atomic::$atomic {
                    type NonAtomic = $non_atomic;

                    fn load(&self, ordering: Ordering) -> Self::NonAtomic {
                        self.load(ordering)
                    }

                    fn cas(&self, current: Self::NonAtomic, new: Self::NonAtomic, success: Ordering, failure: Ordering) -> Result<Self::NonAtomic, Self::NonAtomic> {
                        self.compare_exchange_weak(current, new, success, failure)
                    }
                }
            )*
        };
    }

    impl_atomic!(
        AtomicBool = bool = "8",
        AtomicIsize = isize = "ptr",
        AtomicUsize = usize = "ptr",
        AtomicI8 = i8 = "8",
        AtomicI16 = i16 = "16",
        AtomicI32 = i32 = "32",
        AtomicI64 = i64 = "64",
        AtomicU8 = u8 = "8",
        AtomicU16 = u16 = "16",
        AtomicU32 = u32 = "32",
        AtomicU64 = u64 = "64",
    );

    #[cfg(target_has_atomic = "ptr")]
    impl<T> Atomic for AtomicPtr<T> {
        type NonAtomic = usize;

        fn load(&self, ordering: Ordering) -> Self::NonAtomic {
            self.load(ordering) as usize
        }

        fn cas(
            &self,
            current: Self::NonAtomic,
            new: Self::NonAtomic,
            success: Ordering,
            failure: Ordering,
        ) -> Result<Self::NonAtomic, Self::NonAtomic> {
            self.compare_exchange(current as *mut T, new as *mut T, success, failure)
                .map(|ptr| ptr as usize)
                .map_err(|ptr| ptr as usize)
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate std;

    use super::*;

    use std::{
        boxed::Box,
        sync::atomic::{AtomicBool, AtomicI32, AtomicPtr, Ordering},
        sync::Arc,
        thread,
        time::Duration,
        vec::Vec,
    };

    #[test]
    fn block_on_write_waits_until_atomic_bool_changes() {
        let value = Arc::new(AtomicBool::new(false));
        let watched = Arc::clone(&value);

        let handle = thread::spawn(move || {
            block_on_write(
                watched.as_ref(),
                MonitorControl::Power,
                MonitorSupport::platform(),
            );
        });

        thread::sleep(Duration::from_millis(50));
        assert!(!handle.is_finished());

        value.store(true, Ordering::Relaxed);

        handle.join().unwrap();
    }

    #[test]
    fn block_on_write_waits_until_atomic_i32_changes() {
        let value = Arc::new(AtomicI32::new(42));
        let watched = Arc::clone(&value);

        let handle = thread::spawn(move || {
            block_on_write(
                watched.as_ref(),
                MonitorControl::Power,
                MonitorSupport::platform(),
            );
        });

        thread::sleep(Duration::from_millis(50));
        assert!(!handle.is_finished());

        value.store(43, Ordering::Relaxed);

        handle.join().unwrap();
    }

    #[test]
    fn block_on_write_pointer_changes_are_detected() {
        let first = Box::into_raw(Box::new(1u8));
        let second = Box::into_raw(Box::new(2u8));

        let value = Arc::new(AtomicPtr::new(first));
        let watched = Arc::clone(&value);

        let handle = thread::spawn(move || {
            block_on_write(
                watched.as_ref(),
                MonitorControl::Power,
                MonitorSupport::platform(),
            );
        });

        thread::sleep(Duration::from_millis(50));
        assert!(!handle.is_finished());

        value.store(second, Ordering::Relaxed);

        handle.join().unwrap();

        unsafe {
            drop(Box::from_raw(first));
            drop(Box::from_raw(second));
        }
    }

    #[test]
    fn block_on_write_does_not_return_for_same_value() {
        let value = Arc::new(AtomicI32::new(123));
        let watched = Arc::clone(&value);

        let handle = thread::spawn(move || {
            block_on_write(
                watched.as_ref(),
                MonitorControl::Power,
                MonitorSupport::platform(),
            );
        });

        thread::sleep(Duration::from_millis(50));
        value.store(123, Ordering::Relaxed);

        thread::sleep(Duration::from_millis(50));
        assert!(!handle.is_finished());

        value.store(124, Ordering::Relaxed);

        handle.join().unwrap();
    }

    #[test]
    fn cond_wait_ne_returns_immediately_when_value_is_already_different() {
        let value = AtomicI32::new(42);

        cond_wait_ne(
            &value,
            41,
            MonitorControl::Power,
            MonitorSupport::platform(),
        );

        assert_eq!(value.load(Ordering::Relaxed), 42);
    }

    #[test]
    fn cond_wait_ne_waits_until_value_changes() {
        use std::sync::mpsc;

        let value = Arc::new(AtomicI32::new(42));
        let watched = Arc::clone(&value);

        let (started_tx, started_rx) = mpsc::channel();
        let (finished_tx, finished_rx) = mpsc::channel();

        let handle = thread::spawn(move || {
            started_tx.send(()).unwrap();

            cond_wait_ne(
                watched.as_ref(),
                42,
                MonitorControl::Power,
                MonitorSupport::platform(),
            );

            finished_tx.send(()).unwrap();
        });

        started_rx.recv().unwrap();

        assert!(finished_rx.try_recv().is_err());

        value.store(43, Ordering::Release);

        finished_rx
            .recv_timeout(Duration::from_secs(1))
            .expect("cond_wait_ne did not return after the value changed");

        handle.join().unwrap();
    }

    #[test]
    fn cond_wait_eq_returns_immediately_when_value_is_already_different() {
        let value = AtomicI32::new(42);

        cond_wait_eq(
            &value,
            42,
            MonitorControl::Power,
            MonitorSupport::platform(),
        );

        assert_eq!(value.load(Ordering::Relaxed), 42);
    }

    #[test]
    fn cond_wait_eq_waits_until_value_changes() {
        use std::sync::mpsc;

        let value = Arc::new(AtomicI32::new(42));
        let watched = Arc::clone(&value);

        let (started_tx, started_rx) = mpsc::channel();
        let (finished_tx, finished_rx) = mpsc::channel();

        let handle = thread::spawn(move || {
            started_tx.send(()).unwrap();

            cond_wait_eq(
                watched.as_ref(),
                43,
                MonitorControl::Power,
                MonitorSupport::platform(),
            );

            finished_tx.send(()).unwrap();
        });

        started_rx.recv().unwrap();

        assert!(finished_rx.try_recv().is_err());

        value.store(43, Ordering::Release);

        finished_rx
            .recv_timeout(Duration::from_secs(1))
            .expect("cond_wait_eq did not return after the value changed");

        handle.join().unwrap();
    }

    #[test]
    fn swap_on_write_returns_previous_value_and_stores_new_value() {
        let value = AtomicI32::new(42);

        let previous = swap_on_write(
            &value,
            123,
            MonitorControl::Power,
            MonitorSupport::platform(),
        );

        assert_eq!(previous, 42);
        assert_eq!(value.load(Ordering::Acquire), 123);
    }

    #[test]
    fn swap_on_write_can_write_the_same_value() {
        let value = AtomicI32::new(42);

        let previous = swap_on_write(
            &value,
            42,
            MonitorControl::Power,
            MonitorSupport::platform(),
        );

        assert_eq!(previous, 42);
        assert_eq!(value.load(Ordering::Acquire), 42);
    }

    #[test]
    fn swap_on_write_handles_racing_writers() {
        use std::sync::Barrier;

        const WRITERS: i32 = 8;

        let value = Arc::new(AtomicI32::new(0));
        let barrier = Arc::new(Barrier::new(WRITERS as usize));

        let mut handles = Vec::new();

        for new_value in 1..=WRITERS {
            let value = Arc::clone(&value);
            let barrier = Arc::clone(&barrier);

            handles.push(thread::spawn(move || {
                barrier.wait();

                let previous = swap_on_write(
                    value.as_ref(),
                    new_value,
                    MonitorControl::Power,
                    MonitorSupport::platform(),
                );

                (new_value, previous)
            }));
        }

        let results: Vec<(i32, i32)> = handles
            .into_iter()
            .map(|handle| handle.join().unwrap())
            .collect();

        let final_value = value.load(Ordering::Acquire);

        assert!((1..=WRITERS).contains(&final_value));

        assert_eq!(
            results
                .iter()
                .filter(|(_, previous)| *previous == 0)
                .count(),
            1
        );

        for (_, previous) in &results {
            assert!(
                *previous == 0 || (1..=WRITERS).contains(previous),
                "unexpected previous value: {previous}"
            );
        }
    }
}