Skip to content

fix Mutex and Rwlock primitives #329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
arihant2math wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix Mutex and Rwlock primitives #329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 13 additions & 42 deletions libkernel/src/sync/mutex.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@ pub struct Mutex<T: ?Sized, CPU: CpuOps> {
///
/// When an `AsyncMutexGuard` is dropped, it automatically releases the lock and
/// wakes up the next task.
///
/// The guard is `Send` when `T: Send`, and `Sync` when `T: Sync`, independently
/// of the CPU marker type. Unlike the mutex itself, sharing a guard exposes
/// shared references to the protected value, so a guard over `Cell` cannot be
/// shared between threads.
#[must_use = "if unused, the Mutex will immediately unlock"]
pub struct AsyncMutexGuard<'a, T: ?Sized, CPU: CpuOps> {
mutex: &'a Mutex<T, CPU>,
Expand Down Expand Up @@ -134,46 +139,12 @@ impl<T: ?Sized, CPU: CpuOps> DerefMut for AsyncMutexGuard<'_, T, CPU> {
unsafe impl<T: ?Sized + Send, CPU: CpuOps> Send for Mutex<T, CPU> {}
unsafe impl<T: ?Sized + Send, CPU: CpuOps> Sync for Mutex<T, CPU> {}

impl<CPU: CpuOps> Mutex<(), CPU> {
/// Acquires the mutex lock without caring about the data.
pub(crate) fn acquire(&self) -> MutexAcquireFuture<'_, CPU> {
MutexAcquireFuture { mutex: self }
}

/// Releases the mutex lock without caring about the data.
///
/// # Safety
/// The caller must ensure that they have previously called [`Self::acquire()`].
pub(crate) unsafe fn release(&self) {
let mut state = self.state.lock_save_irq();

if let Some(next_waker) = state.waiters.pop_front() {
next_waker.wake();
}

state.is_locked = false;
}
}

/// A future that resolves to a locked mutex
pub struct MutexAcquireFuture<'a, CPU: CpuOps> {
mutex: &'a Mutex<(), CPU>,
}

impl<CPU: CpuOps> Future for MutexAcquireFuture<'_, CPU> {
type Output = ();

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut state = self.mutex.state.lock_save_irq();
// SAFETY: Moving the guard transfers exclusive access to T. The guard holds no
// saved interrupt state; its drop acquires and releases the state spinlock on
// the current CPU, so it is not tied to the CPU that acquired the mutex.
unsafe impl<T: ?Sized + Send, CPU: CpuOps> Send for AsyncMutexGuard<'_, T, CPU> {}

if !state.is_locked {
state.is_locked = true;
Poll::Ready(())
} else {
if state.waiters.iter().all(|w| !w.will_wake(cx.waker())) {
state.waiters.push_back(cx.waker().clone());
}
Poll::Pending
}
}
}
// SAFETY: Shared access to a guard exposes only &T, which may cross threads when
// T is Sync. Mutable access and dropping the guard still require ownership.
// This explicit impl prevents inheriting the mutex's weaker T: Send Sync bound.
unsafe impl<T: ?Sized + Sync, CPU: CpuOps> Sync for AsyncMutexGuard<'_, T, CPU> {}
219 changes: 185 additions & 34 deletions libkernel/src/sync/rwlock.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,27 +1,165 @@
//! Async-aware readers–writer lock.

use super::spinlock::SpinLockIrq;
use super::{spinlock::SpinLockIrq, waker_set::WakerSet};
use crate::CpuOps;
use crate::sync::mutex::Mutex;
use core::cell::UnsafeCell;
use core::future::Future;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::task::{Context, Poll, Waker};

struct RwlockState<CPU: CpuOps> {
num_readers: SpinLockIrq<usize, CPU>,
writer_lock: Mutex<(), CPU>,
struct RwlockState {
num_readers: usize,
writer_active: bool,
waiting_writers: usize,
read_waiters: WakerSet,
write_waiters: WakerSet,
}

// Wakers may re-enter the lock, so extract notifications under the spinlock
// but run their callbacks only after releasing it and restoring interrupts.
enum Wakeup {
None,
Writer(Waker),
Readers(WakerSet),
}

impl Wakeup {
fn wake(self) {
match self {
Self::None => {}
Self::Writer(waker) => waker.wake(),
Self::Readers(mut waiters) => waiters.wake_all(),
}
}
}

impl RwlockState {
fn take_waiters(&mut self) -> Wakeup {
if self.writer_active {
return Wakeup::None;
}

if self.waiting_writers != 0 {
if self.num_readers == 0
&& let Some((waker, ())) = self.write_waiters.take_first()
{
return Wakeup::Writer(waker);
}
Wakeup::None
} else {
Wakeup::Readers(self.read_waiters.take_all())
}
}
}

// Pending acquisitions own only a waiter registration, never a lock or a
// reader slot. A writer remains counted even after its waker has been removed
// by a notification, until it acquires the lock or is cancelled.
struct AcquireFuture<'a, T: ?Sized, CPU: CpuOps> {
rwlock: &'a Rwlock<T, CPU>,
write: bool,
token: Option<u64>,
}

impl<T: ?Sized, CPU: CpuOps> Future for AcquireFuture<'_, T, CPU> {
type Output = ();

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
let waker = cx.waker().clone();
let mut state = this.rwlock.state.lock_save_irq();
let can_acquire = !state.writer_active
&& if this.write {
state.num_readers == 0
} else {
state.waiting_writers == 0
};

if can_acquire {
let retired = this.token.take().and_then(|token| {
if this.write {
state.waiting_writers -= 1;
state.write_waiters.take(token)
} else {
state.read_waiters.take(token)
}
});

if this.write {
state.writer_active = true;
} else {
state.num_readers = state
.num_readers
.checked_add(1)
.expect("rwlock reader count overflow");
}
drop(state);
drop(retired);
return Poll::Ready(());
}

// Refresh the waker on every pending poll. Notifications remove
// registrations, and another task may take the lock before we run.
let retired = if this.write {
if let Some(token) = this.token {
state.write_waiters.take(token)
} else {
state.waiting_writers = state
.waiting_writers
.checked_add(1)
.expect("rwlock waiting writer count overflow");
None
}
} else {
this.token.and_then(|token| state.read_waiters.take(token))
};
let waiters = if this.write {
&mut state.write_waiters
} else {
&mut state.read_waiters
};
this.token = Some(waiters.insert(waker, ()));
drop(state);
drop(retired);

Poll::Pending
}
}

impl<T: ?Sized, CPU: CpuOps> Drop for AcquireFuture<'_, T, CPU> {
fn drop(&mut self) {
if let Some(token) = self.token {
let mut state = self.rwlock.state.lock_save_irq();
let retired = if self.write {
state.waiting_writers -= 1;
state.write_waiters.take(token)
} else {
state.read_waiters.take(token)
};
// Cancellation of a queued or already-notified writer must not
// leave other waiters asleep while the lock is available.
let wakeup = state.take_waiters();
drop(state);
drop(retired);
wakeup.wake();
}
}
}

/// An asynchronous, rwlock primitive.
///
/// This rwlock can be used to protect shared data across asynchronous tasks.
/// `lock()` returns a future that resolves to a guard. When the guard is
/// dropped, the lock is released.
/// [`read()`](Self::read) and [`write()`](Self::write) return futures that resolve
/// to guards. When a guard is dropped, its lock is released.
///
/// Waiting writers take precedence over new readers.
pub struct Rwlock<T: ?Sized, CPU: CpuOps> {
state: RwlockState<CPU>,
state: SpinLockIrq<RwlockState, CPU>,
data: UnsafeCell<T>,
}

/// A guard that provides read-only access to the data in an `AsyncRwlock`.
/// A guard that provides read-only access to the data in a [`Rwlock`].
///
/// When an `AsyncRwlockReadGuard` is dropped, it automatically decreases the
/// read count and wakes up the next task if necessary.
Expand All @@ -30,10 +168,10 @@ pub struct AsyncRwlockReadGuard<'a, T: ?Sized, CPU: CpuOps> {
rwlock: &'a Rwlock<T, CPU>,
}

/// A guard that provides exclusive access to the data in an `AsyncRwlock`.
/// A guard that provides exclusive access to the data in a [`Rwlock`].
///
/// When an `AsyncRwlockWriteGuard` is dropped, it automatically releases the lock and
/// wakes up the next task.
/// When an `AsyncRwlockWriteGuard` is dropped, it automatically releases the lock
/// and wakes up the next task.
#[must_use = "if unused, the Rwlock will immediately unlock"]
pub struct AsyncRwlockWriteGuard<'a, T: ?Sized, CPU: CpuOps> {
rwlock: &'a Rwlock<T, CPU>,
Expand All @@ -43,10 +181,13 @@ impl<T, CPU: CpuOps> Rwlock<T, CPU> {
/// Creates a new asynchronous rwlock in an unlocked state.
pub fn new(data: T) -> Self {
Self {
state: RwlockState {
num_readers: SpinLockIrq::new(0),
writer_lock: Mutex::new(()),
},
state: SpinLockIrq::new(RwlockState {
num_readers: 0,
writer_active: false,
waiting_writers: 0,
read_waiters: WakerSet::new(),
write_waiters: WakerSet::new(),
}),
data: UnsafeCell::new(data),
}
}
Expand All @@ -66,11 +207,12 @@ impl<T: ?Sized, CPU: CpuOps> Rwlock<T, CPU> {
/// Returns a guard asynchronously. The guard is released when the
/// returned [`AsyncRwlockReadGuard`] is dropped.
pub async fn read(&self) -> AsyncRwlockReadGuard<'_, T, CPU> {
let mut num_readers = self.state.num_readers.lock_save_irq();
*num_readers += 1;
if *num_readers == 1 {
self.state.writer_lock.acquire().await;
AcquireFuture {
rwlock: self,
write: false,
token: None,
}
.await;
AsyncRwlockReadGuard { rwlock: self }
}

Expand All @@ -79,52 +221,61 @@ impl<T: ?Sized, CPU: CpuOps> Rwlock<T, CPU> {
/// Returns a guard asynchronously. The guard is released when the
/// returned [`AsyncRwlockWriteGuard`] is dropped.
pub async fn write(&self) -> AsyncRwlockWriteGuard<'_, T, CPU> {
self.state.writer_lock.acquire().await;
AcquireFuture {
rwlock: self,
write: true,
token: None,
}
.await;
AsyncRwlockWriteGuard { rwlock: self }
}
}

impl<T: ?Sized, CPU: CpuOps> Drop for AsyncRwlockReadGuard<'_, T, CPU> {
fn drop(&mut self) {
let mut num_readers = self.rwlock.state.num_readers.lock_save_irq();
*num_readers -= 1;
if *num_readers == 0 {
unsafe { self.rwlock.state.writer_lock.release() };
}
let mut state = self.rwlock.state.lock_save_irq();
state.num_readers -= 1;
let wakeup = state.take_waiters();
drop(state);
wakeup.wake();
}
}

impl<T: ?Sized, CPU: CpuOps> Deref for AsyncRwlockReadGuard<'_, T, CPU> {
type Target = T;
fn deref(&self) -> &T {
// SAFETY: This is safe because the existence of this guard guarantees
// we have read access to the data without any writers.
// SAFETY: The existence of this guard guarantees read access without
// any writers. Concurrent read access requires T: Sync to share a lock.
unsafe { &*self.rwlock.data.get() }
}
}

impl<T: ?Sized, CPU: CpuOps> Drop for AsyncRwlockWriteGuard<'_, T, CPU> {
fn drop(&mut self) {
unsafe { self.rwlock.state.writer_lock.release() };
let mut state = self.rwlock.state.lock_save_irq();
state.writer_active = false;
let wakeup = state.take_waiters();
drop(state);
wakeup.wake();
}
}

impl<T: ?Sized, CPU: CpuOps> Deref for AsyncRwlockWriteGuard<'_, T, CPU> {
type Target = T;
fn deref(&self) -> &T {
// SAFETY: This is safe because the existence of this guard guarantees
// we have exclusive access to the data.
// SAFETY: The existence of this guard guarantees exclusive access.
unsafe { &*self.rwlock.data.get() }
}
}

impl<T: ?Sized, CPU: CpuOps> DerefMut for AsyncRwlockWriteGuard<'_, T, CPU> {
fn deref_mut(&mut self) -> &mut T {
// SAFETY: This is safe because the existence of this guard guarantees
// we have exclusive access to the data.
// SAFETY: The existence of this guard guarantees exclusive access.
unsafe { &mut *self.rwlock.data.get() }
}
}

// SAFETY: Moving an owned lock transfers ownership of T. Sharing a lock also
// permits concurrent readers to access &T, so T must be Sync as well as Send.
unsafe impl<T: ?Sized + Send, CPU: CpuOps> Send for Rwlock<T, CPU> {}
unsafe impl<T: ?Sized + Send, CPU: CpuOps> Sync for Rwlock<T, CPU> {}
unsafe impl<T: ?Sized + Send + Sync, CPU: CpuOps> Sync for Rwlock<T, CPU> {}
Loading
Loading