You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

102 lines
3.0 KiB
Rust

use core::{cell::UnsafeCell, hint, ops::{Deref, DerefMut}, sync::atomic::{AtomicBool, Ordering}};
pub struct Spinlock<T> {
data: UnsafeCell<T>,
locked: AtomicBool,
}
impl<T> Spinlock<T> {
pub const fn new(data: T) -> Self {
Self {
data: UnsafeCell::new(data),
locked: AtomicBool::new(false),
}
}
pub fn lock(&self) -> SpinlockGuard<T> {
// If we observe `locked` was `false`, then:
// - Use acquire ordering, so nothing inside the critical section gets reordered before we
// observed the `false` value.
// - Store `true`, so nothing else can enter the critical section until we exit it.
// Otherwise, spin until we observe a `false` value.
while self.locked
.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
hint::spin_loop();
}
SpinlockGuard { lock: self }
}
/// # Safety
/// There must be no "active" `SpinlockGuards` for this lock, i.e. a `SpinlockGuard` which can be
/// used to obtain a reference to the spinlock-protected data.
unsafe fn unlock(&self) {
// Unset `locked` with release ordering so that nothing inside the critical section gets
// reordered to after we stored `false`.
self.locked.store(false, Ordering::Release);
}
/// # Safety
/// No mutable references to the spinlock-protected data may exist.
unsafe fn get<'s, 'a>(&'s self) -> &'a T
where
's: 'a,
{
unsafe { &*self.data.get() }
}
/// # Safety
/// No references to the spinlock-protected data may exist.
unsafe fn get_mut<'s, 'a>(&'s self) -> &'a mut T
where
's: 'a,
{
unsafe { &mut *self.data.get() }
}
}
unsafe impl<T> Sync for Spinlock<T>
where
T: Send,
{}
pub struct SpinlockGuard<'a, T> {
lock: &'a Spinlock<T>,
}
impl<'a, T> Deref for SpinlockGuard<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &Self::Target {
// SAFETY:
// For the entire lifetime of the `SpinlockGuard`, `locked` remains `true`, so we have
// exclusive access to `data`, so no mutable references to `data` can exist.
unsafe { self.lock.get() }
}
}
impl<'a, T> DerefMut for SpinlockGuard<'a, T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
// SAFETY:
// For the entire lifetime of the `SpinlockGuard`, `locked` remains `true`, so we have
// exclusive access to `data`, so no other references to `data` can exist.
unsafe { self.lock.get_mut() }
}
}
impl<'a, T> Drop for SpinlockGuard<'a, T> {
fn drop(&mut self) {
// SAFETY:
// Only one `SpinlockGuard` can exist at a time for a particular lock, since we set `locked`
// to true before creating a guard and refuse to create any new ones until it is `false` again.
// Therefore, we are the only `SpinlockGuard` for the lock. Since this is the destructor, and
// we don't access the spinlock-protected data here, there are therefore no "active"
// `SpinlockGuard`s remaining for the lock.
unsafe { self.lock.unlock() }
}
}