Fix leak in VecDeque::drain when drop panics

This commit is contained in:
Jonas Schievink 2019-12-13 19:32:16 +01:00
parent dc492452da
commit b04ca13873
2 changed files with 86 additions and 34 deletions

View File

@ -2575,47 +2575,61 @@ unsafe impl<T: Send> Send for Drain<'_, T> {}
#[stable(feature = "drain", since = "1.6.0")]
impl<T> Drop for Drain<'_, T> {
fn drop(&mut self) {
self.for_each(drop);
struct DropGuard<'r, 'a, T>(&'r mut Drain<'a, T>);
let source_deque = unsafe { self.deque.as_mut() };
impl<'r, 'a, T> Drop for DropGuard<'r, 'a, T> {
fn drop(&mut self) {
self.0.for_each(drop);
// T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head
//
// T t h H
// [. . . o o x x o o . . .]
//
let orig_tail = source_deque.tail;
let drain_tail = source_deque.head;
let drain_head = self.after_tail;
let orig_head = self.after_head;
let source_deque = unsafe { self.0.deque.as_mut() };
let tail_len = count(orig_tail, drain_tail, source_deque.cap());
let head_len = count(drain_head, orig_head, source_deque.cap());
// T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head
//
// T t h H
// [. . . o o x x o o . . .]
//
let orig_tail = source_deque.tail;
let drain_tail = source_deque.head;
let drain_head = self.0.after_tail;
let orig_head = self.0.after_head;
// Restore the original head value
source_deque.head = orig_head;
let tail_len = count(orig_tail, drain_tail, source_deque.cap());
let head_len = count(drain_head, orig_head, source_deque.cap());
match (tail_len, head_len) {
(0, 0) => {
source_deque.head = 0;
source_deque.tail = 0;
}
(0, _) => {
source_deque.tail = drain_head;
}
(_, 0) => {
source_deque.head = drain_tail;
}
_ => unsafe {
if tail_len <= head_len {
source_deque.tail = source_deque.wrap_sub(drain_head, tail_len);
source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len);
} else {
source_deque.head = source_deque.wrap_add(drain_tail, head_len);
source_deque.wrap_copy(drain_tail, drain_head, head_len);
// Restore the original head value
source_deque.head = orig_head;
match (tail_len, head_len) {
(0, 0) => {
source_deque.head = 0;
source_deque.tail = 0;
}
(0, _) => {
source_deque.tail = drain_head;
}
(_, 0) => {
source_deque.head = drain_tail;
}
_ => unsafe {
if tail_len <= head_len {
source_deque.tail = source_deque.wrap_sub(drain_head, tail_len);
source_deque.wrap_copy(source_deque.tail, orig_tail, tail_len);
} else {
source_deque.head = source_deque.wrap_add(drain_tail, head_len);
source_deque.wrap_copy(drain_tail, drain_head, head_len);
}
},
}
},
}
}
while let Some(item) = self.next() {
let guard = DropGuard(self);
drop(item);
mem::forget(guard);
}
DropGuard(self);
}
}

View File

@ -1606,3 +1606,41 @@ fn truncate_leak() {
assert_eq!(unsafe { DROPS }, 7);
}
#[test]
fn test_drain_leak() {
static mut DROPS: i32 = 0;
#[derive(Debug, PartialEq)]
struct D(u32, bool);
impl Drop for D {
fn drop(&mut self) {
unsafe {
DROPS += 1;
}
if self.1 {
panic!("panic in `drop`");
}
}
}
let mut v = VecDeque::new();
v.push_back(D(4, false));
v.push_back(D(5, false));
v.push_back(D(6, false));
v.push_front(D(3, false));
v.push_front(D(2, true));
v.push_front(D(1, false));
v.push_front(D(0, false));
catch_unwind(AssertUnwindSafe(|| {
v.drain(1..=4);
})).ok();
assert_eq!(unsafe { DROPS }, 4);
assert_eq!(v.len(), 3);
drop(v);
assert_eq!(unsafe { DROPS }, 7);
}