use parking_lot::Mutex as RegularMutex;
use parking_lot::RwLock as RegularRwLock;
use parking_lot::RwLockReadGuard as RegularRwLockReadGuard;
use parking_lot::RwLockWriteGuard as RegularRwLockWriteGuard;
use slab::Slab;
use stable_deref_trait::StableDeref;
use std::cell::UnsafeCell;
use std::future::Future;
use std::mem;
use std::ops::{Deref, DerefMut};
use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::time::Duration;
const WAIT_KEY_NONE: usize = std::usize::MAX;
enum Waiter {
Waiting(Waker),
Woken,
}
impl Waiter {
#[inline]
fn register(&mut self, w: &Waker) {
match self {
Waiter::Waiting(waker) if w.will_wake(waker) => {}
_ => *self = Waiter::Waiting(w.clone()),
}
}
#[inline]
fn wake(&mut self) {
match mem::replace(self, Waiter::Woken) {
Waiter::Waiting(waker) => waker.wake(),
Waiter::Woken => {}
}
}
}
pub struct RwLockReadGuard<'a, T> {
_inner_guard: Option<RegularRwLockReadGuard<'a, ()>>,
lock: &'a RwLock<T>,
}
impl<'a, T> Deref for RwLockReadGuard<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &Self::Target {
unsafe { &*self.lock.data.get() }
}
}
impl<'a, T> Drop for RwLockReadGuard<'a, T> {
#[inline]
fn drop(&mut self) {
drop(self._inner_guard.take());
let mut waiters = self.lock.waiters.lock();
if let Some((_i, waiter)) = waiters.iter_mut().next() {
waiter.wake();
}
}
}
pub struct RwLockWriteGuard<'a, T> {
_inner_guard: Option<RegularRwLockWriteGuard<'a, ()>>,
lock: &'a RwLock<T>,
}
impl<'a, T> Deref for RwLockWriteGuard<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &Self::Target {
unsafe { &*self.lock.data.get() }
}
}
impl<'a, T> DerefMut for RwLockWriteGuard<'a, T> {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { &mut *self.lock.data.get() }
}
}
impl<'a, T> Drop for RwLockWriteGuard<'a, T> {
#[inline]
fn drop(&mut self) {
drop(self._inner_guard.take());
let mut waiters = self.lock.waiters.lock();
if let Some((_i, waiter)) = waiters.iter_mut().next() {
waiter.wake();
}
}
}
pub struct RwLock<T> {
lock: RegularRwLock<()>,
waiters: RegularMutex<Slab<Waiter>>,
data: UnsafeCell<T>,
}
impl<T> RwLock<T> {
fn remove_waker(&self, wait_key: usize, wake_another: bool) {
if wait_key != WAIT_KEY_NONE {
let mut waiters = self.waiters.lock();
match waiters.remove(wait_key) {
Waiter::Waiting(_) => {}
Waiter::Woken => {
if wake_another {
if let Some((_i, waiter)) = waiters.iter_mut().next() {
waiter.wake();
}
}
}
}
}
}
pub fn new(data: T) -> Self {
Self {
lock: RegularRwLock::new(()),
waiters: RegularMutex::new(Slab::new()),
data: UnsafeCell::new(data),
}
}
#[inline]
pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
self.lock.try_read().map(|guard| RwLockReadGuard {
_inner_guard: Some(guard),
lock: self,
})
}
#[inline]
pub fn try_read_for(&self, d: Duration) -> Option<RwLockReadGuard<'_, T>> {
self.lock.try_read_for(d).map(|guard| RwLockReadGuard {
_inner_guard: Some(guard),
lock: self,
})
}
#[inline]
pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
self.lock.try_write().map(|guard| RwLockWriteGuard {
_inner_guard: Some(guard),
lock: self,
})
}
#[inline]
pub fn try_write_for(&self, d: Duration) -> Option<RwLockWriteGuard<'_, T>> {
self.lock.try_write_for(d).map(|guard| RwLockWriteGuard {
_inner_guard: Some(guard),
lock: self,
})
}
#[inline]
pub fn read(&self) -> RwLockReadGuard<'_, T> {
RwLockReadGuard {
_inner_guard: Some(self.lock.read()),
lock: self,
}
}
#[inline]
pub fn write(&self) -> RwLockWriteGuard<'_, T> {
RwLockWriteGuard {
_inner_guard: Some(self.lock.write()),
lock: self,
}
}
#[inline]
pub fn async_read(&self) -> RwLockReadFuture<'_, T> {
RwLockReadFuture {
lock: Some(self),
wait_key: WAIT_KEY_NONE,
}
}
#[inline]
pub fn async_write(&self) -> RwLockWriteFuture<'_, T> {
RwLockWriteFuture {
lock: Some(self),
wait_key: WAIT_KEY_NONE,
}
}
}
pub struct RwLockReadFuture<'a, T> {
lock: Option<&'a RwLock<T>>,
wait_key: usize,
}
impl<'a, T> Drop for RwLockReadFuture<'a, T> {
#[inline]
fn drop(&mut self) {
if let Some(lock) = self.lock {
lock.remove_waker(self.wait_key, true);
}
}
}
impl<'a, T> Future for RwLockReadFuture<'a, T> {
type Output = RwLockReadGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let lock = self.lock.expect("polled after completion");
if let Some(guard) = lock.try_read() {
lock.remove_waker(self.wait_key, false);
self.lock = None;
return Poll::Ready(guard);
}
{
let mut waiters = lock.waiters.lock();
if self.wait_key == WAIT_KEY_NONE {
self.wait_key = waiters.insert(Waiter::Waiting(cx.waker().clone()));
} else {
waiters[self.wait_key].register(cx.waker())
}
}
if let Some(guard) = lock.try_read() {
lock.remove_waker(self.wait_key, false);
self.lock = None;
return Poll::Ready(guard);
}
Poll::Pending
}
}
pub struct RwLockWriteFuture<'a, T> {
lock: Option<&'a RwLock<T>>,
wait_key: usize,
}
impl<'a, T> Drop for RwLockWriteFuture<'a, T> {
#[inline]
fn drop(&mut self) {
if let Some(lock) = self.lock {
lock.remove_waker(self.wait_key, true);
}
}
}
impl<'a, T> Future for RwLockWriteFuture<'a, T> {
type Output = RwLockWriteGuard<'a, T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let lock = self.lock.expect("polled after completion");
if let Some(guard) = lock.try_write() {
lock.remove_waker(self.wait_key, false);
self.lock = None;
return Poll::Ready(guard);
}
{
let mut waiters = lock.waiters.lock();
if self.wait_key == WAIT_KEY_NONE {
self.wait_key = waiters.insert(Waiter::Waiting(cx.waker().clone()));
} else {
waiters[self.wait_key].register(cx.waker())
}
}
if let Some(guard) = lock.try_write() {
lock.remove_waker(self.wait_key, false);
self.lock = None;
return Poll::Ready(guard);
}
Poll::Pending
}
}
unsafe impl<T: Send> Send for RwLock<T> {}
unsafe impl<T: Sync> Sync for RwLock<T> {}
unsafe impl<T: Send> Send for RwLockReadFuture<'_, T> {}
unsafe impl<T: Send> Sync for RwLockReadFuture<'_, T> {}
unsafe impl<T: Send> Send for RwLockWriteFuture<'_, T> {}
unsafe impl<T: Send> Sync for RwLockWriteFuture<'_, T> {}
unsafe impl<T: Send> Send for RwLockReadGuard<'_, T> {}
unsafe impl<T: Sync> Sync for RwLockReadGuard<'_, T> {}
unsafe impl<T: Send> Send for RwLockWriteGuard<'_, T> {}
unsafe impl<T: Sync> Sync for RwLockWriteGuard<'_, T> {}
unsafe impl<T> StableDeref for RwLockReadGuard<'_, T> {}
unsafe impl<T> StableDeref for RwLockWriteGuard<'_, T> {}