housekeeping 0.0.3

//! Core logic for QSBR.
//!
//! This is an internal module, implementing the core algorithm used by
//! [`housekeeping`]. It has been made public because it is very difficult to
//! implement correctly, and it may have applications outside [`housekeeping`].
//! This implementation has been thoroughly reviewed and tested, and should be
//! built upon rather than implementing QSBR (or EBR) from scratch.
//!
//! [`housekeeping`]: crate
//!
//! # Description
//!
//! QSBR (Quiescent State Based Reclamation) is, most generally, a technique for
//! measuring and progressing _synchronization_ (as in concurrent programming)
//! between threads.
//!
//! QSBR measures progress in units of _phases_. At any time, every thread is
//! _using_ a particular phase. Threads may be all on the same phase, or may be
//! moving from one phase to the next (i.e. two phases are in use; the older is
//! losing users).
//!
//! Each phase is associated with a particular snapshot of global memory; all
//! threads within that phase will observe memory at least as recent as that
//! snapshot. This is based on the [`Acquire`] ordering for atomic operations.
//!
//! [`Acquire`]: core::sync::atomic::Ordering::Acquire
//!
//! Threads periodically enter a _quiescent state_ (QS), during which they try
//! to progress to the next phase. If the thread detects that it is not on the
//! latest phase, it will move there, and update its global view of memory. If
//! all threads are on the same phase, a QS will create a new phase.
//!
//! QSBR offers the following guarantee: *if a thread performs a write operation
//! within one phase, that operation is guaranteed to be visible to after two
//! phases*. That is: if a thread is on phase `p`, it has observed all writes
//! performed on phase `p - 2`. In addition, _all_ threads have observed all
//! writes performed on phase `p - 3`.
//
// TODO: Is there a formal proof of this somewhere?
//!
//! # Implementation
//!
//! This module provides [`Schedule`], encasing the core logic for QSBR. Every
//! thread participating in QSBR is prescribed a [`User`], through which it can
//! enter quiescent states and observe the current phase number.

use core::fmt;

use crate::loomish::{
    hint,
    sync::atomic::{self, AtomicUsize},
};

//----------- Schedule ---------------------------------------------------------

/// A QSBR schedule.
///
/// [`Schedule`] tracks three QSBR phases in a ring buffer. These phases are
/// identified by their indices (0, 1, or 2). Phases are ordered by ascending
/// indices, i.e. phase 0 comes before phase 1 (ignoring wrap-around).
///
/// See [the module-level documentation](self) for more information.
///
/// ## Invariants
///
/// Each phase is in one of the following states:
///
/// - `Primary`: The phase has zero or more users. It can gain or lose users.
///
/// - `Older`: The phase has at least one user. It is losing users (it cannot
///   gain any). The next phase is `Primary`.
///
/// - `Dead`: The phase has exactly zero users.
///
/// From oldest to newest, the three phases can have one of the following
/// combinations of states:
///
/// - `Dead-Dead-Primary`: There is only one phase in use.
///
///   If a thread undergoes a QS, it will create a new phase and switch to it.
///
///   This state is also used when there are no users at all.
///
/// - `Dead-Older-Primary`: There are two phases in use, and threads are
///   switching from the older phase to the primary one.
///
///   If a thread on the older phase undergoes a QS, it will switch to the
///   primary one. If a thread on the primary phase undergoes a QS, it's a no-op.
///
/// - `Dead-Primary-Primary`: There are two phases in use, and the older one is
///   about to be marked `Older`.
///
///   This is a temporary state that occurs when a new phase is being created.
///   If a thread on the older phase undergoes a QS, it will switch to the newer
///   one. If a thread on the newer phase undergoes a QS, it's a no-op.
///
///   During this time, newly registered users may be added to either primary
///   phase. This does not affect correctness.
#[derive(Debug)]
pub struct Schedule {
    /// The states of the phases.
    ///
    /// These atomic variables are synchronized with the global view of memory
    /// the corresponding phases represent. Performing an [`Acquire`] load will
    /// appropriately synchronize the current thread.
    //
    // TODO: Cache-line padding?
    states: [AtomicState; 3],
}

impl Schedule {
    /// Construct a new [`Schedule`].
    ///
    /// Phase 0 is selected as the primary phase; phases 1 and 2 are dead.
    pub fn new() -> Self {
        Self {
            states: [
                AtomicState::new(State::EMPTY_PRIMARY),
                AtomicState::new(State::DEAD),
                AtomicState::new(State::DEAD),
            ],
        }
    }

    /// Register a new [`User`].
    ///
    /// The primary phase will be found and a user will be registered against
    /// it; the index of that phase is returned.
    ///
    /// ## Panics
    ///
    /// May panic (directly or in other users) if `usize::MAX / 2` (or more)
    /// users are registered against this [`Schedule`].
    pub fn register(&self) -> User {
        // NOTE: There is always at least one primary phase.
        loop {
            for (index, phase) in self.states.iter().enumerate() {
                let state = phase.load(atomic::Ordering::Relaxed);

                if !state.is_primary() {
                    continue;
                }

                // NOTE(memory ordering): 'Acquire' on success in order to
                // synchronize this thread with the view of global memory
                // represented by the now-registered phase.
                //
                // TODO: I believe this has to be non-weak because IIRC the ARM
                // and RISC-V architectures only guarantee forward progress if
                // weak CAS operations are used within a tight loop. But someone
                // should check that.
                match phase.compare_exchange(
                    state,
                    state.add_user(),
                    atomic::Ordering::Acquire,
                    atomic::Ordering::Relaxed,
                ) {
                    Ok(_) => return User { index: index as u8 },
                    Err(_) => continue,
                }
            }

            // A primary phase could not be found, presumably due to concurrent
            // modifications. Hint to the CPU (and the compiler) that this is a
            // spin loop; this should (hopefully) help flush caches.
            hint::spin_loop();
        }
    }
}

impl Default for Schedule {
    fn default() -> Self {
        Self::new()
    }
}

//----------- User -------------------------------------------------------------

/// A user registered in a QSBR [`Schedule`].
#[derive(Debug)]
pub struct User {
    /// The index of the phase this user is registered in.
    ///
    /// This value must be exactly 0, 1, or 2.
    index: u8,
}

impl User {
    /// The index of the phase the user is in.
    ///
    /// This is exactly 0, 1, or 2.
    #[must_use]
    pub const fn index(&self) -> usize {
        self.index as usize
    }

    /// Progress the QSBR schedule.
    ///
    /// If all threads have progressed to the current phase, a new phase is
    /// initialized and moved to. If another thread has already initialized the
    /// next phase, it is moved to. Otherwise, no progress is made.
    ///
    /// If a new phase is moved to, and the old phase is being left, an instance
    /// of [`Leaving`] is returned; this indicates that progress was made, and
    /// allows the user to control how the phase is left. See [`Leaving`] for
    /// more information.
    pub fn progress<'s>(&mut self, schedule: &'s Schedule) -> Option<Leaving<'s>> {
        // Check the state of the currently registered phase.
        //
        // NOTE(memory ordering): A write to 'curr' to make it non-primary uses
        // 'Release' ordering and happens-after a write to initialize the next
        // phase. This read uses 'Acquire' ordering to observe that phase.
        let curr = &schedule.states[self.index as usize];
        let state = curr.load(atomic::Ordering::Acquire);

        if state.is_primary() {
            // The phase remains primary; check whether a new phase is needed.

            // Check whether the previous phase is dead.
            //
            // NOTE(memory ordering): A write to 'prev' to deregister a user
            // happens-after all writes by that user in that phase. The write
            // that makes 'prev' dead happens-after all writes performed during
            // that phase. These writes will be synchronized into the new phase.
            let prev = &schedule.states[(self.index as usize + 2) % 3];
            let state = prev.load(atomic::Ordering::Acquire);

            // If the previous phase is primary, then the current one is also
            // primary, and the previous phase is being switched away from. But
            // the previous phase must then have a user registered against it
            // which will mark it as non-primary.
            debug_assert_ne!(state, State::EMPTY_PRIMARY);

            if !state.is_dead() {
                // The previous state is not dead; there is nothing to do.
                return None;
            }

            // Race to initialize the next phase.
            //
            // NOTE(memory ordering): Synchronize the new phase with all the
            // writes from the previous (not current) phase. If the new phase
            // has already been initialized, then no synchronization is needed.
            let next = &schedule.states[(self.index as usize + 1) % 3];
            let initialized = next
                .compare_exchange(
                    State::DEAD,
                    State::EMPTY_PRIMARY.add_user(),
                    atomic::Ordering::Release,
                    atomic::Ordering::Relaxed,
                )
                .is_ok();

            // Register against the next phase.
            //
            // NOTE(memory ordering): The current thread is already synchronized
            // with the phase, so no further synchronization is needed.
            if !initialized {
                next.update(
                    atomic::Ordering::Relaxed,
                    atomic::Ordering::Relaxed,
                    |state| state.add_user(),
                );
            }

            // Mark the current phase as non-primary.
            //
            // NOTE(memory ordering): When a reader of 'curr' observes it is
            // non-primary, it will synchronize with this write to observe the
            // initialized state of 'next'.
            if initialized {
                curr.update(
                    atomic::Ordering::Release,
                    atomic::Ordering::Relaxed,
                    |state| state.with_primary(false),
                );
            }

            let leaving_index = self.index;
            self.index = (self.index + 1) % 3;

            Some(Leaving {
                schedule,
                index: leaving_index,
            })
        } else {
            // The phase is not primary; switch away from it.

            // Register against the next phase.
            let next = &schedule.states[(self.index as usize + 1) % 3];
            next.update(
                atomic::Ordering::Relaxed,
                atomic::Ordering::Relaxed,
                |state| state.add_user(),
            );

            let leaving_index = self.index;
            self.index = (self.index + 1) % 3;

            Some(Leaving {
                schedule,
                index: leaving_index,
            })
        }
    }

    /// Deregister this user, leaving the current phase.
    ///
    /// The user is consumed and a [`Leaving`] is returned, through which the
    /// caller can leave the current phase (and control how leaving occurs).
    pub fn deregister<'s>(self, schedule: &'s Schedule) -> Leaving<'s> {
        // Check the state of the currently registered phase.
        //
        // NOTE(memory ordering): A write to 'curr' to make it non-primary uses
        // 'Release' ordering and happens-after a write to initialize the next
        // phase. This read uses 'Acquire' ordering to observe that phase.
        let curr = &schedule.states[self.index as usize];
        let state = curr.load(atomic::Ordering::Acquire);

        if !state.is_primary() {
            // The current phase has been marked non-primary, by another thread.
            // There is no need to make further progress.
            return Leaving {
                schedule,
                index: self.index,
            };
        }

        // The phase remains primary; check whether a new phase is needed.

        // Check whether the previous phase is dead.
        //
        // NOTE(memory ordering): A write to 'prev' to deregister a user
        // happens-after all writes by that user in that phase. The write that
        // makes 'prev' dead happens-after all writes performed during that
        // phase. These writes will be synchronized into the new phase.
        let prev = &schedule.states[(self.index as usize + 2) % 3];
        let state = prev.load(atomic::Ordering::Acquire);

        // If the previous phase is primary, then the current one is also
        // primary, and the previous phase is being switched away from. But the
        // previous phase must then have a user registered against it which will
        // mark it as non-primary.
        debug_assert_ne!(state, State::EMPTY_PRIMARY);

        if !state.is_dead() {
            // The previous state is not dead; there is nothing to do.
            return Leaving {
                schedule,
                index: self.index,
            };
        }

        // Race to initialize the next phase.
        //
        // NOTE(memory ordering): Synchronize the new phase with all the writes
        // from the previous (not current) phase. If the new phase has already
        // been initialized, then no synchronization is needed.
        let next = &schedule.states[(self.index as usize + 1) % 3];
        let initialized = next
            .compare_exchange(
                State::DEAD,
                State::EMPTY_PRIMARY,
                atomic::Ordering::Release,
                atomic::Ordering::Relaxed,
            )
            .is_ok();

        // Mark the current phase as non-primary.
        //
        // NOTE(memory ordering): When a reader of 'curr' observes it is
        // non-primary, it will synchronize with this write to observe the
        // initialized state of 'next'.
        if initialized {
            curr.update(
                atomic::Ordering::Release,
                atomic::Ordering::Relaxed,
                |state| state.with_primary(false),
            );
        }

        Leaving {
            schedule,
            index: self.index,
        }
    }
}

//----------- Leaving ----------------------------------------------------------

/// A [`User`] leaving a phase.
///
/// This can be returned by [`User::progress()`] or [`User::deregister()`]. In
/// the former case, the user is registered against two consecutive phases, and
/// is leaving the older one; in the latter case, the user is leaving the only
/// phase it is registered against.
///
/// [`Leaving`] allows the caller to choose _how_ to leave the phase. Phases
/// are reference-counted, so the caller can check whether it is the last one
/// to leave the phase, and gain unique access to it. [`Leaving::leave_last()`]
/// will try acquiring unique access to the phase, and leave it otherwise. If
/// such functionality is not required, [`Leaving::leave()`] can be called (it
/// is also called automatically on drop).
///
/// **NOTE:** Forgetting a [`Leaving`] will deadlock the [`Schedule`],
/// preventing any further progress.
pub struct Leaving<'s> {
    /// The underlying schedule.
    schedule: &'s Schedule,

    /// The index of the phase being left.
    ///
    /// This value must be exactly 0, 1, or 2.
    index: u8,
}

impl<'s> Leaving<'s> {
    /// The associated QSBR schedule.
    #[must_use]
    pub const fn schedule(&self) -> &'s Schedule {
        self.schedule
    }

    /// The index of the phase being left.
    ///
    /// This is exactly 0, 1, or 2.
    #[must_use]
    pub const fn index(&self) -> usize {
        self.index as usize
    }

    /// Leave the phase.
    ///
    /// This is equivalent to dropping `self`, but may be preferable as it is
    /// more explicit.
    pub fn leave(self) {
        // Execute the drop hook.
        let _ = self;
    }

    /// Try being the last user to leave the phase.
    ///
    /// The phase is left _unless_ `self` is its last user, in which case an
    /// instance of [`LastLeaving`] is returned.
    pub fn leave_last(self) -> Option<LastLeaving<'s>> {
        let (schedule, index) = (self.schedule, self.index);
        // Don't run the drop hook.
        core::mem::forget(self);

        // Leave the phase unless it would become dead.
        //
        // NOTE(memory ordering): Synchronize 'state' with all the writes
        // performed by this thread in that phase. When a thread creates a
        // new phase, it will synchronize with 'state' and so carry forward the
        // updated view of memory.
        let state = &schedule.states[index as usize];
        match state.try_update(
            atomic::Ordering::Release,
            atomic::Ordering::Acquire,
            |state| Some(state.del_user()).filter(|s| !s.is_dead()),
        ) {
            // Left the phase successfully.
            Ok(_) => None,
            // Did not leave the phase because it would become dead.
            Err(_) => Some(LastLeaving { schedule, index }),
        }
    }
}

impl Drop for Leaving<'_> {
    fn drop(&mut self) {
        // Leave the phase without fuss.
        //
        // NOTE(memory ordering): Synchronize 'state' with all the writes
        // performed by this thread in that phase. When a thread creates a new
        // phase, it will synchronize with 'state' and so carry forward the
        // updated view of memory.
        let state = &self.schedule.states[self.index as usize];
        state.update(
            atomic::Ordering::Release,
            atomic::Ordering::Relaxed,
            |state| state.del_user(),
        );
    }
}

//----------- LastLeaving ------------------------------------------------------

/// The last user to leave a phase.
///
/// This can be returned by [`Leaving::leave_last()`]. It expresses unique
/// access to a phase; all other users are guaranteed to be on the next phase.
///
/// The phase will be released when the object is dropped. Until then, the
/// caller has unique access to any data associated with the phase.
///
/// **NOTE:** Forgetting a [`Leaving`] will deadlock the [`Schedule`],
/// preventing any further progress.
pub struct LastLeaving<'s> {
    /// The underlying schedule.
    schedule: &'s Schedule,

    /// The index of the phase being left.
    ///
    /// This value must be exactly 0, 1, or 2.
    index: u8,
}

impl<'s> LastLeaving<'s> {
    /// The associated QSBR schedule.
    #[must_use]
    pub const fn schedule(&self) -> &'s Schedule {
        self.schedule
    }

    /// The index of the phase being left.
    ///
    /// This is exactly 0, 1, or 2.
    #[must_use]
    pub const fn index(&self) -> usize {
        self.index as usize
    }
}

impl Drop for LastLeaving<'_> {
    fn drop(&mut self) {
        // Leave the phase.
        //
        // NOTE(memory ordering): Synchronize 'state' with all the writes
        // performed by this thread in that phase. When a thread creates a new
        // phase, it will synchronize with 'state' and so carry forward the
        // updated view of memory.
        let state = &self.schedule.states[self.index as usize];
        state.store(State::DEAD, atomic::Ordering::Release);
    }
}

//----------- State ------------------------------------------------------------

/// State for a QSBR phase.
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
#[repr(transparent)]
struct State(usize);

impl State {
    /// An older phase with no users.
    pub const DEAD: Self = Self(0);

    /// The primary phase, but with no users.
    pub const EMPTY_PRIMARY: Self = Self(1);

    /// Whether this is the primary phase.
    #[must_use]
    pub const fn is_primary(&self) -> bool {
        self.0 & 1 != 0
    }

    /// Whether the phase is dead.
    #[must_use]
    pub const fn is_dead(&self) -> bool {
        self.0 == 0
    }

    /// How many threads are using the phase.
    #[must_use]
    pub const fn users(&self) -> usize {
        self.0 / 2
    }

    /// Set whether this phase is primary.
    #[must_use]
    pub const fn with_primary(self, primary: bool) -> Self {
        Self(self.0 & !1 | primary as usize)
    }

    /// Add a user to the phase.
    #[must_use]
    pub const fn add_user(self) -> Self {
        Self(self.0.strict_add(2))
    }

    /// Remove a user from the epoch.
    #[must_use]
    pub const fn del_user(self) -> Self {
        Self(self.0.strict_sub(2))
    }
}

impl fmt::Debug for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("State")
            .field("primary", &self.is_primary())
            .field("users", &self.users())
            .finish()
    }
}

//----------- AtomicState ------------------------------------------------------

/// Atomic state for a QSBR phase.
///
/// This is a wrapper around [`AtomicUsize`] that translates its operations to
/// [`State`]s.
#[repr(transparent)]
struct AtomicState(AtomicUsize);

impl AtomicState {
    /// Construct a new [`AtomicState`].
    pub fn new(state: State) -> Self {
        Self(AtomicUsize::new(state.0))
    }

    /// Load the state.
    pub fn load(&self, order: atomic::Ordering) -> State {
        State(self.0.load(order))
    }

    /// Store the state.
    pub fn store(&self, val: State, order: atomic::Ordering) {
        self.0.store(val.0, order);
    }

    /// Swap the state from an expected value to a new one.
    pub fn compare_exchange(
        &self,
        current: State,
        new: State,
        success: atomic::Ordering,
        failure: atomic::Ordering,
    ) -> Result<State, State> {
        match self.0.compare_exchange(current.0, new.0, success, failure) {
            Ok(state) => Ok(State(state)),
            Err(state) => Err(State(state)),
        }
    }

    /// Swap the state from an expected value to a new one.
    pub fn compare_exchange_weak(
        &self,
        current: State,
        new: State,
        success: atomic::Ordering,
        failure: atomic::Ordering,
    ) -> Result<State, State> {
        match self
            .0
            .compare_exchange_weak(current.0, new.0, success, failure)
        {
            Ok(state) => Ok(State(state)),
            Err(state) => Err(State(state)),
        }
    }

    /// Update the state based on its current value, fallibly.
    pub fn try_update<F>(
        &self,
        set_order: atomic::Ordering,
        fetch_order: atomic::Ordering,
        mut f: F,
    ) -> Result<State, State>
    where
        F: FnMut(State) -> Option<State>,
    {
        let mut curr = self.load(fetch_order);
        loop {
            let Some(next) = (f)(curr) else {
                break Err(curr);
            };
            match self.compare_exchange_weak(curr, next, set_order, fetch_order) {
                Ok(curr) => break Ok(curr),
                Err(actual) => curr = actual,
            }
        }
    }

    /// Update the state based on its current value.
    pub fn update<F>(
        &self,
        set_order: atomic::Ordering,
        fetch_order: atomic::Ordering,
        mut f: F,
    ) -> State
    where
        F: FnMut(State) -> State,
    {
        let mut curr = self.load(fetch_order);
        loop {
            let next = (f)(curr);
            match self.compare_exchange_weak(curr, next, set_order, fetch_order) {
                Ok(curr) => break curr,
                Err(actual) => curr = actual,
            }
        }
    }
}

impl fmt::Debug for AtomicState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.load(atomic::Ordering::Relaxed).fmt(f)
    }
}