diff --git a/hercules_rt/src/lib.rs b/hercules_rt/src/lib.rs index a5954ca03c32441709a2ce5452e9b2584f9c6155..9265808be132b0137293332fea7ba4a378bd1bd9 100644 --- a/hercules_rt/src/lib.rs +++ b/hercules_rt/src/lib.rs @@ -5,6 +5,7 @@ use std::future::Future; use std::marker::PhantomData; use std::ptr::{copy_nonoverlapping, write_bytes, NonNull}; use std::slice::{from_raw_parts, from_raw_parts_mut}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::OnceLock; /* @@ -928,3 +929,30 @@ unsafe impl GlobalAlloc for AlignedAlloc { #[global_allocator] static A: AlignedAlloc = AlignedAlloc; + +pub struct SpinBarrier { + num: usize, + waiting: AtomicUsize, + gen: AtomicUsize, +} + +impl SpinBarrier { + pub const fn new(num: usize) -> Self { + SpinBarrier { + num, + waiting: AtomicUsize::new(0), + gen: AtomicUsize::new(0), + } + } + + pub fn wait(&self) { + let old_gen = self.gen.load(Ordering::Acquire); + let old_waiting = self.waiting.fetch_add(1, Ordering::Relaxed); + if old_waiting + 1 == self.num { + self.waiting.store(0, Ordering::Relaxed); + self.gen.fetch_add(1, Ordering::Release); + } else { + while old_gen == self.gen.load(Ordering::Acquire) {} + } + } +}