slight refactor of Lock

This commit is contained in:
Janis 2025-07-06 14:26:39 +02:00
parent 4c70dbfc71
commit 5a3b6447f3

View file

@ -3,9 +3,6 @@ use core::{
sync::atomic::{AtomicU32, Ordering}, sync::atomic::{AtomicU32, Ordering},
}; };
const LOCKED_BIT: u32 = 0b001;
const EMPTY: u32 = 0;
/// A simple lock implementation using an atomic u32. /// A simple lock implementation using an atomic u32.
#[repr(transparent)] #[repr(transparent)]
pub struct Lock { pub struct Lock {
@ -13,6 +10,9 @@ pub struct Lock {
} }
impl Lock { impl Lock {
const LOCKED: u32 = 0b001;
const EMPTY: u32 = 0;
/// Creates a new lock in the unlocked state. /// Creates a new lock in the unlocked state.
pub const fn new() -> Self { pub const fn new() -> Self {
Self { Self {
@ -35,7 +35,12 @@ impl Lock {
// attempt acquiring the lock with no contention. // attempt acquiring the lock with no contention.
if self if self
.inner .inner
.compare_exchange_weak(EMPTY, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed) .compare_exchange_weak(
Self::EMPTY,
Self::LOCKED,
Ordering::Acquire,
Ordering::Relaxed,
)
.is_ok() .is_ok()
{ {
// We successfully acquired the lock. // We successfully acquired the lock.
@ -48,7 +53,7 @@ impl Lock {
pub fn unlock(&self) { pub fn unlock(&self) {
// use release semantics to ensure that all previous writes are // use release semantics to ensure that all previous writes are
// available to other threads. // available to other threads.
self.inner.fetch_and(!LOCKED_BIT, Ordering::Release); self.inner.fetch_and(!Self::LOCKED, Ordering::Release);
} }
fn lock_slow(&self) { fn lock_slow(&self) {
@ -58,11 +63,11 @@ impl Lock {
let mut state = self.inner.load(Ordering::Acquire); let mut state = self.inner.load(Ordering::Acquire);
loop { loop {
// If the lock isn't locked, we can try to acquire it. // If the lock isn't locked, we can try to acquire it.
if state & LOCKED_BIT == 0 { if state & Self::LOCKED == 0 {
// Try to acquire the lock. // Try to acquire the lock.
match self.inner.compare_exchange_weak( match self.inner.compare_exchange_weak(
state, state,
state | LOCKED_BIT, state | Self::LOCKED,
Ordering::Acquire, Ordering::Acquire,
Ordering::Relaxed, Ordering::Relaxed,
) { ) {
@ -97,13 +102,13 @@ impl Lock {
} }
// If we reach here, we need to park the thread. // If we reach here, we need to park the thread.
atomic_wait::wait(&self.inner, LOCKED_BIT); atomic_wait::wait(&self.inner, Self::LOCKED);
if self if self
.inner .inner
.compare_exchange_weak( .compare_exchange_weak(
state, state,
state | LOCKED_BIT, state | Self::LOCKED,
Ordering::Acquire, Ordering::Acquire,
Ordering::Relaxed, Ordering::Relaxed,
) )