Auto merge of #132597 - lukas-code:btree-plug-leak, r=jhpratt

btree: don't leak value if destructor of key panics

This PR fixes a regression from #84904.

The `BTreeMap` already attempts to handle panicking destructors of the key-value pairs by continuing to execute the remaining destructors after one destructor panicked. However, after #84904 the destructor of a value in a key-value pair gets skipped if the destructor of the key panics, only continuing with the next key-value pair. This PR reverts to the behavior before #84904 to also drop the corresponding value if the destructor of a key panics.

This avoids potential memory leaks and can fix the soundness of programs that rely on the destructors being executed (even though this should not be relied upon, because the std collections currently do not guarantee that the remaining elements are dropped after a panic in a destructor).

cc `@Amanieu` because you had opinions on panicking destructors

library/alloc/src/collections/btree/map/tests.rs

@@ -2270,6 +2270,54 @@ fn test_into_iter_drop_leak_height_0() {
     assert_eq!(e.dropped(), 1);
 }
 
+#[test]
+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
+fn test_into_iter_drop_leak_kv_panic_in_key() {
+    let a_k = CrashTestDummy::new(0);
+    let a_v = CrashTestDummy::new(1);
+    let b_k = CrashTestDummy::new(2);
+    let b_v = CrashTestDummy::new(3);
+    let c_k = CrashTestDummy::new(4);
+    let c_v = CrashTestDummy::new(5);
+    let mut map = BTreeMap::new();
+    map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never));
+    map.insert(b_k.spawn(Panic::InDrop), b_v.spawn(Panic::Never));
+    map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never));
+
+    catch_unwind(move || drop(map.into_iter())).unwrap_err();
+
+    assert_eq!(a_k.dropped(), 1);
+    assert_eq!(a_v.dropped(), 1);
+    assert_eq!(b_k.dropped(), 1);
+    assert_eq!(b_v.dropped(), 1);
+    assert_eq!(c_k.dropped(), 1);
+    assert_eq!(c_v.dropped(), 1);
+}
+
+#[test]
+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
+fn test_into_iter_drop_leak_kv_panic_in_val() {
+    let a_k = CrashTestDummy::new(0);
+    let a_v = CrashTestDummy::new(1);
+    let b_k = CrashTestDummy::new(2);
+    let b_v = CrashTestDummy::new(3);
+    let c_k = CrashTestDummy::new(4);
+    let c_v = CrashTestDummy::new(5);
+    let mut map = BTreeMap::new();
+    map.insert(a_k.spawn(Panic::Never), a_v.spawn(Panic::Never));
+    map.insert(b_k.spawn(Panic::Never), b_v.spawn(Panic::InDrop));
+    map.insert(c_k.spawn(Panic::Never), c_v.spawn(Panic::Never));
+
+    catch_unwind(move || drop(map.into_iter())).unwrap_err();
+
+    assert_eq!(a_k.dropped(), 1);
+    assert_eq!(a_v.dropped(), 1);
+    assert_eq!(b_k.dropped(), 1);
+    assert_eq!(b_v.dropped(), 1);
+    assert_eq!(c_k.dropped(), 1);
+    assert_eq!(c_v.dropped(), 1);
+}
+
 #[test]
 #[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
 fn test_into_iter_drop_leak_height_1() {

library/alloc/src/collections/btree/node.rs

@@ -1173,11 +1173,25 @@ impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV>
     /// The node that the handle refers to must not yet have been deallocated.
     #[inline]
     pub unsafe fn drop_key_val(mut self) {
+        // Run the destructor of the value even if the destructor of the key panics.
+        struct Dropper<'a, T>(&'a mut MaybeUninit<T>);
+        impl<T> Drop for Dropper<'_, T> {
+            #[inline]
+            fn drop(&mut self) {
+                unsafe {
+                    self.0.assume_init_drop();
+                }
+            }
+        }
+
         debug_assert!(self.idx < self.node.len());
         let leaf = self.node.as_leaf_dying();
         unsafe {
-            leaf.keys.get_unchecked_mut(self.idx).assume_init_drop();
-            leaf.vals.get_unchecked_mut(self.idx).assume_init_drop();
+            let key = leaf.keys.get_unchecked_mut(self.idx);
+            let val = leaf.vals.get_unchecked_mut(self.idx);
+            let _guard = Dropper(val);
+            key.assume_init_drop();
+            // dropping the guard will drop the value
         }
     }
 }