RedBear-OS/src/sync/barrier.rs

use core::{
    num::NonZeroU32,
    sync::atomic::{AtomicU32, Ordering},
};

pub struct Barrier {
    original_count: NonZeroU32,
    // 4
    lock: crate::sync::Mutex<Inner>,
    // 16
    cvar: FutexState,
    // 24
}
#[derive(Debug)]
struct Inner {
    _unused0: u32,
    _unused1: u32,
}

struct FutexState {
    count: AtomicU32,
    sense: AtomicU32,
}

impl FutexState {
    const fn new(count: u32) -> Self {
        Self {
            count: AtomicU32::new(count),
            sense: AtomicU32::new(0),
        }
    }
}

pub enum WaitResult {
    Waited,
    NotifiedAll,
}

impl Barrier {
    pub fn new(count: NonZeroU32) -> Self {
        Self {
            original_count: count,
            lock: crate::sync::Mutex::new(Inner {
                _unused0: 0,
                _unused1: 0,
            }),
            cvar: FutexState::new(count.get()),
        }
    }
    pub fn wait(&self) -> WaitResult {
        let _ = &self.lock;
        let sense = self.cvar.sense.load(Ordering::Acquire);

        if self.cvar.count.fetch_sub(1, Ordering::AcqRel) == 1 {
            self.cvar
                .count
                .store(self.original_count.get(), Ordering::Relaxed);
            self.cvar
                .sense
                .store(sense.wrapping_add(1), Ordering::Release);
            crate::sync::futex_wake(&self.cvar.sense, i32::MAX);

            WaitResult::NotifiedAll
        } else {
            // SMP fix: wait directly on the barrier generation word instead of routing through the
            // condvar unlock->futex_wait path. If the last thread flips `sense` after we load it
            // but before our futex wait starts, the futex observes a stale value and returns
            // immediately instead of sleeping forever after a missed broadcast wakeup.
            while self.cvar.sense.load(Ordering::Acquire) == sense {
                let _ = crate::sync::futex_wait(&self.cvar.sense, sense, None);
            }

            WaitResult::Waited
        }
    }
}