Fix binary_heap::DrainSorted drop leak on panics

Author: jonas-schievink

src/liballoc/collections/binary_heap.rs

@@ -147,7 +147,7 @@
 
 use core::fmt;
 use core::iter::{FromIterator, FusedIterator, TrustedLen};
-use core::mem::{size_of, swap, ManuallyDrop};
+use core::mem::{self, size_of, swap, ManuallyDrop};
 use core::ops::{Deref, DerefMut};
 use core::ptr;
 
@@ -1239,7 +1239,19 @@ pub struct DrainSorted<'a, T: Ord> {
 impl<'a, T: Ord> Drop for DrainSorted<'a, T> {
     /// Removes heap elements in heap order.
     fn drop(&mut self) {
-        while let Some(_) = self.inner.pop() {}
+        struct DropGuard<'r, 'a, T: Ord>(&'r mut DrainSorted<'a, T>);
+
+        impl<'r, 'a, T: Ord> Drop for DropGuard<'r, 'a, T> {
+            fn drop(&mut self) {
+                while let Some(_) = self.0.inner.pop() {}
+            }
+        }
+
+        while let Some(item) = self.inner.pop() {
+            let guard = DropGuard(self);
+            drop(item);
+            mem::forget(guard);
+        }
     }
 }
 

src/liballoc/tests/binary_heap.rs

@@ -1,6 +1,8 @@
 use std::collections::binary_heap::{Drain, PeekMut};
 use std::collections::BinaryHeap;
 use std::iter::TrustedLen;
+use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::sync::atomic::{AtomicU32, Ordering};
 
 #[test]
 fn test_iterator() {
@@ -275,6 +277,37 @@ fn test_drain_sorted() {
     assert!(q.is_empty());
 }
 
+#[test]
+fn test_drain_sorted_leak() {
+    static DROPS: AtomicU32 = AtomicU32::new(0);
+
+    #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
+    struct D(u32, bool);
+
+    impl Drop for D {
+        fn drop(&mut self) {
+            DROPS.fetch_add(1, Ordering::SeqCst);
+
+            if self.1 {
+                panic!("panic in `drop`");
+            }
+        }
+    }
+
+    let mut q = BinaryHeap::from(vec![
+        D(0, false),
+        D(1, false),
+        D(2, false),
+        D(3, true),
+        D(4, false),
+        D(5, false),
+    ]);
+
+    catch_unwind(AssertUnwindSafe(|| drop(q.drain_sorted()))).ok();
+
+    assert_eq!(DROPS.load(Ordering::SeqCst), 6);
+}
+
 #[test]
 fn test_extend_ref() {
     let mut a = BinaryHeap::new();