Skip to content

Commit

Permalink
Fix & test leak of some BTreeMap nodes on panic during into_iter
Browse files Browse the repository at this point in the history
  • Loading branch information
ssomers committed Mar 6, 2020
1 parent 865b44a commit 44c97c4
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
11 changes: 10 additions & 1 deletion src/liballoc/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1477,6 +1477,14 @@ impl<K, V> Drop for IntoIter<K, V> {
// Continue the same loop we perform below. This only runs when unwinding, so we
// don't have to care about panics this time (they'll abort).
while let Some(_) = self.0.next() {}

// No need to avoid the shared root, because the tree was definitely not empty.
unsafe {
let mut node = ptr::read(&self.0.front).into_node().forget_type();
while let Some(parent) = node.deallocate_and_ascend() {
node = parent.into_node().forget_type();
}
}
}
}

Expand All @@ -1491,7 +1499,8 @@ impl<K, V> Drop for IntoIter<K, V> {
if node.is_shared_root() {
return;
}

// Most of the nodes have been deallocated while traversing
// but one pile from a leaf up to the root is left standing.
while let Some(parent) = node.deallocate_and_ascend() {
node = parent.into_node().forget_type();
}
Expand Down
26 changes: 25 additions & 1 deletion src/liballoc/tests/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1021,7 +1021,7 @@ fn test_split_off_large_random_sorted() {
}

#[test]
fn test_into_iter_drop_leak() {
fn test_into_iter_drop_leak_1() {
static DROPS: AtomicU32 = AtomicU32::new(0);

struct D;
Expand All @@ -1045,3 +1045,27 @@ fn test_into_iter_drop_leak() {

assert_eq!(DROPS.load(Ordering::SeqCst), 5);
}

#[test]
fn test_into_iter_drop_leak_2() {
let size = 12; // to obtain tree with 2 levels (having edges to leaf nodes)
static DROPS: AtomicU32 = AtomicU32::new(0);
static PANIC_POINT: AtomicU32 = AtomicU32::new(0);

struct D;
impl Drop for D {
fn drop(&mut self) {
if DROPS.fetch_add(1, Ordering::SeqCst) == PANIC_POINT.load(Ordering::SeqCst) {
panic!("panic in `drop`");
}
}
}

for panic_point in vec![0, 1, size - 2, size - 1] {
DROPS.store(0, Ordering::SeqCst);
PANIC_POINT.store(panic_point, Ordering::SeqCst);
let map: BTreeMap<_, _> = (0..size).map(|i| (i, D)).collect();
catch_unwind(move || drop(map.into_iter())).ok();
assert_eq!(DROPS.load(Ordering::SeqCst), size);
}
}

0 comments on commit 44c97c4

Please sign in to comment.