use std::{
fmt,
mem::{align_of, offset_of, size_of},
sync::atomic::{AtomicBool, AtomicI64, AtomicU32, AtomicU64, Ordering},
};
use log::trace;
pub const MAX_USER_WORDS: usize = 5;
#[repr(C, align(64))]
pub struct EpochEntry {
epoch: AtomicU64,
thread_id: AtomicU64,
reentrant: AtomicU32,
reserved: AtomicBool,
_pad: [u8; 3],
user_words: [AtomicI64; MAX_USER_WORDS],
}
const _: () = {
assert!(size_of::<EpochEntry>() == 64);
assert!(align_of::<EpochEntry>() == 64);
assert!(offset_of!(EpochEntry, epoch) == 0);
assert!(offset_of!(EpochEntry, thread_id) == 8);
assert!(offset_of!(EpochEntry, reentrant) == 16);
assert!(offset_of!(EpochEntry, reserved) == 20);
assert!(offset_of!(EpochEntry, user_words) == 24);
};
impl EpochEntry {
pub const fn new() -> Self {
Self {
epoch: AtomicU64::new(0),
thread_id: AtomicU64::new(0),
reentrant: AtomicU32::new(0),
reserved: AtomicBool::new(false),
_pad: [0; 3],
user_words: [const { AtomicI64::new(0) }; MAX_USER_WORDS],
}
}
#[inline]
pub fn try_reserve(&self) -> bool {
if self.epoch.load(Ordering::SeqCst) != 0 {
return false;
}
if self
.reserved
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
return false;
}
if self.epoch.load(Ordering::SeqCst) != 0 {
self.reserved.store(false, Ordering::SeqCst);
return false;
}
true
}
#[inline]
pub fn release_reserve(&self) {
trace!("释放 EpochEntry 条目槽位");
self.reset();
self.reserved.store(false, Ordering::SeqCst);
}
#[inline]
pub fn reset(&self) {
self.reentrant.store(0, Ordering::Relaxed);
self.thread_id.store(0, Ordering::Relaxed);
self.epoch.store(0, Ordering::Release);
}
#[inline]
pub fn enter_with_tid(&self, current_epoch: &AtomicU64, thread_id: u64) -> u64 {
let prev = self.reentrant.load(Ordering::Relaxed);
if prev > 0 {
self.reentrant.store(prev + 1, Ordering::Relaxed);
return self.epoch.load(Ordering::Acquire);
}
let epoch = current_epoch.load(Ordering::Acquire);
self.thread_id.store(thread_id, Ordering::Relaxed);
self.epoch.store(epoch, Ordering::Release);
self.reentrant.store(1, Ordering::Relaxed);
epoch
}
#[inline]
pub fn exit(&self) -> bool {
let prev = self.reentrant.load(Ordering::Relaxed);
if prev == 0 {
return false;
}
if prev > 1 {
self.reentrant.store(prev - 1, Ordering::Relaxed);
return false;
}
self.reset();
true
}
#[inline]
pub fn try_claim(&self, thread_id: u64, current_epoch: &AtomicU64) -> bool {
if self.epoch.load(Ordering::SeqCst) != 0 || self.reserved.load(Ordering::SeqCst) {
return false;
}
let epoch = current_epoch.load(Ordering::Acquire);
if self
.epoch
.compare_exchange(0, epoch, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
return false;
}
if self.reserved.load(Ordering::SeqCst) {
self.epoch.store(0, Ordering::SeqCst);
return false;
}
self.thread_id.store(thread_id, Ordering::Release);
self.reentrant.store(1, Ordering::Release);
true
}
#[inline]
pub fn refresh_epoch(&self, new_epoch: u64) {
self.epoch.store(new_epoch, Ordering::Release);
}
#[inline]
pub fn protected_epoch(&self) -> u64 {
self.epoch.load(Ordering::Acquire)
}
#[inline]
pub fn is_protected(&self) -> bool {
self.protected_epoch() != 0
}
#[inline]
pub fn reentrant_count(&self) -> u32 {
self.reentrant.load(Ordering::Relaxed)
}
#[inline]
pub fn inc_reentrant(&self) {
let prev = self.reentrant.load(Ordering::Relaxed);
self.reentrant.store(prev + 1, Ordering::Relaxed);
}
#[inline]
pub fn thread_id(&self) -> u64 {
self.thread_id.load(Ordering::Acquire)
}
#[inline]
pub unsafe fn user_word_unchecked(&self, idx: usize) -> i64 {
debug_assert!(idx < MAX_USER_WORDS);
unsafe { self.user_words.get_unchecked(idx) }.load(Ordering::Acquire)
}
#[inline]
pub unsafe fn set_user_word_unchecked(&self, idx: usize, val: i64) {
debug_assert!(idx < MAX_USER_WORDS);
unsafe { self.user_words.get_unchecked(idx) }.store(val, Ordering::Release);
}
#[inline]
pub unsafe fn user_word_atomic_unchecked(&self, idx: usize) -> &AtomicI64 {
debug_assert!(idx < MAX_USER_WORDS);
unsafe { self.user_words.get_unchecked(idx) }
}
}
impl Default for EpochEntry {
fn default() -> Self {
Self::new()
}
}
impl fmt::Debug for EpochEntry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EpochEntry")
.field("epoch", &self.epoch.load(Ordering::Relaxed))
.field("thread_id", &self.thread_id.load(Ordering::Relaxed))
.field("reentrant", &self.reentrant.load(Ordering::Relaxed))
.field("reserved", &self.reserved.load(Ordering::Relaxed))
.finish()
}
}