feat: use async-lock for RwLock and Barrier

pull/933/head
Koxiaet 3 years ago committed by GitHub
parent 1184a7f97f
commit 4a3f963810
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -52,8 +52,8 @@ std = [
"slab",
"wasm-bindgen-futures",
"futures-channel",
"async-mutex",
"async-channel",
"async-lock",
]
alloc = [
"futures-core/alloc",
@ -64,7 +64,7 @@ tokio03 = ["async-global-executor/tokio03"]
[dependencies]
async-attributes = { version = "1.1.1", optional = true }
async-mutex = { version = "1.1.3", optional = true }
async-lock = { version = "2.3.0", optional = true }
crossbeam-utils = { version = "0.8.0", optional = true }
futures-core = { version = "0.3.4", optional = true, default-features = false }
futures-io = { version = "0.3.4", optional = true }

@ -1,229 +0,0 @@
use crate::sync::{Condvar,Mutex};
/// A barrier enables multiple tasks to synchronize the beginning
/// of some computation.
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::{Arc, Barrier};
/// use async_std::task;
///
/// let mut handles = Vec::with_capacity(10);
/// let barrier = Arc::new(Barrier::new(10));
/// for _ in 0..10 {
/// let c = barrier.clone();
/// // The same messages will be printed together.
/// // You will NOT see any interleaving.
/// handles.push(task::spawn(async move {
/// println!("before wait");
/// c.wait().await;
/// println!("after wait");
/// }));
/// }
/// // Wait for the other futures to finish.
/// for handle in handles {
/// handle.await;
/// }
/// # });
/// ```
#[cfg(feature = "unstable")]
#[cfg_attr(feature = "docs", doc(cfg(unstable)))]
#[derive(Debug)]
pub struct Barrier {
state: Mutex<BarrierState>,
cvar: Condvar,
num_tasks: usize,
}
// The inner state of a double barrier
#[derive(Debug)]
struct BarrierState {
count: usize,
generation_id: usize,
}
/// A `BarrierWaitResult` is returned by `wait` when all threads in the `Barrier` have rendezvoused.
///
/// [`wait`]: struct.Barrier.html#method.wait
/// [`Barrier`]: struct.Barrier.html
///
/// # Examples
///
/// ```
/// use async_std::sync::Barrier;
///
/// let barrier = Barrier::new(1);
/// let barrier_wait_result = barrier.wait();
/// ```
#[cfg(feature = "unstable")]
#[cfg_attr(feature = "docs", doc(cfg(unstable)))]
#[derive(Debug, Clone)]
pub struct BarrierWaitResult(bool);
impl Barrier {
/// Creates a new barrier that can block a given number of tasks.
///
/// A barrier will block `n`-1 tasks which call [`wait`] and then wake up
/// all tasks at once when the `n`th task calls [`wait`].
///
/// [`wait`]: #method.wait
///
/// # Examples
///
/// ```
/// use std::sync::Barrier;
///
/// let barrier = Barrier::new(10);
/// ```
pub fn new(n: usize) -> Barrier {
Barrier {
state: Mutex::new(BarrierState {
count: 0,
generation_id: 1,
}),
cvar: Condvar::new(),
num_tasks: n,
}
}
/// Blocks the current task until all tasks have rendezvoused here.
///
/// Barriers are re-usable after all tasks have rendezvoused once, and can
/// be used continuously.
///
/// A single (arbitrary) task will receive a [`BarrierWaitResult`] that
/// returns `true` from [`is_leader`] when returning from this function, and
/// all other tasks will receive a result that will return `false` from
/// [`is_leader`].
///
/// [`BarrierWaitResult`]: struct.BarrierWaitResult.html
/// [`is_leader`]: struct.BarrierWaitResult.html#method.is_leader
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::{Arc, Barrier};
/// use async_std::task;
///
/// let mut handles = Vec::with_capacity(10);
/// let barrier = Arc::new(Barrier::new(10));
/// for _ in 0..10 {
/// let c = barrier.clone();
/// // The same messages will be printed together.
/// // You will NOT see any interleaving.
/// handles.push(task::spawn(async move {
/// println!("before wait");
/// c.wait().await;
/// println!("after wait");
/// }));
/// }
/// // Wait for the other futures to finish.
/// for handle in handles {
/// handle.await;
/// }
/// # });
/// ```
pub async fn wait(&self) -> BarrierWaitResult {
let mut state = self.state.lock().await;
let local_gen = state.generation_id;
state.count += 1;
if state.count < self.num_tasks {
while local_gen == state.generation_id && state.count < self.num_tasks {
state = self.cvar.wait(state).await;
}
BarrierWaitResult(false)
} else {
state.count = 0;
state.generation_id = state.generation_id.wrapping_add(1);
self.cvar.notify_all();
BarrierWaitResult(true)
}
}
}
impl BarrierWaitResult {
/// Returns `true` if this task from [`wait`] is the "leader task".
///
/// Only one task will have `true` returned from their result, all other
/// tasks will have `false` returned.
///
/// [`wait`]: struct.Barrier.html#method.wait
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::Barrier;
///
/// let barrier = Barrier::new(1);
/// let barrier_wait_result = barrier.wait().await;
/// println!("{:?}", barrier_wait_result.is_leader());
/// # });
/// ```
pub fn is_leader(&self) -> bool {
self.0
}
}
#[cfg(all(test, not(target_os = "unknown")))]
mod test {
use futures::channel::mpsc::unbounded;
use futures::sink::SinkExt;
use futures::stream::StreamExt;
use crate::sync::{Arc, Barrier};
use crate::task;
#[test]
fn test_barrier() {
// NOTE(dignifiedquire): Based on the test in std, I was seeing some
// race conditions, so running it in a loop to make sure things are
// solid.
for _ in 0..1_000 {
task::block_on(async move {
const N: usize = 10;
let barrier = Arc::new(Barrier::new(N));
let (tx, mut rx) = unbounded();
for _ in 0..N - 1 {
let c = barrier.clone();
let mut tx = tx.clone();
task::spawn(async move {
let res = c.wait().await;
tx.send(res.is_leader()).await.unwrap();
});
}
// At this point, all spawned threads should be blocked,
// so we shouldn't get anything from the port
let res = rx.try_next();
assert!(match res {
Err(_err) => true,
_ => false,
});
let mut leader_found = barrier.wait().await.is_leader();
// Now, the barrier is cleared and we should get data.
for _ in 0..N - 1 {
if rx.next().await.unwrap() {
assert!(!leader_found);
leader_found = true;
}
}
assert!(leader_found);
});
}
}
}

@ -177,22 +177,20 @@
pub use std::sync::{Arc, Weak};
#[doc(inline)]
pub use async_mutex::{Mutex, MutexGuard};
pub use async_lock::{Mutex, MutexGuard, MutexGuardArc};
pub use rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard};
mod rwlock;
#[doc(inline)]
pub use async_lock::{RwLock, RwLockReadGuard, RwLockUpgradableReadGuard, RwLockWriteGuard};
cfg_unstable! {
pub use barrier::{Barrier, BarrierWaitResult};
pub use async_lock::{Barrier, BarrierWaitResult};
#[allow(deprecated)]
pub use channel::{channel, Sender, Receiver, RecvError, TryRecvError, TrySendError};
pub use condvar::Condvar;
pub(crate) use waker_set::WakerSet;
mod barrier;
mod condvar;
mod channel;
}
pub(crate) mod waker_set;
pub(crate) use waker_set::WakerSet;
pub(crate) mod waker_set;
}

@ -1,463 +0,0 @@
use std::cell::UnsafeCell;
use std::fmt;
use std::future::Future;
use std::isize;
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::process;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::sync::WakerSet;
use crate::task::{Context, Poll};
/// Set if a write lock is held.
#[allow(clippy::identity_op)]
const WRITE_LOCK: usize = 1 << 0;
/// The value of a single blocked read contributing to the read count.
const ONE_READ: usize = 1 << 1;
/// The bits in which the read count is stored.
const READ_COUNT_MASK: usize = !(ONE_READ - 1);
/// A reader-writer lock for protecting shared data.
///
/// This type is an async version of [`std::sync::RwLock`].
///
/// [`std::sync::RwLock`]: https://doc.rust-lang.org/std/sync/struct.RwLock.html
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(5);
///
/// // Multiple read locks can be held at a time.
/// let r1 = lock.read().await;
/// let r2 = lock.read().await;
/// assert_eq!(*r1, 5);
/// assert_eq!(*r2, 5);
/// drop((r1, r2));
///
/// // Only one write locks can be held at a time.
/// let mut w = lock.write().await;
/// *w += 1;
/// assert_eq!(*w, 6);
/// #
/// # })
/// ```
pub struct RwLock<T: ?Sized> {
state: AtomicUsize,
read_wakers: WakerSet,
write_wakers: WakerSet,
value: UnsafeCell<T>,
}
unsafe impl<T: ?Sized + Send> Send for RwLock<T> {}
unsafe impl<T: ?Sized + Send + Sync> Sync for RwLock<T> {}
impl<T> RwLock<T> {
/// Creates a new reader-writer lock.
///
/// # Examples
///
/// ```
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(0);
/// ```
pub fn new(t: T) -> RwLock<T> {
RwLock {
state: AtomicUsize::new(0),
read_wakers: WakerSet::new(),
write_wakers: WakerSet::new(),
value: UnsafeCell::new(t),
}
}
}
impl<T: ?Sized> RwLock<T> {
/// Acquires a read lock.
///
/// Returns a guard that releases the lock when dropped.
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(1);
///
/// let n = lock.read().await;
/// assert_eq!(*n, 1);
///
/// assert!(lock.try_read().is_some());
/// #
/// # })
/// ```
pub async fn read(&self) -> RwLockReadGuard<'_, T> {
pub struct ReadFuture<'a, T: ?Sized> {
lock: &'a RwLock<T>,
opt_key: Option<usize>,
}
impl<'a, T: ?Sized> Future for ReadFuture<'a, T> {
type Output = RwLockReadGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
// If the current task is in the set, remove it.
if let Some(key) = self.opt_key.take() {
self.lock.read_wakers.remove(key);
}
// Try acquiring a read lock.
match self.lock.try_read() {
Some(guard) => return Poll::Ready(guard),
None => {
// Insert this lock operation.
self.opt_key = Some(self.lock.read_wakers.insert(cx));
// If the lock is still acquired for writing, return.
if self.lock.state.load(Ordering::SeqCst) & WRITE_LOCK != 0 {
return Poll::Pending;
}
}
}
}
}
}
impl<T: ?Sized> Drop for ReadFuture<'_, T> {
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 {
self.lock.read_wakers.cancel(key);
// If there are no active readers, notify a blocked writer if none were
// notified already.
if self.lock.state.load(Ordering::SeqCst) & READ_COUNT_MASK == 0 {
self.lock.write_wakers.notify_any();
}
}
}
}
ReadFuture {
lock: self,
opt_key: None,
}
.await
}
/// Attempts to acquire a read lock.
///
/// If a read lock could not be acquired at this time, then [`None`] is returned. Otherwise, a
/// guard is returned that releases the lock when dropped.
///
/// [`None`]: https://doc.rust-lang.org/std/option/enum.Option.html#variant.None
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(1);
///
/// let n = lock.read().await;
/// assert_eq!(*n, 1);
///
/// assert!(lock.try_read().is_some());
/// #
/// # })
/// ```
pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
let mut state = self.state.load(Ordering::SeqCst);
loop {
// If a write lock is currently held, then a read lock cannot be acquired.
if state & WRITE_LOCK != 0 {
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.
match self.state.compare_exchange_weak(
state,
state + ONE_READ,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(_) => return Some(RwLockReadGuard(self)),
Err(s) => state = s,
}
}
}
/// Acquires a write lock.
///
/// Returns a guard that releases the lock when dropped.
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(1);
///
/// let mut n = lock.write().await;
/// *n = 2;
///
/// assert!(lock.try_read().is_none());
/// #
/// # })
/// ```
pub async fn write(&self) -> RwLockWriteGuard<'_, T> {
pub struct WriteFuture<'a, T: ?Sized> {
lock: &'a RwLock<T>,
opt_key: Option<usize>,
}
impl<'a, T: ?Sized> Future for WriteFuture<'a, T> {
type Output = RwLockWriteGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
// If the current task is in the set, remove it.
if let Some(key) = self.opt_key.take() {
self.lock.write_wakers.remove(key);
}
// Try acquiring a write lock.
match self.lock.try_write() {
Some(guard) => return Poll::Ready(guard),
None => {
// Insert this lock operation.
self.opt_key = Some(self.lock.write_wakers.insert(cx));
// If the lock is still acquired for reading or writing, return.
if self.lock.state.load(Ordering::SeqCst) != 0 {
return Poll::Pending;
}
}
}
}
}
}
impl<T: ?Sized> Drop for WriteFuture<'_, T> {
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 !self.lock.write_wakers.cancel(key) {
// If no other blocked reader was notified, notify all readers.
self.lock.read_wakers.notify_all();
}
}
}
}
WriteFuture {
lock: self,
opt_key: None,
}
.await
}
/// Attempts to acquire a write lock.
///
/// If a write lock could not be acquired at this time, then [`None`] is returned. Otherwise, a
/// guard is returned that releases the lock when dropped.
///
/// [`None`]: https://doc.rust-lang.org/std/option/enum.Option.html#variant.None
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(1);
///
/// let n = lock.read().await;
/// assert_eq!(*n, 1);
///
/// assert!(lock.try_write().is_none());
/// #
/// # })
/// ```
pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
if self
.state
.compare_exchange(0, WRITE_LOCK, Ordering::SeqCst, Ordering::SeqCst)
.is_ok()
{
Some(RwLockWriteGuard(self))
} else {
None
}
}
/// Consumes the lock, returning the underlying data.
///
/// # Examples
///
/// ```
/// use async_std::sync::RwLock;
///
/// let lock = RwLock::new(10);
/// assert_eq!(lock.into_inner(), 10);
/// ```
pub fn into_inner(self) -> T
where
T: Sized,
{
self.value.into_inner()
}
/// Returns a mutable reference to the underlying data.
///
/// Since this call borrows the lock mutably, no actual locking takes place -- the mutable
/// borrow statically guarantees no locks exist.
///
/// # Examples
///
/// ```
/// # async_std::task::block_on(async {
/// #
/// use async_std::sync::RwLock;
///
/// let mut lock = RwLock::new(0);
/// *lock.get_mut() = 10;
/// assert_eq!(*lock.write().await, 10);
/// #
/// # })
/// ```
pub fn get_mut(&mut self) -> &mut T {
unsafe { &mut *self.value.get() }
}
}
impl<T: ?Sized + fmt::Debug> fmt::Debug for RwLock<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
struct Locked;
impl fmt::Debug for Locked {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("<locked>")
}
}
match self.try_read() {
None => f.debug_struct("RwLock").field("data", &Locked).finish(),
Some(guard) => f.debug_struct("RwLock").field("data", &&*guard).finish(),
}
}
}
impl<T> From<T> for RwLock<T> {
fn from(val: T) -> RwLock<T> {
RwLock::new(val)
}
}
impl<T: ?Sized + Default> Default for RwLock<T> {
fn default() -> RwLock<T> {
RwLock::new(Default::default())
}
}
/// A guard that releases the read lock when dropped.
pub struct RwLockReadGuard<'a, T: ?Sized>(&'a RwLock<T>);
unsafe impl<T: ?Sized + Send> Send for RwLockReadGuard<'_, T> {}
unsafe impl<T: ?Sized + Sync> Sync for RwLockReadGuard<'_, T> {}
impl<T: ?Sized> Drop for RwLockReadGuard<'_, T> {
fn drop(&mut self) {
let state = self.0.state.fetch_sub(ONE_READ, Ordering::SeqCst);
// If this was the last reader, notify a blocked writer if none were notified already.
if state & READ_COUNT_MASK == ONE_READ {
self.0.write_wakers.notify_any();
}
}
}
impl<T: ?Sized + fmt::Debug> fmt::Debug for RwLockReadGuard<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T: ?Sized + fmt::Display> fmt::Display for RwLockReadGuard<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
(**self).fmt(f)
}
}
impl<T: ?Sized> Deref for RwLockReadGuard<'_, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
/// A guard that releases the write lock when dropped.
pub struct RwLockWriteGuard<'a, T: ?Sized>(&'a RwLock<T>);
unsafe impl<T: ?Sized + Send> Send for RwLockWriteGuard<'_, T> {}
unsafe impl<T: ?Sized + Sync> Sync for RwLockWriteGuard<'_, T> {}
impl<T: ?Sized> Drop for RwLockWriteGuard<'_, T> {
fn drop(&mut self) {
self.0.state.store(0, Ordering::SeqCst);
// Notify all blocked readers.
if !self.0.read_wakers.notify_all() {
// If there were no blocked readers, notify a blocked writer if none were notified
// already.
self.0.write_wakers.notify_any();
}
}
}
impl<T: ?Sized + fmt::Debug> fmt::Debug for RwLockWriteGuard<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T: ?Sized + fmt::Display> fmt::Display for RwLockWriteGuard<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
(**self).fmt(f)
}
}
impl<T: ?Sized> Deref for RwLockWriteGuard<'_, T> {
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
impl<T: ?Sized> DerefMut for RwLockWriteGuard<'_, T> {
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.0.value.get() }
}
}

@ -127,21 +127,6 @@ impl WakerSet {
false
}
/// Notifies a blocked operation if none have been notified already.
///
/// Returns `true` if an operation was notified.
#[inline]
pub fn notify_any(&self) -> bool {
// Use `SeqCst` ordering to synchronize with `Lock::drop()`.
let flag = self.flag.load(Ordering::SeqCst);
if flag & NOTIFIED == 0 && flag & NOTIFIABLE != 0 {
self.notify(Notify::Any)
} else {
false
}
}
/// Notifies one additional blocked operation.
///
/// Returns `true` if an operation was notified.

Loading…
Cancel
Save