use std::cell::UnsafeCell; use std::fmt; use std::ops::{Deref, DerefMut}; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::task::{Context, Poll}; use crate::signal::Signal; use crossbeam_utils::Backoff; use futures::io::{self, AsyncRead, AsyncWrite}; /// A lock that implements async I/O traits. /// /// This is a blocking mutex that adds the following impls: /// /// - `impl AsyncRead for Lock where &T: AsyncRead + Unpin {}` /// - `impl AsyncRead for &Lock where &T: AsyncRead + Unpin {}` /// - `impl AsyncWrite for Lock where &T: AsyncWrite + Unpin {}` /// - `impl AsyncWrite for &Lock where &T: AsyncWrite + Unpin {}` /// /// This lock is ensures fairness by handling lock operations in the first-in first-out order. /// /// While primarily designed for wrapping async I/O objects, this lock can also be used as a /// regular blocking mutex. It's not quite as efficient as [`parking_lot::Mutex`], but it's still /// an improvement over [`std::sync::Mutex`]. /// /// [`parking_lot::Mutex`]: https://docs.rs/parking_lot/0.10.0/parking_lot/type.Mutex.html /// [`std::sync::Mutex`]: https://doc.rust-lang.org/std/sync/struct.Mutex.html /// /// # Examples /// /// ``` /// use futures::io; /// use futures::prelude::*; /// use piper::Lock; /// /// // Reads data from a stream and echoes it back. /// async fn echo(stream: impl AsyncRead + AsyncWrite + Unpin) -> io::Result { /// let stream = Lock::new(stream); /// io::copy(&stream, &mut &stream).await /// } /// ``` pub struct Lock { /// Set to `true` when the lock is acquired by a `LockGuard`. locked: AtomicBool, /// Lock operations waiting for the lock to get unlocked. lock_ops: Signal, /// The value inside the lock. data: UnsafeCell, } unsafe impl Send for Lock {} unsafe impl Sync for Lock {} impl Lock { /// Creates a new lock. /// /// # Examples /// /// ``` /// use piper::Lock; /// /// let lock = Lock::new(10); /// ``` pub fn new(data: T) -> Lock { Lock { locked: AtomicBool::new(false), lock_ops: Signal::new(), data: UnsafeCell::new(data), } } /// Acquires the lock, blocking the current thread until it is able to do so. /// /// Returns a guard that releases the lock when dropped. /// /// # Examples /// /// ``` /// use piper::Lock; /// /// let lock = Lock::new(10); /// let guard = lock.lock(); /// assert_eq!(*guard, 10); /// ``` pub fn lock(&self) -> LockGuard<'_, T> { loop { // Try locking the lock. let backoff = Backoff::new(); loop { if let Some(guard) = self.try_lock() { return guard; } if backoff.is_completed() { break; } backoff.snooze(); } // Start watching for notifications and try locking again. let l = self.lock_ops.listen(); if let Some(guard) = self.try_lock() { return guard; } l.wait(); } } /// Attempts to acquire the lock. /// /// If the 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 /// /// ``` /// use piper::Lock; /// /// let lock = Lock::new(10); /// if let Ok(guard) = lock.try_lock() { /// assert_eq!(*guard, 10); /// } /// # ; /// ``` #[inline] pub fn try_lock(&self) -> Option> { if !self.locked.compare_and_swap(false, true, Ordering::Acquire) { Some(LockGuard(self)) } else { None } } /// Consumes the lock, returning the underlying data. /// /// # Examples /// /// ``` /// use piper::Lock; /// /// let lock = Lock::new(10); /// assert_eq!(lock.into_inner(), 10); /// ``` pub fn into_inner(self) -> T { self.data.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 the lock is not already acquired. /// /// # Examples /// /// ``` /// use piper::Lock; /// /// let mut lock = Lock::new(0); /// *lock.get_mut() = 10; /// assert_eq!(*lock.lock(), 10); /// ``` pub fn get_mut(&mut self) -> &mut T { unsafe { &mut *self.data.get() } } } impl fmt::Debug for Lock { 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("") } } match self.try_lock() { None => f.debug_struct("Lock").field("data", &Locked).finish(), Some(guard) => f.debug_struct("Lock").field("data", &&*guard).finish(), } } } impl From for Lock { fn from(val: T) -> Lock { Lock::new(val) } } impl Default for Lock { fn default() -> Lock { Lock::new(Default::default()) } } impl AsyncRead for Lock { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll> { Pin::new(&mut *self.lock()).poll_read(cx, buf) } } impl AsyncRead for &Lock { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll> { Pin::new(&mut *self.lock()).poll_read(cx, buf) } } impl AsyncWrite for Lock { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut *self.lock()).poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut *self.lock()).poll_flush(cx) } fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut *self.lock()).poll_close(cx) } } impl AsyncWrite for &Lock { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut *self.lock()).poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut *self.lock()).poll_flush(cx) } fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut *self.lock()).poll_close(cx) } } /// A guard that releases the lock when dropped. pub struct LockGuard<'a, T>(&'a Lock); unsafe impl Send for LockGuard<'_, T> {} unsafe impl Sync for LockGuard<'_, T> {} impl Drop for LockGuard<'_, T> { fn drop(&mut self) { self.0.locked.store(false, Ordering::Release); self.0.lock_ops.notify_one(); } } impl fmt::Debug for LockGuard<'_, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&**self, f) } } impl fmt::Display for LockGuard<'_, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { (**self).fmt(f) } } impl Deref for LockGuard<'_, T> { type Target = T; fn deref(&self) -> &T { unsafe { &*self.0.data.get() } } } impl DerefMut for LockGuard<'_, T> { fn deref_mut(&mut self) -> &mut T { unsafe { &mut *self.0.data.get() } } }