use crate::inner::Inner;
use std::fmt::Debug;
use std::future::Future;
use std::pin::Pin;
use std::ptr::null_mut;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
#[derive(Debug)]
pub struct RwLock<T: ?Sized> {
readers: AtomicUsize,
inner: Inner<T>,
}
impl<T> RwLock<T> {
#[inline]
pub const fn new(data: T) -> RwLock<T> {
RwLock {
readers: AtomicUsize::new(0),
inner: Inner::new(data),
}
}
}
impl<T: ?Sized> RwLock<T> {
#[inline]
pub fn write(&self) -> RwLockWriteGuardFuture<T> {
RwLockWriteGuardFuture {
mutex: &self,
is_realized: false,
}
}
#[inline]
pub fn write_owned(self: &Arc<Self>) -> RwLockWriteOwnedGuardFuture<T> {
RwLockWriteOwnedGuardFuture {
mutex: self.clone(),
is_realized: false,
}
}
#[inline]
pub fn read(&self) -> RwLockReadGuardFuture<T> {
RwLockReadGuardFuture {
mutex: &self,
is_realized: false,
}
}
#[inline]
pub fn read_owned(self: &Arc<Self>) -> RwLockReadOwnedGuardFuture<T> {
RwLockReadOwnedGuardFuture {
mutex: self.clone(),
is_realized: false,
}
}
#[inline]
fn unlock_reader(&self) {
if self.readers.fetch_sub(1, Ordering::Release) == 1 {
self.inner.unlock()
} else {
self.inner.try_wake(null_mut())
}
}
#[inline]
fn add_reader(&self) {
self.readers.fetch_add(1, Ordering::Release);
}
#[inline]
fn try_acquire_reader(&self) -> bool {
self.readers.load(Ordering::Acquire) > 0 || self.inner.try_acquire()
}
}
#[derive(Debug)]
pub struct RwLockWriteGuard<'a, T: ?Sized> {
mutex: &'a RwLock<T>,
}
#[derive(Debug)]
pub struct RwLockWriteGuardFuture<'a, T: ?Sized> {
mutex: &'a RwLock<T>,
is_realized: bool,
}
#[derive(Debug)]
pub struct RwLockWriteOwnedGuard<T: ?Sized> {
mutex: Arc<RwLock<T>>,
}
#[derive(Debug)]
pub struct RwLockWriteOwnedGuardFuture<T: ?Sized> {
mutex: Arc<RwLock<T>>,
is_realized: bool,
}
#[derive(Debug)]
pub struct RwLockReadGuard<'a, T: ?Sized> {
mutex: &'a RwLock<T>,
}
#[derive(Debug)]
pub struct RwLockReadGuardFuture<'a, T: ?Sized> {
mutex: &'a RwLock<T>,
is_realized: bool,
}
#[derive(Debug)]
pub struct RwLockReadOwnedGuard<T: ?Sized> {
mutex: Arc<RwLock<T>>,
}
#[derive(Debug)]
pub struct RwLockReadOwnedGuardFuture<T: ?Sized> {
mutex: Arc<RwLock<T>>,
is_realized: bool,
}
impl<'a, T: ?Sized> Future for RwLockWriteGuardFuture<'a, T> {
type Output = RwLockWriteGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.mutex.inner.try_acquire() {
self.is_realized = true;
Poll::Ready(RwLockWriteGuard { mutex: self.mutex })
} else {
self.mutex.inner.store_waker(cx.waker());
Poll::Pending
}
}
}
impl<T: ?Sized> Future for RwLockWriteOwnedGuardFuture<T> {
type Output = RwLockWriteOwnedGuard<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.mutex.inner.try_acquire() {
self.is_realized = true;
Poll::Ready(RwLockWriteOwnedGuard {
mutex: self.mutex.clone(),
})
} else {
self.mutex.inner.store_waker(cx.waker());
Poll::Pending
}
}
}
impl<'a, T: ?Sized> Future for RwLockReadGuardFuture<'a, T> {
type Output = RwLockReadGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.mutex.try_acquire_reader() {
self.is_realized = true;
self.mutex.add_reader();
Poll::Ready(RwLockReadGuard { mutex: self.mutex })
} else {
self.mutex.inner.store_waker(cx.waker());
Poll::Pending
}
}
}
impl<T: ?Sized> Future for RwLockReadOwnedGuardFuture<T> {
type Output = RwLockReadOwnedGuard<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.mutex.try_acquire_reader() {
self.is_realized = true;
self.mutex.add_reader();
Poll::Ready(RwLockReadOwnedGuard {
mutex: self.mutex.clone(),
})
} else {
self.mutex.inner.store_waker(cx.waker());
Poll::Pending
}
}
}
crate::impl_send_sync_rwlock!(
RwLock,
RwLockReadGuard,
RwLockReadOwnedGuard,
RwLockWriteGuard,
RwLockWriteOwnedGuard
);
crate::impl_deref_mut!(RwLockWriteGuard, 'a);
crate::impl_deref_mut!(RwLockWriteOwnedGuard);
crate::impl_deref!(RwLockReadGuard, 'a);
crate::impl_deref!(RwLockReadOwnedGuard);
crate::impl_drop_guard!(RwLockWriteGuard, 'a, unlock);
crate::impl_drop_guard!(RwLockWriteOwnedGuard, unlock);
crate::impl_drop_guard_self!(RwLockReadGuard, 'a, unlock_reader);
crate::impl_drop_guard_self!(RwLockReadOwnedGuard, unlock_reader);
crate::impl_drop_guard_future!(RwLockWriteGuardFuture, 'a, unlock);
crate::impl_drop_guard_future!(RwLockWriteOwnedGuardFuture, unlock);
crate::impl_drop_guard_future!(RwLockReadGuardFuture, 'a, unlock);
crate::impl_drop_guard_future!(RwLockReadOwnedGuardFuture, unlock);
#[cfg(test)]
mod tests {
use crate::rwlock::{RwLock, RwLockReadGuard, RwLockWriteGuard, RwLockWriteOwnedGuard};
use futures::executor::block_on;
use futures::{FutureExt, StreamExt, TryStreamExt};
use std::ops::AddAssign;
use std::sync::Arc;
use tokio::time::{sleep, Duration};
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_mutex() {
let c = RwLock::new(0);
futures::stream::iter(0..10000)
.for_each_concurrent(None, |_| async {
let mut co: RwLockWriteGuard<i32> = c.write().await;
*co += 1;
})
.await;
let co = c.write().await;
assert_eq!(*co, 10000)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_mutex_delay() {
let expected_result = 100;
let c = RwLock::new(0);
futures::stream::iter(0..expected_result)
.then(|i| c.write().map(move |co| (i, co)))
.for_each_concurrent(None, |(i, mut co)| async move {
sleep(Duration::from_millis(expected_result - i)).await;
*co += 1;
})
.await;
let co = c.write().await;
assert_eq!(*co, expected_result)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_owned_mutex() {
let c = Arc::new(RwLock::new(0));
futures::stream::iter(0..10000)
.for_each_concurrent(None, |_| async {
let mut co: RwLockWriteOwnedGuard<i32> = c.write_owned().await;
*co += 1;
})
.await;
let co = c.write_owned().await;
assert_eq!(*co, 10000)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_container() {
let c = RwLock::new(String::from("lol"));
let mut co: RwLockWriteGuard<String> = c.write().await;
co.add_assign("lol");
assert_eq!(*co, "lollol");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_timeout() {
let c = RwLock::new(String::from("lol"));
let co: RwLockWriteGuard<String> = c.write().await;
futures::stream::iter(0..10000i32)
.then(|_| tokio::time::timeout(Duration::from_nanos(1), c.write()))
.try_for_each_concurrent(None, |_c| futures::future::ok(()))
.await
.expect_err("timout must be");
drop(co);
let mut co: RwLockWriteGuard<String> = c.write().await;
co.add_assign("lol");
assert_eq!(*co, "lollol");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_concurrent_reading() {
let c = RwLock::new(String::from("lol"));
let co: RwLockReadGuard<String> = c.read().await;
futures::stream::iter(0..10000i32)
.then(|_| c.read())
.inspect(|c| assert_eq!(*co, **c))
.for_each_concurrent(None, |_c| futures::future::ready(()))
.await;
assert!(matches!(
tokio::time::timeout(Duration::from_millis(1), c.write()).await,
Err(_)
));
let co2: RwLockReadGuard<String> = c.read().await;
assert_eq!(*co, *co2);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 12)]
async fn test_concurrent_reading_writing() {
let c = RwLock::new(String::from("lol"));
let co: RwLockReadGuard<String> = c.read().await;
let co2: RwLockReadGuard<String> = c.read().await;
assert_eq!(*co, *co2);
drop(co);
drop(co2);
let mut co: RwLockWriteGuard<String> = c.write().await;
assert!(matches!(
tokio::time::timeout(Duration::from_millis(1), c.read()).await,
Err(_)
));
*co += "lol";
drop(co);
let co: RwLockReadGuard<String> = c.read().await;
let co2: RwLockReadGuard<String> = c.read().await;
assert_eq!(*co, "lollol");
assert_eq!(*co, *co2);
}
#[test]
fn multithreading_test() {
let num = 100;
let mutex = Arc::new(RwLock::new(0));
let ths: Vec<_> = (0..num)
.map(|i| {
let mutex = mutex.clone();
std::thread::spawn(move || {
block_on(async {
if i % 2 == 0 {
let mut lock = mutex.write().await;
*lock += 1;
drop(lock)
} else {
let lock1 = mutex.read().await;
let lock2 = mutex.read().await;
assert_eq!(*lock1, *lock2);
drop(lock1);
drop(lock2);
}
})
})
})
.collect();
for thread in ths {
thread.join().unwrap();
}
block_on(async {
let lock = mutex.read().await;
assert_eq!(num / 2, *lock)
})
}
}