77 lines
2.0 KiB
Rust
77 lines
2.0 KiB
Rust
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
|
|
}
|
|
}
|
|
}
|