use std::{
cell::UnsafeCell,
fmt, mem,
num::NonZeroUsize,
ops::{Deref, DerefMut},
sync::{self, Arc, PoisonError},
};
use super::WaitStart;
use crate::Event;
pub struct RwLock<T>
where
T: ?Sized,
{
state: sync::Mutex<RwState>,
readers_may_enter: Event,
writer_may_enter: Event,
value: UnsafeCell<T>,
}
unsafe impl<T> Sync for RwLock<T> where T: ?Sized + Send + Sync {}
impl<T> RwLock<T> {
pub const fn new(value: T) -> Self {
Self {
state: sync::Mutex::new(RwState {
owner: Owner::Unlocked,
writers_waiting: 0,
writers_starved: 0,
readers_starved: 0,
}),
readers_may_enter: Event::new(),
writer_may_enter: Event::new(),
value: UnsafeCell::new(value),
}
}
pub fn into_inner(self) -> T {
self.value.into_inner()
}
}
impl<T> RwLock<T>
where
T: ?Sized,
{
pub async fn read(&self) -> RwLockReadGuard<'_, T> {
self.acquire_read(|| RwLockReadGuard(self)).await
}
pub async fn write(&self) -> RwLockWriteGuard<'_, T> {
self.acquire_write(|| RwLockWriteGuard(self)).await
}
pub fn try_read(&self) -> Option<RwLockReadGuard<'_, T>> {
self.try_acquire_read().then(|| RwLockReadGuard(self))
}
pub fn try_write(&self) -> Option<RwLockWriteGuard<'_, T>> {
self.try_acquire_write().then(|| RwLockWriteGuard(self))
}
pub async fn read_arc(self: &Arc<Self>) -> RwLockReadGuardArc<T> {
self.acquire_read(|| RwLockReadGuardArc(self.clone())).await
}
pub fn try_read_arc(self: &Arc<Self>) -> Option<RwLockReadGuardArc<T>> {
self.try_acquire_read()
.then(|| RwLockReadGuardArc(self.clone()))
}
pub async fn write_arc(self: &Arc<Self>) -> RwLockWriteGuardArc<T> {
self.acquire_write(|| RwLockWriteGuardArc(self.clone()))
.await
}
pub fn try_write_arc(self: &Arc<Self>) -> Option<RwLockWriteGuardArc<T>> {
self.try_acquire_write()
.then(|| RwLockWriteGuardArc(self.clone()))
}
pub fn get_mut(&mut self) -> &mut T {
self.value.get_mut()
}
async fn acquire_read<F, G>(&self, guard: F) -> G
where
F: FnOnce() -> G,
{
let mut reader = WaitingReader {
rwlock: self,
since: None,
starved: false,
};
let listener = loop {
if reader.try_read() {
break None;
}
reader.lost();
let listener = self.readers_may_enter.listen_unfenced();
if reader.try_read() {
break Some(listener);
}
reader.since.get_or_insert_with(WaitStart::now);
listener.await;
};
let guard = guard();
drop(listener);
drop(reader);
guard
}
async fn acquire_write<F, G>(&self, guard: F) -> G
where
F: FnOnce() -> G,
{
let mut writer = WaitingWriter::register(self);
let listener = loop {
if writer.try_write() {
break None;
}
writer.lost();
let listener = self.writer_may_enter.listen_unfenced();
if writer.try_write() {
break Some(listener);
}
writer.since.get_or_insert_with(WaitStart::now);
listener.await;
};
let guard = guard();
drop(listener);
drop(writer);
guard
}
fn try_acquire_read(&self) -> bool {
let mut state = lock(&self.state);
state.writers_waiting == 0 && state.admit_reader()
}
fn try_acquire_write(&self) -> bool {
let mut state = lock(&self.state);
state.writers_starved == 0 && state.admit_writer()
}
fn read_unlock(&self) {
let mut state = lock(&self.state);
let Owner::Reading(readers) = state.owner else {
unreachable!("a read guard exists only while the owner is reading");
};
if let Some(readers) = NonZeroUsize::new(readers.get() - 1) {
state.owner = Owner::Reading(readers);
return;
}
state.owner = Owner::Unlocked;
drop(state);
self.writer_may_enter.notify_unfenced(1);
}
fn write_unlock(&self) {
lock(&self.state).owner = Owner::Unlocked;
self.readers_may_enter.notify_unfenced(usize::MAX);
self.writer_may_enter.notify_unfenced(1);
}
}
impl<T> Default for RwLock<T>
where
T: Default,
{
fn default() -> Self {
Self::new(T::default())
}
}
impl<T> From<T> for RwLock<T> {
fn from(value: T) -> Self {
Self::new(value)
}
}
impl<T> fmt::Debug for RwLock<T>
where
T: ?Sized + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut s = f.debug_struct("RwLock");
match self.try_read() {
Some(guard) => s.field("value", &&*guard),
None => s.field("value", &format_args!("<locked>")),
};
s.finish()
}
}
#[must_use = "if unused the RwLock will immediately unlock"]
pub struct RwLockReadGuard<'a, T>(&'a RwLock<T>)
where
T: ?Sized;
unsafe impl<T> Send for RwLockReadGuard<'_, T> where T: ?Sized + Sync {}
unsafe impl<T> Sync for RwLockReadGuard<'_, T> where T: ?Sized + Sync {}
impl<T> Deref for RwLockReadGuard<'_, T>
where
T: ?Sized,
{
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
impl<T> fmt::Debug for RwLockReadGuard<'_, T>
where
T: ?Sized + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T> Drop for RwLockReadGuard<'_, T>
where
T: ?Sized,
{
fn drop(&mut self) {
self.0.read_unlock();
}
}
#[must_use = "if unused the RwLock will immediately unlock"]
pub struct RwLockWriteGuard<'a, T>(&'a RwLock<T>)
where
T: ?Sized;
unsafe impl<T> Send for RwLockWriteGuard<'_, T> where T: ?Sized + Send {}
unsafe impl<T> Sync for RwLockWriteGuard<'_, T> where T: ?Sized + Sync {}
impl<T> Deref for RwLockWriteGuard<'_, T>
where
T: ?Sized,
{
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
impl<T> DerefMut for RwLockWriteGuard<'_, T>
where
T: ?Sized,
{
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.0.value.get() }
}
}
impl<T> fmt::Debug for RwLockWriteGuard<'_, T>
where
T: ?Sized + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T> Drop for RwLockWriteGuard<'_, T>
where
T: ?Sized,
{
fn drop(&mut self) {
self.0.write_unlock();
}
}
#[must_use = "if unused the RwLock will immediately unlock"]
pub struct RwLockReadGuardArc<T>(Arc<RwLock<T>>)
where
T: ?Sized;
impl<T> Deref for RwLockReadGuardArc<T>
where
T: ?Sized,
{
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
impl<T> fmt::Debug for RwLockReadGuardArc<T>
where
T: ?Sized + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T> Drop for RwLockReadGuardArc<T>
where
T: ?Sized,
{
fn drop(&mut self) {
self.0.read_unlock();
}
}
#[must_use = "if unused the RwLock will immediately unlock"]
pub struct RwLockWriteGuardArc<T>(Arc<RwLock<T>>)
where
T: ?Sized;
impl<T> Deref for RwLockWriteGuardArc<T>
where
T: ?Sized,
{
type Target = T;
fn deref(&self) -> &T {
unsafe { &*self.0.value.get() }
}
}
impl<T> DerefMut for RwLockWriteGuardArc<T>
where
T: ?Sized,
{
fn deref_mut(&mut self) -> &mut T {
unsafe { &mut *self.0.value.get() }
}
}
impl<T> fmt::Debug for RwLockWriteGuardArc<T>
where
T: ?Sized + fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T> Drop for RwLockWriteGuardArc<T>
where
T: ?Sized,
{
fn drop(&mut self) {
self.0.write_unlock();
}
}
struct RwState {
owner: Owner,
writers_waiting: usize,
writers_starved: usize,
readers_starved: usize,
}
impl RwState {
fn admit_reader(&mut self) -> bool {
self.owner = match self.owner {
Owner::Writing => return false,
Owner::Unlocked => Owner::Reading(NonZeroUsize::MIN),
Owner::Reading(readers) => Owner::Reading(
readers
.checked_add(1)
.expect("more readers than a usize counts"),
),
};
true
}
fn admit_writer(&mut self) -> bool {
if !matches!(self.owner, Owner::Unlocked) || self.readers_starved > 0 {
return false;
}
self.owner = Owner::Writing;
true
}
}
enum Owner {
Unlocked,
Reading(NonZeroUsize),
Writing,
}
struct WaitingReader<'a, T>
where
T: ?Sized,
{
rwlock: &'a RwLock<T>,
since: Option<WaitStart>,
starved: bool,
}
impl<T> WaitingReader<'_, T>
where
T: ?Sized,
{
fn try_read(&mut self) -> bool {
let mut state = lock(&self.rwlock.state);
if (!self.starved && state.writers_waiting > 0) || !state.admit_reader() {
return false;
}
if mem::take(&mut self.starved) {
state.readers_starved -= 1;
}
true
}
fn lost(&mut self) {
let Some(since) = self.since else {
return;
};
if self.starved || !since.waited_long() {
return;
}
self.starved = true;
lock(&self.rwlock.state).readers_starved += 1;
}
}
impl<T> Drop for WaitingReader<'_, T>
where
T: ?Sized,
{
fn drop(&mut self) {
if !self.starved {
return;
}
let mut state = lock(&self.rwlock.state);
state.readers_starved -= 1;
if state.readers_starved > 0 || !matches!(state.owner, Owner::Unlocked) {
return;
}
drop(state);
self.rwlock.writer_may_enter.notify_unfenced(1);
}
}
struct WaitingWriter<'a, T>
where
T: ?Sized,
{
rwlock: &'a RwLock<T>,
counted: bool,
since: Option<WaitStart>,
starved: bool,
}
impl<'a, T> WaitingWriter<'a, T>
where
T: ?Sized,
{
fn register(rwlock: &'a RwLock<T>) -> Self {
lock(&rwlock.state).writers_waiting += 1;
Self {
rwlock,
counted: true,
since: None,
starved: false,
}
}
fn try_write(&mut self) -> bool {
let mut state = lock(&self.rwlock.state);
if (self.since.is_none() && state.writers_starved > 0) || !state.admit_writer() {
return false;
}
self.counted = false;
state.writers_waiting -= 1;
if mem::take(&mut self.starved) {
state.writers_starved -= 1;
}
true
}
fn lost(&mut self) {
let Some(since) = self.since else {
return;
};
if self.starved || !since.waited_long() {
return;
}
self.starved = true;
lock(&self.rwlock.state).writers_starved += 1;
}
}
impl<T> Drop for WaitingWriter<'_, T>
where
T: ?Sized,
{
fn drop(&mut self) {
if !self.counted {
return;
}
let mut state = lock(&self.rwlock.state);
state.writers_waiting -= 1;
if self.starved {
state.writers_starved -= 1;
}
let wake_readers = state.writers_waiting == 0 && !matches!(state.owner, Owner::Writing);
let wake_writer = self.starved && matches!(state.owner, Owner::Unlocked);
drop(state);
if wake_readers {
self.rwlock.readers_may_enter.notify_unfenced(usize::MAX);
}
if wake_writer {
self.rwlock.writer_may_enter.notify_unfenced(1);
}
}
}
fn lock(state: &sync::Mutex<RwState>) -> sync::MutexGuard<'_, RwState> {
state.lock().unwrap_or_else(PoisonError::into_inner)
}