Specify behavior if the closure passed to *Guard::*map panics.

Author: zachs18

Diff for: library/std/src/sync/mutex.rs

@@ -612,10 +612,14 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> {
         F: FnOnce(&mut T) -> &mut U,
         U: ?Sized,
     {
-        let mut orig = ManuallyDrop::new(orig);
-        let value = NonNull::from(f(&mut *orig));
+        // SAFETY: the conditions of `MutedGuard::new` were satisfied when the original guard
+        // was created, and have been upheld throughout `map` and/or `try_map`.
+        // The signature of the closure guarantees that it will not "leak" the lifetime of the reference
+        // passed to it. If the closure panics, the guard will be dropped.
+        let data = NonNull::from(f(unsafe { &mut *orig.lock.data.get() }));
+        let orig = ManuallyDrop::new(orig);
         MappedMutexGuard {
-            data: value,
+            data,
             inner: &orig.lock.inner,
             poison_flag: &orig.lock.poison,
             poison: orig.poison.clone(),
@@ -639,16 +643,23 @@ impl<'a, T: ?Sized> MutexGuard<'a, T> {
         F: FnOnce(&mut T) -> Option<&mut U>,
         U: ?Sized,
     {
-        let mut orig = ManuallyDrop::new(orig);
-        match f(&mut *orig).map(NonNull::from) {
-            Some(value) => Ok(MappedMutexGuard {
-                data: value,
-                inner: &orig.lock.inner,
-                poison_flag: &orig.lock.poison,
-                poison: orig.poison.clone(),
-                _variance: PhantomData,
-            }),
-            None => Err(ManuallyDrop::into_inner(orig)),
+        // SAFETY: the conditions of `MutexGuard::new` were satisfied when the original guard
+        // was created, and have been upheld throughout `map` and/or `try_map`.
+        // The signature of the closure guarantees that it will not "leak" the lifetime of the reference
+        // passed to it. If the closure panics, the guard will be dropped.
+        match f(unsafe { &mut *orig.lock.data.get() }) {
+            Some(data) => {
+                let data = NonNull::from(data);
+                let orig = ManuallyDrop::new(orig);
+                Ok(MappedMutexGuard {
+                    data,
+                    inner: &orig.lock.inner,
+                    poison_flag: &orig.lock.poison,
+                    poison: orig.poison.clone(),
+                    _variance: PhantomData,
+                })
+            }
+            None => Err(orig),
         }
     }
 }
@@ -704,15 +715,19 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> {
     /// `MappedMutexGuard::map(...)`. A method would interfere with methods of the
     /// same name on the contents of the `MutexGuard` used through `Deref`.
     #[unstable(feature = "mapped_lock_guards", issue = "117108")]
-    pub fn map<U, F>(orig: Self, f: F) -> MappedMutexGuard<'a, U>
+    pub fn map<U, F>(mut orig: Self, f: F) -> MappedMutexGuard<'a, U>
     where
         F: FnOnce(&mut T) -> &mut U,
         U: ?Sized,
     {
-        let mut orig = ManuallyDrop::new(orig);
-        let value = NonNull::from(f(&mut *orig));
+        // SAFETY: the conditions of `MutedGuard::new` were satisfied when the original guard
+        // was created, and have been upheld throughout `map` and/or `try_map`.
+        // The signature of the closure guarantees that it will not "leak" the lifetime of the reference
+        // passed to it. If the closure panics, the guard will be dropped.
+        let data = NonNull::from(f(unsafe { orig.data.as_mut() }));
+        let orig = ManuallyDrop::new(orig);
         MappedMutexGuard {
-            data: value,
+            data,
             inner: orig.inner,
             poison_flag: orig.poison_flag,
             poison: orig.poison.clone(),
@@ -731,21 +746,28 @@ impl<'a, T: ?Sized> MappedMutexGuard<'a, T> {
     /// same name on the contents of the `MutexGuard` used through `Deref`.
     #[doc(alias = "filter_map")]
     #[unstable(feature = "mapped_lock_guards", issue = "117108")]
-    pub fn try_map<U, F>(orig: Self, f: F) -> Result<MappedMutexGuard<'a, U>, Self>
+    pub fn try_map<U, F>(mut orig: Self, f: F) -> Result<MappedMutexGuard<'a, U>, Self>
     where
         F: FnOnce(&mut T) -> Option<&mut U>,
         U: ?Sized,
     {
-        let mut orig = ManuallyDrop::new(orig);
-        match f(&mut *orig).map(NonNull::from) {
-            Some(value) => Ok(MappedMutexGuard {
-                data: value,
-                inner: orig.inner,
-                poison_flag: orig.poison_flag,
-                poison: orig.poison.clone(),
-                _variance: PhantomData,
-            }),
-            None => Err(ManuallyDrop::into_inner(orig)),
+        // SAFETY: the conditions of `MutedGuard::new` were satisfied when the original guard
+        // was created, and have been upheld throughout `map` and/or `try_map`.
+        // The signature of the closure guarantees that it will not "leak" the lifetime of the reference
+        // passed to it. If the closure panics, the guard will be dropped.
+        match f(unsafe { orig.data.as_mut() }) {
+            Some(data) => {
+                let data = NonNull::from(data);
+                let orig = ManuallyDrop::new(orig);
+                Ok(MappedMutexGuard {
+                    data,
+                    inner: orig.inner,
+                    poison_flag: orig.poison_flag,
+                    poison: orig.poison.clone(),
+                    _variance: PhantomData,
+                })
+            }
+            None => Err(orig),
         }
     }
 }

Diff for: library/std/src/sync/mutex/tests.rs

@@ -1,6 +1,6 @@
 use crate::sync::atomic::{AtomicUsize, Ordering};
 use crate::sync::mpsc::channel;
-use crate::sync::{Arc, Condvar, MappedMutexGuard, Mutex, MutexGuard};
+use crate::sync::{Arc, Condvar, MappedMutexGuard, Mutex, MutexGuard, TryLockError};
 use crate::thread;
 
 struct Packet<T>(Arc<(Mutex<T>, Condvar)>);
@@ -264,3 +264,64 @@ fn test_mapping_mapped_guard() {
     drop(guard);
     assert_eq!(*lock.get_mut().unwrap(), [0, 42, 0, 0]);
 }
+
+#[test]
+fn panic_while_mapping_unlocked_poison() {
+    let lock = Mutex::new(());
+
+    let _ = crate::panic::catch_unwind(|| {
+        let guard = lock.lock().unwrap();
+        let _guard = MutexGuard::map::<(), _>(guard, |_| panic!());
+    });
+
+    match lock.try_lock() {
+        Ok(_) => panic!("panicking in a MutexGuard::map closure should poison the Mutex"),
+        Err(TryLockError::WouldBlock) => {
+            panic!("panicking in a MutexGuard::map closure should unlock the mutex")
+        }
+        Err(TryLockError::Poisoned(_)) => {}
+    }
+
+    let _ = crate::panic::catch_unwind(|| {
+        let guard = lock.lock().unwrap();
+        let _guard = MutexGuard::try_map::<(), _>(guard, |_| panic!());
+    });
+
+    match lock.try_lock() {
+        Ok(_) => panic!("panicking in a MutexGuard::try_map closure should poison the Mutex"),
+        Err(TryLockError::WouldBlock) => {
+            panic!("panicking in a MutexGuard::try_map closure should unlock the mutex")
+        }
+        Err(TryLockError::Poisoned(_)) => {}
+    }
+
+    let _ = crate::panic::catch_unwind(|| {
+        let guard = lock.lock().unwrap();
+        let guard = MutexGuard::map::<(), _>(guard, |val| val);
+        let _guard = MappedMutexGuard::map::<(), _>(guard, |_| panic!());
+    });
+
+    match lock.try_lock() {
+        Ok(_) => panic!("panicking in a MappedMutexGuard::map closure should poison the Mutex"),
+        Err(TryLockError::WouldBlock) => {
+            panic!("panicking in a MappedMutexGuard::map closure should unlock the mutex")
+        }
+        Err(TryLockError::Poisoned(_)) => {}
+    }
+
+    let _ = crate::panic::catch_unwind(|| {
+        let guard = lock.lock().unwrap();
+        let guard = MutexGuard::map::<(), _>(guard, |val| val);
+        let _guard = MappedMutexGuard::try_map::<(), _>(guard, |_| panic!());
+    });
+
+    match lock.try_lock() {
+        Ok(_) => panic!("panicking in a MappedMutexGuard::try_map closure should poison the Mutex"),
+        Err(TryLockError::WouldBlock) => {
+            panic!("panicking in a MappedMutexGuard::try_map closure should unlock the mutex")
+        }
+        Err(TryLockError::Poisoned(_)) => {}
+    }
+
+    drop(lock);
+}