diff --git a/futures-util/src/lock/mod.rs b/futures-util/src/lock/mod.rs index 8ca0ff625..4b7d9f06a 100644 --- a/futures-util/src/lock/mod.rs +++ b/futures-util/src/lock/mod.rs @@ -14,7 +14,8 @@ pub use self::bilock::{BiLock, BiLockAcquire, BiLockGuard, ReuniteError}; #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))] #[cfg(feature = "std")] pub use self::mutex::{ - MappedMutexGuard, Mutex, MutexGuard, MutexLockFuture, OwnedMutexGuard, OwnedMutexLockFuture, + MappedMutexGuard, Mutex, MutexGuard, MutexLockFuture, OwnedMappedMutexGuard, OwnedMutexGuard, + OwnedMutexLockFuture, }; #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))] diff --git a/futures-util/src/lock/mutex.rs b/futures-util/src/lock/mutex.rs index 79272f4be..8ea4526bc 100644 --- a/futures-util/src/lock/mutex.rs +++ b/futures-util/src/lock/mutex.rs @@ -1,3 +1,4 @@ +use core::ptr; use std::cell::UnsafeCell; use std::marker::PhantomData; use std::ops::{Deref, DerefMut}; @@ -283,6 +284,51 @@ pub struct OwnedMutexGuard { mutex: Arc>, } +/// A owned handle to a held `Mutex`. +struct OwnedMutexGuardExtendedLock { + mutex: Arc>, +} + +impl OwnedMutexGuard { + /// `skip_drop` prevents the `OwnedMutexGuard` from being automatically dropped, allowing manual control over the drop process. + /// This method returns an `OwnedMutexGuardExtendedLock` which contains the original `Arc>`. + fn skip_drop(self) -> OwnedMutexGuardExtendedLock { + // Prevents automatic drop by wrapping in `ManuallyDrop` + let man = mem::ManuallyDrop::new(self); + // Unsafely reads the `Arc>` from the `ManuallyDrop` wrapper + OwnedMutexGuardExtendedLock { mutex: unsafe { ptr::read(&man.mutex) } } + } + + /// Returns a locked view over a portion of the locked data. + #[inline] + pub fn map(mut this: Self, f: F) -> OwnedMappedMutexGuard + where + U: ?Sized, + F: FnOnce(&mut T) -> &mut U, + { + let value = f(&mut *this) as *mut U; + let mutex = this.skip_drop(); + OwnedMappedMutexGuard { value, mutex: mutex.mutex } + } + + /// Returns a locked view over a portion of the locked data. + /// If the closure returns `None`, the original `OwnedMappedMutexGuard` is returned. + #[inline] + pub fn try_map(mut this: Self, f: F) -> Result, Self> + where + U: ?Sized, + F: FnOnce(&mut T) -> Option<&mut U>, + { + let value = match f(&mut *this) { + Some(data) => data as *mut U, + None => return Err(this), + }; + let mutex = this.skip_drop(); + + Ok(OwnedMappedMutexGuard { value, mutex: mutex.mutex }) + } +} + impl fmt::Debug for OwnedMutexGuard { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("OwnedMutexGuard") @@ -311,6 +357,60 @@ impl DerefMut for OwnedMutexGuard { } } +/// An RAII guard returned by the `OwnedMutexGuard::map` and `OwnedMappedMutexGuard::map` methods. +/// When this structure is dropped (falls out of scope), the lock will be unlocked. +pub struct OwnedMappedMutexGuard { + value: *mut U, + mutex: Arc>, +} + +/// A owned handle to a held `Mutex`. +struct OwnedMappedMutexGuardExtendedLock { + _value: *mut U, + mutex: Arc>, +} + +impl OwnedMappedMutexGuard { + /// `skip_drop` prevents the `OwnedMutexGuard` from being automatically dropped, allowing manual control over the drop process. + /// This method returns an `OwnedMappedMutexGuardExtendedLock` which contains the original `Arc>`. + fn skip_drop(self) -> OwnedMappedMutexGuardExtendedLock { + // Prevents automatic drop by wrapping in `ManuallyDrop` + let selfi = mem::ManuallyDrop::new(self); + // Unsafely reads the `Arc>` from the `ManuallyDrop` wrapper + OwnedMappedMutexGuardExtendedLock { + mutex: unsafe { ptr::read(&selfi.mutex) }, + _value: selfi.value, + } + } + + /// Returns a locked view over a portion of the locked data. + #[inline] + pub fn map(mut this: Self, f: F) -> OwnedMappedMutexGuard + where + F: FnOnce(&mut U) -> &mut C, + { + let value = f(&mut *this) as *mut C; + let mutex = this.skip_drop(); + OwnedMappedMutexGuard { value, mutex: mutex.mutex } + } + + /// Returns a locked view over a portion of the locked data. + /// If the closure returns `None`, the original `OwnedMappedMutexGuard` is returned. + #[inline] + pub fn try_map(mut this: Self, f: F) -> Result, Self> + where + F: FnOnce(&mut U) -> Option<&mut C>, + { + let value = match f(&mut *this) { + Some(data) => data as *mut C, + None => return Err(this), + }; + let mutex = this.skip_drop(); + + Ok(OwnedMappedMutexGuard { value, mutex: mutex.mutex }) + } +} + /// A future which resolves when the target mutex has been successfully acquired. pub struct MutexLockFuture<'a, T: ?Sized> { // `None` indicates that the mutex was successfully acquired. @@ -544,6 +644,52 @@ unsafe impl Sync for OwnedMutexGuard {} unsafe impl Send for MappedMutexGuard<'_, T, U> {} unsafe impl Sync for MappedMutexGuard<'_, T, U> {} +impl Drop for OwnedMappedMutexGuard { + fn drop(&mut self) { + self.mutex.unlock() + } +} + +unsafe impl Sync for OwnedMappedMutexGuard +where + T: ?Sized + Send + Sync, + U: ?Sized + Send + Sync, +{ +} +unsafe impl Send for OwnedMappedMutexGuard +where + T: ?Sized + Send, + U: ?Sized + Send, +{ +} + +impl Deref for OwnedMappedMutexGuard { + type Target = U; + /// Deref implementation to allow access to the mapped value. + fn deref(&self) -> &Self::Target { + unsafe { &*self.value } + } +} + +impl DerefMut for OwnedMappedMutexGuard { + /// DerefMut implementation to allow mutable access to the mapped value. + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.value } + } +} + +impl fmt::Debug for OwnedMappedMutexGuard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&**self, f) + } +} + +impl fmt::Display for OwnedMappedMutexGuard { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&**self, f) + } +} + #[cfg(test)] mod tests { use super::*;