eater
/
async-std
forked from mirror/async-std
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Simplify RwLock using WakerSet (#440)

Browse Source
poc-serde-support
Stjepan Glavina 5 years ago committed by GitHub
parent 78614c6c1d
commit 20cdf73bb0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 115 additions and 211 deletions
Show Stats Download Patch File Download Diff File

290
src/sync/rwlock.rs
Unescape Escape View File

@ -1,26 +1,21 @@
use std::cell::UnsafeCell; use std::cell::UnsafeCell;
use std::fmt; use std::fmt;
use std::isize;
use std::ops::{Deref, DerefMut}; use std::ops::{Deref, DerefMut};
use std::pin::Pin; use std::pin::Pin;
use std::process;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use slab::Slab;
use crate::future::Future; use crate::future::Future;
use crate::task::{Context, Poll, Waker}; use crate::sync::WakerSet;
use crate::task::{Context, Poll};
/// Set if a write lock is held. /// Set if a write lock is held.
#[allow(clippy::identity_op)] #[allow(clippy::identity_op)]
const WRITE_LOCK: usize = 1 << 0; const WRITE_LOCK: usize = 1 << 0;
/// Set if there are read operations blocked on the lock.
const BLOCKED_READS: usize = 1 << 1;
/// Set if there are write operations blocked on the lock.
const BLOCKED_WRITES: usize = 1 << 2;
/// The value of a single blocked read contributing to the read count. /// The value of a single blocked read contributing to the read count.
const ONE_READ: usize = 1 << 3; const ONE_READ: usize = 1 << 1;
/// The bits in which the read count is stored. /// The bits in which the read count is stored.
const READ_COUNT_MASK: usize = !(ONE_READ - 1); const READ_COUNT_MASK: usize = !(ONE_READ - 1);
@ -56,8 +51,8 @@ const READ_COUNT_MASK: usize = !(ONE_READ - 1);
/// ``` /// ```
pub struct RwLock<T> { pub struct RwLock<T> {
state: AtomicUsize, state: AtomicUsize,
reads: std::sync::Mutex<Slab<Option<Waker>>>, read_wakers: WakerSet,
writes: std::sync::Mutex<Slab<Option<Waker>>>, write_wakers: WakerSet,
value: UnsafeCell<T>, value: UnsafeCell<T>,
} }
@ -77,8 +72,8 @@ impl<T> RwLock<T> {
pub fn new(t: T) -> RwLock<T> { pub fn new(t: T) -> RwLock<T> {
RwLock { RwLock {
state: AtomicUsize::new(0), state: AtomicUsize::new(0),
reads: std::sync::Mutex::new(Slab::new()), read_wakers: WakerSet::new(),
writes: std::sync::Mutex::new(Slab::new()), write_wakers: WakerSet::new(),
value: UnsafeCell::new(t), value: UnsafeCell::new(t),
} }
} }
@ -104,100 +99,61 @@ impl<T> RwLock<T> {
/// # }) /// # })
/// ``` /// ```
pub async fn read(&self) -> RwLockReadGuard<'_, T> { pub async fn read(&self) -> RwLockReadGuard<'_, T> {
pub struct LockFuture<'a, T> { pub struct ReadFuture<'a, T> {
lock: &'a RwLock<T>, lock: &'a RwLock<T>,
opt_key: Option<usize>, opt_key: Option<usize>,
acquired: bool,
} }
impl<'a, T> Future for LockFuture<'a, T> { impl<'a, T> Future for ReadFuture<'a, T> {
type Output = RwLockReadGuard<'a, T>; type Output = RwLockReadGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.lock.try_read() { let poll = match self.lock.try_read() {
Some(guard) => { Some(guard) => Poll::Ready(guard),
self.acquired = true;
Poll::Ready(guard)
}
None => { None => {
let mut reads = self.lock.reads.lock().unwrap(); // Insert this lock operation.
// Register the current task.
match self.opt_key { match self.opt_key {
None => { None => self.opt_key = Some(self.lock.read_wakers.insert(cx)),
// Insert a new entry into the list of blocked reads. Some(key) => self.lock.read_wakers.update(key, cx),
let w = cx.waker().clone();
let key = reads.insert(Some(w));
self.opt_key = Some(key);
if reads.len() == 1 {
self.lock.state.fetch_or(BLOCKED_READS, Ordering::Relaxed);
}
}
Some(key) => {
// There is already an entry in the list of blocked reads. Just
// reset the waker if it was removed.
if reads[key].is_none() {
let w = cx.waker().clone();
reads[key] = Some(w);
}
}
} }
// Try locking again because it's possible the lock got unlocked just // Try locking again because it's possible the lock got unlocked just
// before the current task was registered as a blocked task. // before the current task was inserted into the waker set.
match self.lock.try_read() { match self.lock.try_read() {
Some(guard) => { Some(guard) => Poll::Ready(guard),
self.acquired = true;
Poll::Ready(guard)
}
None => Poll::Pending, None => Poll::Pending,
} }
} }
};
if poll.is_ready() {
// If the current task is in the set, remove it.
if let Some(key) = self.opt_key.take() {
self.lock.read_wakers.complete(key);
}
} }
poll
} }
} }
impl<T> Drop for LockFuture<'_, T> { impl<T> Drop for ReadFuture<'_, T> {
fn drop(&mut self) { fn drop(&mut self) {
// If the current task is still in the set, that means it is being cancelled now.
if let Some(key) = self.opt_key { if let Some(key) = self.opt_key {
let mut reads = self.lock.reads.lock().unwrap(); self.lock.read_wakers.cancel(key);
let opt_waker = reads.remove(key);
if reads.is_empty() {
self.lock.state.fetch_and(!BLOCKED_READS, Ordering::Relaxed);
}
if opt_waker.is_none() { // If there are no active readers, wake one of the writers.
// We were awoken. Wake up another blocked read. if self.lock.state.load(Ordering::SeqCst) & READ_COUNT_MASK == 0 {
if let Some((_, opt_waker)) = reads.iter_mut().next() { self.lock.write_wakers.notify_one();
if let Some(w) = opt_waker.take() {
w.wake();
return;
}
}
drop(reads);
if !self.acquired {
// We didn't acquire the lock and didn't wake another blocked read.
// Wake a blocked write instead.
let mut writes = self.lock.writes.lock().unwrap();
if let Some((_, opt_waker)) = writes.iter_mut().next() {
if let Some(w) = opt_waker.take() {
w.wake();
return;
}
}
}
} }
} }
} }
} }
LockFuture { ReadFuture {
lock: self, lock: self,
opt_key: None, opt_key: None,
acquired: false,
} }
.await .await
} }
@ -226,7 +182,7 @@ impl<T> RwLock<T> {
/// # }) /// # })
/// ``` /// ```
pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> { pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
let mut state = self.state.load(Ordering::Acquire); let mut state = self.state.load(Ordering::SeqCst);
loop { loop {
// If a write lock is currently held, then a read lock cannot be acquired. // If a write lock is currently held, then a read lock cannot be acquired.
@ -234,12 +190,17 @@ impl<T> RwLock<T> {
return None; return None;
} }
// Make sure the number of readers doesn't overflow.
if state > isize::MAX as usize {
process::abort();
}
// Increment the number of active reads. // Increment the number of active reads.
match self.state.compare_exchange_weak( match self.state.compare_exchange_weak(
state, state,
state + ONE_READ, state + ONE_READ,
Ordering::AcqRel, Ordering::SeqCst,
Ordering::Acquire, Ordering::SeqCst,
) { ) {
Ok(_) => return Some(RwLockReadGuard(self)), Ok(_) => return Some(RwLockReadGuard(self)),
Err(s) => state = s, Err(s) => state = s,
@ -268,99 +229,59 @@ impl<T> RwLock<T> {
/// # }) /// # })
/// ``` /// ```
pub async fn write(&self) -> RwLockWriteGuard<'_, T> { pub async fn write(&self) -> RwLockWriteGuard<'_, T> {
pub struct LockFuture<'a, T> { pub struct WriteFuture<'a, T> {
lock: &'a RwLock<T>, lock: &'a RwLock<T>,
opt_key: Option<usize>, opt_key: Option<usize>,
acquired: bool,
} }
impl<'a, T> Future for LockFuture<'a, T> { impl<'a, T> Future for WriteFuture<'a, T> {
type Output = RwLockWriteGuard<'a, T>; type Output = RwLockWriteGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.lock.try_write() { let poll = match self.lock.try_write() {
Some(guard) => { Some(guard) => Poll::Ready(guard),
self.acquired = true;
Poll::Ready(guard)
}
None => { None => {
let mut writes = self.lock.writes.lock().unwrap(); // Insert this lock operation.
// Register the current task.
match self.opt_key { match self.opt_key {
None => { None => self.opt_key = Some(self.lock.write_wakers.insert(cx)),
// Insert a new entry into the list of blocked writes. Some(key) => self.lock.write_wakers.update(key, cx),
let w = cx.waker().clone();
let key = writes.insert(Some(w));
self.opt_key = Some(key);
if writes.len() == 1 {
self.lock.state.fetch_or(BLOCKED_WRITES, Ordering::Relaxed);
}
}
Some(key) => {
// There is already an entry in the list of blocked writes. Just
// reset the waker if it was removed.
if writes[key].is_none() {
let w = cx.waker().clone();
writes[key] = Some(w);
}
}
} }
// Try locking again because it's possible the lock got unlocked just // Try locking again because it's possible the lock got unlocked just
// before the current task was registered as a blocked task. // before the current task was inserted into the waker set.
match self.lock.try_write() { match self.lock.try_write() {
Some(guard) => { Some(guard) => Poll::Ready(guard),
self.acquired = true;
Poll::Ready(guard)
}
None => Poll::Pending, None => Poll::Pending,
} }
} }
};
if poll.is_ready() {
// If the current task is in the set, remove it.
if let Some(key) = self.opt_key.take() {
self.lock.write_wakers.complete(key);
}
} }
poll
} }
} }
impl<T> Drop for LockFuture<'_, T> { impl<T> Drop for WriteFuture<'_, T> {
fn drop(&mut self) { fn drop(&mut self) {
// If the current task is still in the set, that means it is being cancelled now.
if let Some(key) = self.opt_key { if let Some(key) = self.opt_key {
let mut writes = self.lock.writes.lock().unwrap(); if !self.lock.write_wakers.cancel(key) {
let opt_waker = writes.remove(key); // If no other blocked reader was notified, notify all readers.
self.lock.read_wakers.notify_all();
if writes.is_empty() {
self.lock
.state
.fetch_and(!BLOCKED_WRITES, Ordering::Relaxed);
}
if opt_waker.is_none() && !self.acquired {
// We were awoken but didn't acquire the lock. Wake up another write.
if let Some((_, opt_waker)) = writes.iter_mut().next() {
if let Some(w) = opt_waker.take() {
w.wake();
return;
}
}
drop(writes);
// There are no blocked writes. Wake a blocked read instead.
let mut reads = self.lock.reads.lock().unwrap();
if let Some((_, opt_waker)) = reads.iter_mut().next() {
if let Some(w) = opt_waker.take() {
w.wake();
return;
}
}
} }
} }
} }
} }
LockFuture { WriteFuture {
lock: self, lock: self,
opt_key: None, opt_key: None,
acquired: false,
} }
.await .await
} }
@ -389,24 +310,10 @@ impl<T> RwLock<T> {
/// # }) /// # })
/// ``` /// ```
pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> { pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
let mut state = self.state.load(Ordering::Acquire); if self.state.compare_and_swap(0, WRITE_LOCK, Ordering::SeqCst) == 0 {
Some(RwLockWriteGuard(self))
loop { } else {
// If any kind of lock is currently held, then a write lock cannot be acquired. None
if state & (WRITE_LOCK | READ_COUNT_MASK) != 0 {
return None;
}
// Set the write lock.
match self.state.compare_exchange_weak(
state,
state | WRITE_LOCK,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return Some(RwLockWriteGuard(self)),
Err(s) => state = s,
}
} }
} }
@ -449,18 +356,15 @@ impl<T> RwLock<T> {
impl<T: fmt::Debug> fmt::Debug for RwLock<T> { impl<T: fmt::Debug> fmt::Debug for RwLock<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.try_read() { struct Locked;
None => { impl fmt::Debug for Locked {
struct LockedPlaceholder; fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
impl fmt::Debug for LockedPlaceholder { f.write_str("<locked>")
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("<locked>")
}
}
f.debug_struct("RwLock")
.field("data", &LockedPlaceholder)
.finish()
} }
}
match self.try_read() {
None => f.debug_struct("RwLock").field("data", &Locked).finish(),
Some(guard) => f.debug_struct("RwLock").field("data", &&*guard).finish(), Some(guard) => f.debug_struct("RwLock").field("data", &&*guard).finish(),
} }
} }
@ -486,18 +390,11 @@ unsafe impl<T: Sync> Sync for RwLockReadGuard<'_, T> {}
impl<T> Drop for RwLockReadGuard<'_, T> { impl<T> Drop for RwLockReadGuard<'_, T> {
fn drop(&mut self) { fn drop(&mut self) {
let state = self.0.state.fetch_sub(ONE_READ, Ordering::AcqRel); let state = self.0.state.fetch_sub(ONE_READ, Ordering::SeqCst);
// If this was the last read and there are blocked writes, wake one of them up.
if (state & READ_COUNT_MASK) == ONE_READ && state & BLOCKED_WRITES != 0 {
let mut writes = self.0.writes.lock().unwrap();
if let Some((_, opt_waker)) = writes.iter_mut().next() { // If this was the last read, wake one of the writers.
// If there is no waker in this entry, that means it was already woken. if state & READ_COUNT_MASK == ONE_READ {
if let Some(w) = opt_waker.take() { self.0.write_wakers.notify_one();
w.wake();
}
}
} }
} }
} }
@ -530,25 +427,12 @@ unsafe impl<T: Sync> Sync for RwLockWriteGuard<'_, T> {}
impl<T> Drop for RwLockWriteGuard<'_, T> { impl<T> Drop for RwLockWriteGuard<'_, T> {
fn drop(&mut self) { fn drop(&mut self) {
let state = self.0.state.fetch_and(!WRITE_LOCK, Ordering::AcqRel); self.0.state.store(0, Ordering::SeqCst);
let mut guard = None;
// Check if there are any blocked reads or writes.
if state & BLOCKED_READS != 0 {
guard = Some(self.0.reads.lock().unwrap());
} else if state & BLOCKED_WRITES != 0 {
guard = Some(self.0.writes.lock().unwrap());
}
// Wake up a single blocked task. // Notify all blocked readers.
if let Some(mut guard) = guard { if !self.0.read_wakers.notify_all() {
if let Some((_, opt_waker)) = guard.iter_mut().next() { // If there were no blocked readers, notify a blocked writer.
// If there is no waker in this entry, that means it was already woken. self.0.write_wakers.notify_one();
if let Some(w) = opt_waker.take() {
w.wake();
}
}
} }
} }
} }

36
src/sync/waker_set.rs
Unescape Escape View File

@ -95,8 +95,11 @@ impl WakerSet {
} }
/// Removes the waker of a cancelled operation. /// Removes the waker of a cancelled operation.
pub fn cancel(&self, key: usize) { ///
/// Returns `true` if another blocked operation from the set was notified.
pub fn cancel(&self, key: usize) -> bool {
let mut inner = self.lock(); let mut inner = self.lock();
if inner.entries.remove(key).is_none() { if inner.entries.remove(key).is_none() {
inner.none_count -= 1; inner.none_count -= 1;
@ -107,33 +110,45 @@ impl WakerSet {
w.wake(); w.wake();
inner.none_count += 1; inner.none_count += 1;
} }
return true;
} }
} }
false
} }
/// Notifies one blocked operation. /// Notifies one blocked operation.
///
/// Returns `true` if an operation was notified.
#[inline] #[inline]
pub fn notify_one(&self) { pub fn notify_one(&self) -> bool {
// Use `SeqCst` ordering to synchronize with `Lock::drop()`. // Use `SeqCst` ordering to synchronize with `Lock::drop()`.
if self.flag.load(Ordering::SeqCst) & NOTIFY_ONE != 0 { if self.flag.load(Ordering::SeqCst) & NOTIFY_ONE != 0 {
self.notify(false); self.notify(false)
} else {
false
} }
} }
/// Notifies all blocked operations. /// Notifies all blocked operations.
// TODO: Delete this attribute when `crate::sync::channel()` is stabilized. ///
#[cfg(feature = "unstable")] /// Returns `true` if at least one operation was notified.
#[inline] #[inline]
pub fn notify_all(&self) { pub fn notify_all(&self) -> bool {
// Use `SeqCst` ordering to synchronize with `Lock::drop()`. // Use `SeqCst` ordering to synchronize with `Lock::drop()`.
if self.flag.load(Ordering::SeqCst) & NOTIFY_ALL != 0 { if self.flag.load(Ordering::SeqCst) & NOTIFY_ALL != 0 {
self.notify(true); self.notify(true)
} else {
false
} }
} }
/// Notifies blocked operations, either one or all of them. /// Notifies blocked operations, either one or all of them.
fn notify(&self, all: bool) { ///
/// Returns `true` if at least one operation was notified.
fn notify(&self, all: bool) -> bool {
let mut inner = &mut *self.lock(); let mut inner = &mut *self.lock();
let mut notified = false;
for (_, opt_waker) in inner.entries.iter_mut() { for (_, opt_waker) in inner.entries.iter_mut() {
// If there is no waker in this entry, that means it was already woken. // If there is no waker in this entry, that means it was already woken.
@ -141,10 +156,15 @@ impl WakerSet {
w.wake(); w.wake();
inner.none_count += 1; inner.none_count += 1;
} }
notified = true;
if !all { if !all {
break; break;
} }
} }
notified
} }
/// Locks the list of entries. /// Locks the list of entries.

Write Preview
Loading…
Cancel
Save