From 03d0cdb7145042ce51c9eb98bfd081a4f10fe8e1 Mon Sep 17 00:00:00 2001 From: Mica White Date: Fri, 8 Mar 2024 12:18:41 -0500 Subject: Keyable --- src/guard.rs | 16 ++++---- src/key.rs | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 4 +- src/lock.rs | 98 ---------------------------------------------- src/lockable.rs | 4 ++ src/mutex.rs | 30 +++++++------- 6 files changed, 149 insertions(+), 121 deletions(-) create mode 100644 src/key.rs delete mode 100644 src/lock.rs (limited to 'src') diff --git a/src/guard.rs b/src/guard.rs index 3c1b636..30e7d4a 100644 --- a/src/guard.rs +++ b/src/guard.rs @@ -1,17 +1,17 @@ use std::ops::{Deref, DerefMut}; -use crate::{lockable::Lockable, ThreadKey}; +use crate::{key::Keyable, lockable::Lockable}; /// A guard for a generic [`Lockable`] type. -pub struct LockGuard<'a, L: Lockable<'a>> { +pub struct LockGuard<'a, L: Lockable<'a>, Key: Keyable> { guard: L::Output, - _key: &'a mut ThreadKey, + _key: Key, } -impl<'a, L: Lockable<'a>> LockGuard<'a, L> { +impl<'a, L: Lockable<'a>, Key: Keyable> LockGuard<'a, L, Key> { /// Locks the lockable type and returns a guard that can be used to access /// the underlying data. - pub fn lock(lock: &'a L, key: &'a mut ThreadKey) -> Self { + pub fn lock(lock: &'a L, key: Key) -> Self { Self { // safety: we have the thread's key guard: unsafe { lock.lock() }, @@ -22,7 +22,7 @@ impl<'a, L: Lockable<'a>> LockGuard<'a, L> { /// Attempts to lock the guard without blocking. If successful, this method /// returns a guard that can be used to access the data. Otherwise, the key /// is given back as an error. - pub fn try_lock(lock: &'a L, key: &'a mut ThreadKey) -> Option { + pub fn try_lock(lock: &'a L, key: Key) -> Option { // safety: we have the thread's key unsafe { lock.try_lock() }.map(|guard| Self { guard, _key: key }) } @@ -35,7 +35,7 @@ impl<'a, L: Lockable<'a>> LockGuard<'a, L> { } } -impl<'a, L: Lockable<'a>> Deref for LockGuard<'a, L> { +impl<'a, L: Lockable<'a>, Key: Keyable> Deref for LockGuard<'a, L, Key> { type Target = L::Output; fn deref(&self) -> &Self::Target { @@ -43,7 +43,7 @@ impl<'a, L: Lockable<'a>> Deref for LockGuard<'a, L> { } } -impl<'a, L: Lockable<'a>> DerefMut for LockGuard<'a, L> { +impl<'a, L: Lockable<'a>, Key: Keyable> DerefMut for LockGuard<'a, L, Key> { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.guard } diff --git a/src/key.rs b/src/key.rs new file mode 100644 index 0000000..92f3b99 --- /dev/null +++ b/src/key.rs @@ -0,0 +1,118 @@ +use std::fmt::{self, Debug}; +use std::marker::PhantomData; +use std::sync::atomic::{AtomicBool, Ordering}; + +use once_cell::sync::Lazy; +use thread_local::ThreadLocal; + +use self::sealed::Sealed; + +mod sealed { + use super::ThreadKey; + pub trait Sealed {} + impl Sealed for ThreadKey {} + impl Sealed for &mut ThreadKey {} +} + +static KEY: Lazy> = Lazy::new(ThreadLocal::new); + +/// The key for the current thread. +/// +/// Only one of these exist per thread. To get the current thread's key, call +/// [`ThreadKey::lock`]. If the `ThreadKey` is dropped, it can be reobtained. +pub type ThreadKey = Key<'static>; + +/// A dumb lock that's just a wrapper for an [`AtomicBool`]. +#[derive(Debug, Default)] +pub struct AtomicLock { + is_locked: AtomicBool, +} + +pub struct Key<'a> { + phantom: PhantomData<*const ()>, // implement !Send and !Sync + lock: &'a AtomicLock, +} + +/// Allows the type to be used as a key for a lock +/// +/// # Safety +/// +/// Only one value which implements this trait may be allowed to exist at a +/// time. Creating a new `Keyable` value requires making any other `Keyable` +/// values invalid. +pub unsafe trait Keyable: Sealed {} +unsafe impl Keyable for ThreadKey {} +unsafe impl Keyable for &mut ThreadKey {} + +impl Debug for ThreadKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ThreadKey") + } +} + +impl<'a> Drop for Key<'a> { + fn drop(&mut self) { + unsafe { self.lock.force_unlock() } + } +} + +impl ThreadKey { + /// Get the current thread's `ThreadKey`, if it's not already taken. + /// + /// The first time this is called, it will successfully return a + /// `ThreadKey`. However, future calls to this function will return + /// [`None`], unless the key is dropped or unlocked first. + #[must_use] + pub fn lock() -> Option { + KEY.get_or_default().try_lock() + } + + /// Unlocks the `ThreadKey`. + /// + /// After this method is called, a call to [`ThreadKey::lock`] will return + /// this `ThreadKey`. + pub fn unlock(key: Self) { + drop(key); + } +} + +impl AtomicLock { + /// Create a new unlocked `AtomicLock`. + #[must_use] + pub const fn new() -> Self { + Self { + is_locked: AtomicBool::new(false), + } + } + + /// Checks whether this `Lock` is currently locked. + pub fn is_locked(&self) -> bool { + self.is_locked.load(Ordering::Relaxed) + } + + /// Attempt to lock the `AtomicLock`. + /// + /// If the lock is already locked, then this'll return false. If it is + /// unlocked, it'll return true. If the lock is currently unlocked, then it + /// will be locked after this function is called. + /// + /// This is not a fair lock. It is not recommended to call this function + /// repeatedly in a loop. + pub fn try_lock(&self) -> Option { + // safety: we just acquired the lock + (!self.is_locked.swap(true, Ordering::Acquire)).then_some(Key { + phantom: PhantomData, + lock: self, + }) + } + + /// Forcibly unlocks the `AtomicLock`. + /// + /// # Safety + /// + /// This should only be called if the key to the lock has been "lost". That + /// means the program no longer has a reference to the key. + pub unsafe fn force_unlock(&self) { + self.is_locked.store(false, Ordering::Release); + } +} diff --git a/src/lib.rs b/src/lib.rs index 457fae0..35f7bb4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,11 +4,11 @@ #![allow(clippy::declare_interior_mutable_const)] mod guard; -mod lock; +mod key; mod lockable; pub mod mutex; pub use guard::LockGuard; -pub use lock::{Key, ThreadKey}; +pub use key::{Key, ThreadKey}; pub use lockable::Lockable; pub use mutex::Mutex; diff --git a/src/lock.rs b/src/lock.rs deleted file mode 100644 index 942474a..0000000 --- a/src/lock.rs +++ /dev/null @@ -1,98 +0,0 @@ -use std::fmt::{self, Debug}; -use std::marker::PhantomData; -use std::sync::atomic::{AtomicBool, Ordering}; - -use once_cell::sync::Lazy; -use thread_local::ThreadLocal; - -static KEY: Lazy> = Lazy::new(ThreadLocal::new); - -/// The key for the current thread. -/// -/// Only one of these exist per thread. To get the current thread's key, call -/// [`ThreadKey::lock`]. If the `ThreadKey` is dropped, it can be reobtained. -pub type ThreadKey = Key<'static>; - -/// A dumb lock that's just a wrapper for an [`AtomicBool`]. -#[derive(Debug, Default)] -pub struct Lock { - is_locked: AtomicBool, -} - -pub struct Key<'a> { - phantom: PhantomData<*const ()>, // implement !Send and !Sync - lock: &'a Lock, -} - -impl Debug for ThreadKey { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "ThreadKey") - } -} - -impl<'a> Drop for Key<'a> { - fn drop(&mut self) { - unsafe { self.lock.force_unlock() } - } -} - -impl ThreadKey { - /// Get the current thread's `ThreadKey`, if it's not already taken. - /// - /// The first time this is called, it will successfully return a - /// `ThreadKey`. However, future calls to this function will return - /// [`None`], unless the key is dropped or unlocked first. - #[must_use] - pub fn lock() -> Option { - KEY.get_or_default().try_lock() - } - - /// Unlocks the `ThreadKey`. - /// - /// After this method is called, a call to [`ThreadKey::lock`] will return - /// this `ThreadKey`. - pub fn unlock(key: Self) { - drop(key); - } -} - -impl Lock { - /// Create a new unlocked `Lock`. - #[must_use] - pub const fn new() -> Self { - Self { - is_locked: AtomicBool::new(false), - } - } - - /// Checks whether this `Lock` is currently locked. - pub fn is_locked(&self) -> bool { - self.is_locked.load(Ordering::Relaxed) - } - - /// Attempt to lock the `Lock`. - /// - /// If the lock is already locked, then this'll return false. If it is - /// unlocked, it'll return true. If the lock is currently unlocked, then it - /// will be locked after this function is called. - /// - /// This is not a fair lock. It is not recommended to call this function - /// repeatedly in a loop. - pub fn try_lock(&self) -> Option { - // safety: we just acquired the lock - (!self.is_locked.swap(true, Ordering::Acquire)).then_some(Key { - phantom: PhantomData, - lock: self, - }) - } - - /// Forcibly unlocks the `Lock`. - /// - /// # Safety - /// - /// This should only be called if the key to the lock has been "lost". That - /// means the program no longer has a reference to the key. - pub unsafe fn force_unlock(&self) { - self.is_locked.store(false, Ordering::Release); - } -} diff --git a/src/lockable.rs b/src/lockable.rs index b6ba7d6..cda4e13 100644 --- a/src/lockable.rs +++ b/src/lockable.rs @@ -31,6 +31,8 @@ pub unsafe trait Lockable<'a>: sealed::Sealed { /// * Use this without ownership or mutable access to the [`ThreadKey`], /// which should last as long as the return value is alive. /// * Call this on multiple locks without unlocking first. + /// + /// [`ThreadKey`]: `crate::key::ThreadKey` unsafe fn lock(&'a self) -> Self::Output; /// Attempt to lock without blocking. @@ -42,6 +44,8 @@ pub unsafe trait Lockable<'a>: sealed::Sealed { /// It is undefined behavior to use this without ownership or mutable /// access to the [`ThreadKey`], which should last as long as the return /// value is alive. + /// + /// [`ThreadKey`]: `crate::key::ThreadKey` unsafe fn try_lock(&'a self) -> Option; /// Release the lock diff --git a/src/mutex.rs b/src/mutex.rs index e806fac..28e1786 100644 --- a/src/mutex.rs +++ b/src/mutex.rs @@ -1,8 +1,7 @@ use std::cell::UnsafeCell; use std::ops::{Deref, DerefMut}; -use crate::lock::Lock; -use crate::ThreadKey; +use crate::key::{AtomicLock, Keyable}; /// A spinning mutex pub type SpinLock = Mutex; @@ -37,11 +36,13 @@ pub unsafe trait RawMutex { /// A raw mutex which just spins pub struct RawSpin { - lock: Lock, + lock: AtomicLock, } unsafe impl RawMutex for RawSpin { - const INIT: Self = Self { lock: Lock::new() }; + const INIT: Self = Self { + lock: AtomicLock::new(), + }; fn lock(&self) { loop { @@ -124,12 +125,12 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexRef<'a, T, R> { /// /// [`lock`]: `Mutex::lock` /// [`try_lock`]: `Mutex::try_lock` -pub struct MutexGuard<'a, T: ?Sized + 'a, R: RawMutex = RawSpin> { +pub struct MutexGuard<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex = RawSpin> { mutex: MutexRef<'a, T, R>, - _thread_key: &'a mut ThreadKey, + _thread_key: Key, } -impl<'a, T: ?Sized + 'a, R: RawMutex> Deref for MutexGuard<'a, T, R> { +impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> Deref for MutexGuard<'a, T, Key, R> { type Target = T; fn deref(&self) -> &Self::Target { @@ -137,16 +138,16 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> Deref for MutexGuard<'a, T, R> { } } -impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexGuard<'a, T, R> { +impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> DerefMut for MutexGuard<'a, T, Key, R> { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.mutex } } -impl<'a, T: ?Sized + 'a, R: RawMutex> MutexGuard<'a, T, R> { +impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> MutexGuard<'a, T, Key, R> { /// Create a guard to the given mutex. Undefined if multiple guards to the /// same mutex exist at once. - unsafe fn new(mutex: &'a Mutex, thread_key: &'a mut ThreadKey) -> Self { + const unsafe fn new(mutex: &'a Mutex, thread_key: Key) -> Self { Self { mutex: MutexRef(mutex), _thread_key: thread_key, @@ -196,7 +197,7 @@ impl Mutex { /// let key = ThreadKey::lock().unwrap(); /// assert_eq!(*mutex.lock(key), 10); /// ``` - pub fn lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> MutexGuard<'_, T, R> { + pub fn lock<'s, 'a: 's, Key: Keyable>(&'s self, key: Key) -> MutexGuard<'_, T, Key, R> { self.raw.lock(); // safety: we just locked the mutex @@ -240,7 +241,10 @@ impl Mutex { /// let key = ThreadKey::lock().unwrap(); /// assert_eq!(*mutex.lock(key), 10); /// ``` - pub fn try_lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> Option> { + pub fn try_lock<'s, 'a: 's, Key: Keyable>( + &'s self, + key: Key, + ) -> Option> { if self.raw.try_lock() { // safety: we just locked the mutex Some(unsafe { MutexGuard::new(self, key) }) @@ -281,7 +285,7 @@ impl Mutex { /// let key = Mutex::unlock(guard); /// ``` #[allow(clippy::missing_const_for_fn)] - pub fn unlock(guard: MutexGuard<'_, T, R>) { + pub fn unlock(guard: MutexGuard<'_, T, impl Keyable, R>) { drop(guard); } } -- cgit v1.2.3