Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- //use futures_core::future::{FusedFuture, Future};
- use futures::future::{FusedFuture, Future};
- //use futures_core::task::{Context, Poll, Waker};
- use futures::task::{Context, Poll, Waker};
- use slab::Slab;
- use std::cell::UnsafeCell;
- use std::marker::PhantomData;
- use std::ops::{Deref, DerefMut};
- use std::pin::Pin;
- use std::sync::atomic::{AtomicUsize, Ordering};
- use std::sync::Mutex as StdMutex;
- use std::{fmt, mem};
- /// A futures-aware mutex.
- pub struct Mutex<T: ?Sized> {
- state: AtomicUsize,
- waiters: StdMutex<Slab<Waiter>>,
- read_waiters: StdMutex<Slab<Waiter>>,
- value: UnsafeCell<T>,
- }
- impl<T: ?Sized> fmt::Debug for Mutex<T> {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- let s = State(self.state.load(Ordering::SeqCst));
- f.debug_struct("Mutex")
- .field("is_locked", &s.is_locked())
- .field("has_waiters", &s.has_waiters())
- .field("has_read_waiters", &s.has_read_waiters())
- .field("is_corrupted", &s.is_corrupted())
- .field("read_lock_count", &s.read_lock_count())
- .finish()
- }
- }
- impl<T> From<T> for Mutex<T> {
- fn from(t: T) -> Self {
- Self::new(t)
- }
- }
- impl<T: Default> Default for Mutex<T> {
- fn default() -> Mutex<T> {
- Mutex::new(Default::default())
- }
- }
- pub enum Waiter {
- Waiting(Waker),
- Woken,
- }
- impl Waiter {
- fn register(&mut self, waker: &Waker) {
- match self {
- Waiter::Waiting(w) if waker.will_wake(w) => {}
- _ => *self = Waiter::Waiting(waker.clone()),
- }
- }
- fn wake(&mut self) {
- match mem::replace(self, Waiter::Woken) {
- Waiter::Waiting(waker) => waker.wake(),
- Waiter::Woken => {}
- }
- }
- }
- macro_rules! bitflag {
- ($name:ident, $bit:expr) => {
- const fn $name(&self) -> bool {
- self.0 & $bit != 0
- }
- };
- ($name:ident, $flag:ident, $bit:expr) => {
- const $flag: usize = $bit;
- bitflag!($name, Self::$flag);
- }
- }
- struct State(usize);
- impl State {
- // set if it's write locked
- bitflag!(is_locked, IS_LOCKED, 1 << 31);
- // set if there are write waiters
- bitflag!(has_waiters, HAS_WAITERS, 1 << 30);
- // set if there are read waiters
- bitflag!(has_read_waiters, HAS_READ_WAITERS, 1 << 29);
- // the bits which are valid for read lock count (supports 16777216 readers)
- const MAX_READ_LOCKS: usize = (1 << 24) - 1;
- // bits which should always be zero (for future use)
- bitflag!(is_corrupted, (1 << 29) - 1 - Self::MAX_READ_LOCKS);
- const fn read_lock_count(&self) -> usize {
- self.0 & Self::MAX_READ_LOCKS
- }
- }
- impl<T> Mutex<T> {
- /// Creates a new futures-aware mutex.
- pub fn new(t: T) -> Mutex<T> {
- Mutex {
- state: AtomicUsize::new(0),
- waiters: StdMutex::new(Slab::new()),
- read_waiters: StdMutex::new(Slab::new()),
- value: UnsafeCell::new(t),
- }
- }
- /// Consumes this mutex, returning the underlying data.
- ///
- /// # Examples
- ///
- /// ```
- /// use futures::lock::Mutex;
- ///
- /// let mutex = Mutex::new(0);
- /// assert_eq!(mutex.into_inner(), 0);
- /// ```
- pub fn into_inner(self) -> T {
- self.value.into_inner()
- }
- }
- pub struct Lock<'a, RW: ReadWrite, T: ?Sized> {
- pub mutex: &'a Mutex<T>,
- _rw: PhantomData<RW>,
- }
- impl<'a, RW: ReadWrite, T: ?Sized> Lock<'a, RW, T> {
- /// Attempt to acquire the lock immediately.
- ///
- /// If the lock is currently held, this will return `None`.
- pub fn try_lock(&self) -> Option<MutexGuard<'a, RW, T>> {
- if RW::IS_WRITE {
- let old_state = State(
- self.mutex
- .state
- .fetch_or(State::IS_LOCKED, Ordering::Acquire),
- );
- if old_state.is_locked() {
- return None;
- }
- } else {
- let old_state = State(self.mutex.state.fetch_add(1, Ordering::Acquire));
- if old_state.read_lock_count() == State::MAX_READ_LOCKS {
- panic!("Overflowed the number of read locks");
- }
- if old_state.is_locked() {
- self.mutex.state.fetch_sub(1, Ordering::Relaxed);
- return None;
- }
- }
- Some(MutexGuard {
- mutex: self.mutex,
- _rw: PhantomData,
- })
- }
- /// Acquire the lock asynchronously.
- ///
- /// This method returns a future that will resolve once the lock has been
- /// successfully acquired.
- pub fn lock(&self) -> MutexLockFuture<'a, RW, T> {
- MutexLockFuture {
- mutex: Some(self.mutex),
- wait_key: WAIT_KEY_NONE,
- _rw: PhantomData,
- }
- }
- }
- impl<'a, T: ?Sized> Mutex<T> {
- pub const fn read(&'a self) -> Lock<'a, Read, T> {
- Lock {
- mutex: self,
- _rw: PhantomData,
- }
- }
- pub const fn write(&'a self) -> Lock<'a, Write, T> {
- Lock {
- mutex: self,
- _rw: PhantomData,
- }
- }
- /// Returns a mutable reference to the underlying data.
- ///
- /// Since this call borrows the `Mutex` mutably, no actual locking needs to
- /// take place -- the mutable borrow statically guarantees no locks exist.
- ///
- /// # Examples
- ///
- /// ```
- /// # futures::executor::block_on(async {
- /// use futures::lock::Mutex;
- ///
- /// let mut mutex = Mutex::new(0);
- /// *mutex.get_mut() = 10;
- /// assert_eq!(*mutex.lock().await, 10);
- /// # });
- /// ```
- pub fn get_mut(&mut self) -> &mut T {
- // We know statically that there are no other references to `self`, so
- // there's no need to lock the inner mutex.
- unsafe { &mut *self.value.get() }
- }
- fn remove_waker<RW:ReadWrite>(&self, wait_key: usize, wake_another: bool) {
- if wait_key == WAIT_KEY_NONE {
- return;
- }
- let mut waiters = RW::waiters(self).lock().unwrap();
- match waiters.remove(wait_key) {
- Waiter::Waiting(_) => {}
- Waiter::Woken => {
- // We were awoken, but then dropped before we could
- // wake up to acquire the lock. Wake up another
- // waiter.
- if RW::IS_WRITE && wake_another {
- if let Some((_i, waiter)) = waiters.iter_mut().next() {
- waiter.wake();
- }
- }
- }
- }
- if waiters.is_empty() {
- self.state.fetch_and(!State::HAS_WAITERS, Ordering::Relaxed); // released by mutex unlock
- // If we're in read mode, wake_another is true, AND we're the last one
- // Lets wake a write lock
- if !RW::IS_WRITE && wake_another {
- let mut write_waiters = self.waiters.lock().unwrap();
- if let Some((_i, waiter)) = write_waiters.iter_mut().next() {
- waiter.wake();
- }
- }
- }
- }
- }
- pub trait ReadWrite: Sized + Unpin {
- const IS_WRITE: bool;
- const STR: &'static str;
- const HAS_WAITERS_FLAG: usize;
- fn lock<'a, T: ?Sized>(m: &'a Mutex<T>) -> Lock<'a, Self, T>;
- fn waiters<'a, T: ?Sized>(m: &Mutex<T>) -> &StdMutex<Slab<Waiter>>;
- }
- pub struct Write;
- pub struct Read;
- impl ReadWrite for Write {
- const IS_WRITE: bool = true;
- const STR: &'static str = "Write";
- const HAS_WAITERS_FLAG: usize = State::HAS_WAITERS;
- fn lock<'a, T: ?Sized>(m: &'a Mutex<T>) -> Lock<'a, Self, T> {
- m.write()
- }
- fn waiters<'a, T: ?Sized>(m: &Mutex<T>) -> &StdMutex<Slab<Waiter>> {
- &m.waiters
- }
- }
- impl ReadWrite for Read {
- const IS_WRITE: bool = true;
- const STR: &'static str = "Read";
- const HAS_WAITERS_FLAG: usize = State::HAS_READ_WAITERS;
- fn lock<'a, T: ?Sized>(m: &'a Mutex<T>) -> Lock<'a, Self, T> {
- m.read()
- }
- fn waiters<'a, T: ?Sized>(m: &Mutex<T>) -> &StdMutex<Slab<Waiter>> {
- &m.read_waiters
- }
- }
- // Sentinel for when no slot in the `Slab` has been dedicated to this object.
- const WAIT_KEY_NONE: usize = usize::max_value();
- /// A future which resolves when the target mutex has been successfully acquired.
- pub struct MutexLockFuture<'a, RW: ReadWrite, T: ?Sized> {
- // `None` indicates that the mutex was successfully acquired.
- mutex: Option<&'a Mutex<T>>,
- wait_key: usize,
- _rw: PhantomData<RW>,
- }
- impl<RW: ReadWrite, T: ?Sized> fmt::Debug for MutexLockFuture<'_, RW, T> {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct(&format!("MutexLockFuture<{}>", RW::STR)[..])
- .field("was_acquired", &self.mutex.is_none())
- .field("mutex", &self.mutex)
- .field(
- "wait_key",
- &(if self.wait_key == WAIT_KEY_NONE {
- None
- } else {
- Some(self.wait_key)
- }),
- )
- .finish()
- }
- }
- impl<'a, RW: ReadWrite, T: ?Sized> FusedFuture for MutexLockFuture<'a, RW, T> {
- fn is_terminated(&self) -> bool {
- self.mutex.is_none()
- }
- }
- impl<'a, RW: ReadWrite, T: ?Sized> Future for MutexLockFuture<'a, RW, T> {
- type Output = MutexGuard<'a, RW, T>;
- fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
- let mutex = self.mutex.expect("polled MutexLockFuture after completion");
- let lock = RW::lock(mutex);
- if let Some(locked) = lock.try_lock() {
- mutex.remove_waker::<RW>(self.wait_key, false);
- self.mutex = None;
- return Poll::Ready(locked);
- }
- {
- let mut waiters = RW::waiters(mutex).lock().unwrap();
- if self.wait_key == WAIT_KEY_NONE {
- self.wait_key = waiters.insert(Waiter::Waiting(cx.waker().clone()));
- if waiters.len() == 1 {
- mutex
- .state
- .fetch_or(RW::HAS_WAITERS_FLAG, Ordering::Relaxed); // released by mutex unlock
- }
- } else {
- waiters[self.wait_key].register(cx.waker());
- }
- }
- // Ensure that we haven't raced `MutexGuard::drop`'s unlock path by
- // attempting to acquire the lock again.
- if let Some(locked) = lock.try_lock() {
- mutex.remove_waker::<RW>(self.wait_key, false);
- self.mutex = None;
- return Poll::Ready(locked);
- }
- Poll::Pending
- }
- }
- impl<RW: ReadWrite, T: ?Sized> Drop for MutexLockFuture<'_, RW, T> {
- fn drop(&mut self) {
- if let Some(mutex) = self.mutex {
- // This future was dropped before it acquired the mutex.
- //
- // Remove ourselves from the map, waking up another waiter if we
- // had been awoken to acquire the lock.
- mutex.remove_waker::<RW>(self.wait_key, true);
- }
- }
- }
- /// An RAII guard returned by the `lock` and `try_lock` methods.
- /// When this structure is dropped (falls out of scope), the lock will be
- /// unlocked.
- pub struct MutexGuard<'a, RW: ReadWrite, T: ?Sized> {
- mutex: &'a Mutex<T>,
- _rw: PhantomData<RW>,
- }
- impl<RW: ReadWrite, T: ?Sized + fmt::Debug> fmt::Debug for MutexGuard<'_, RW, T> {
- fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
- f.debug_struct(&format!("MutexGuard<{}>", RW::STR)[..])
- .field("value", &*self)
- .field("mutex", &self.mutex)
- .finish()
- }
- }
- impl<RW: ReadWrite, T: ?Sized> Drop for MutexGuard<'_, RW, T> {
- fn drop(&mut self) {
- if RW::IS_WRITE {
- let old_state = State(
- self.mutex
- .state
- .fetch_and(!State::IS_LOCKED, Ordering::AcqRel),
- );
- if old_state.has_waiters() {
- // Wake one write waiter
- let mut waiters = self.mutex.waiters.lock().unwrap();
- if let Some((_i, waiter)) = waiters.iter_mut().next() {
- waiter.wake();
- }
- } else if old_state.has_read_waiters() {
- // Wake all read waiters
- let mut waiters = self.mutex.waiters.lock().unwrap();
- for w in waiters.iter_mut() {
- w.1.wake()
- }
- }
- } else {
- let old_state = State(self.mutex.state.fetch_sub(1, Ordering::AcqRel));
- if old_state.read_lock_count() > 1 {
- // We're not the last read lock holder
- } else if old_state.has_waiters() {
- // Wake one write waiter
- let mut waiters = self.mutex.waiters.lock().unwrap();
- if let Some((_i, waiter)) = waiters.iter_mut().next() {
- waiter.wake();
- }
- } else if old_state.has_read_waiters() {
- panic!("Race condition, read waiters were queued while the read lock was held")
- }
- }
- }
- }
- impl<RW: ReadWrite, T: ?Sized> Deref for MutexGuard<'_, RW, T> {
- type Target = T;
- fn deref(&self) -> &T {
- unsafe { &*self.mutex.value.get() }
- }
- }
- impl<T: ?Sized> DerefMut for MutexGuard<'_, Write, T> {
- fn deref_mut(&mut self) -> &mut T {
- unsafe { &mut *self.mutex.value.get() }
- }
- }
- // Mutexes can be moved freely between threads and acquired on any thread so long
- // as the inner value can be safely sent between threads.
- unsafe impl<T: ?Sized + Send> Send for Mutex<T> {}
- unsafe impl<T: ?Sized + Send> Sync for Mutex<T> {}
- // It's safe to switch which thread the acquire is being attempted on so long as
- // `T` can be accessed on that thread.
- unsafe impl<RW: ReadWrite, T: ?Sized + Send> Send for MutexLockFuture<'_, RW, T> {}
- // doesn't have any interesting `&self` methods (only Debug)
- unsafe impl<RW: ReadWrite, T: ?Sized> Sync for MutexLockFuture<'_, RW, T> {}
- // Safe to send since we don't track any thread-specific details-- the inner
- // lock is essentially spinlock-equivalent (attempt to flip an atomic bool)
- unsafe impl<RW: ReadWrite, T: ?Sized + Send> Send for MutexGuard<'_, RW, T> {}
- unsafe impl<RW: ReadWrite, T: ?Sized + Sync> Sync for MutexGuard<'_, RW, T> {}
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement