Skip to content

Prevent Vec::drain_filter from double dropping on panic #61224

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jul 9, 2019
Merged

Prevent Vec::drain_filter from double dropping on panic #61224

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
84ae969
Prevent Vec::drain_filter from double dropping on panic
aloucks May 27, 2019
17a517a
Fix formatting nit
aloucks May 27, 2019
53d46ae
Add drain_filter_unconsumed test
aloucks May 27, 2019
b13ae65
Refactor DrainFilter::next and update comments
aloucks May 27, 2019
f5ab031
Disable drain_filter tests that require catch_unwind on miri
aloucks May 27, 2019
a4a6a67
Remove while loop in DrainFilter::drop and add additional docs
aloucks Jul 7, 2019
df5b32e
Clarify double-drop comment
aloucks Jul 7, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 99 additions & 0 deletions src/liballoc/tests/vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -945,6 +945,105 @@ fn drain_filter_complex() {
}
}

#[test]
fn drain_filter_consumed_panic() {
use std::rc::Rc;
use std::sync::Mutex;

struct Check {
index: usize,
drop_counts: Rc<Mutex<Vec<usize>>>,
};

impl Drop for Check {
fn drop(&mut self) {
self.drop_counts.lock().unwrap()[self.index] += 1;
println!("drop: {}", self.index);
}
}

let check_count = 10;
let drop_counts = Rc::new(Mutex::new(vec![0_usize; check_count]));
let mut data: Vec<Check> = (0..check_count)
.map(|index| Check { index, drop_counts: Rc::clone(&drop_counts) })
.collect();

let _ = std::panic::catch_unwind(move || {
let filter = |c: &mut Check| {
if c.index == 2 {
panic!("panic at index: {}", c.index);
}
// Verify that if the filter could panic again on another element
// that it would not cause a double panic and all elements of the
// vec would still be dropped exactly once.
if c.index == 4 {
panic!("panic at index: {}", c.index);
}
c.index < 6
};
let drain = data.drain_filter(filter);

// NOTE: The DrainFilter is explictly consumed
drain.for_each(drop);
});

let drop_counts = drop_counts.lock().unwrap();
assert_eq!(check_count, drop_counts.len());

for (index, count) in drop_counts.iter().cloned().enumerate() {
assert_eq!(1, count, "unexpected drop count at index: {} (count: {})", index, count);
}
}

#[test]
fn drain_filter_unconsumed_panic() {
use std::rc::Rc;
use std::sync::Mutex;

struct Check {
index: usize,
drop_counts: Rc<Mutex<Vec<usize>>>,
};

impl Drop for Check {
fn drop(&mut self) {
self.drop_counts.lock().unwrap()[self.index] += 1;
println!("drop: {}", self.index);
}
}

let check_count = 10;
let drop_counts = Rc::new(Mutex::new(vec![0_usize; check_count]));
let mut data: Vec<Check> = (0..check_count)
.map(|index| Check { index, drop_counts: Rc::clone(&drop_counts) })
.collect();

let _ = std::panic::catch_unwind(move || {
let filter = |c: &mut Check| {
if c.index == 2 {
panic!("panic at index: {}", c.index);
}
// Verify that if the filter could panic again on another element
// that it would not cause a double panic and all elements of the
// vec would still be dropped exactly once.
if c.index == 4 {
panic!("panic at index: {}", c.index);
}
c.index < 6
};
let _drain = data.drain_filter(filter);

// NOTE: The DrainFilter is dropped without being consumed
});

let drop_counts = drop_counts.lock().unwrap();
assert_eq!(check_count, drop_counts.len());

for (index, count) in drop_counts.iter().cloned().enumerate() {
assert_eq!(1, count, "unexpected drop count at index: {} (count: {})", index, count);
}
}

#[test]
fn test_reserve_exact() {
// This is all the same as test_reserve
Expand Down
73 changes: 63 additions & 10 deletions src/liballoc/vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2120,6 +2120,7 @@ impl<T> Vec<T> {
del: 0,
old_len,
pred: filter,
panic_flag: false,
}
}
}
Expand Down Expand Up @@ -2751,6 +2752,7 @@ pub struct DrainFilter<'a, T, F>
del: usize,
old_len: usize,
pred: F,
panic_flag: bool,
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This has grown enough fields that they should really be documented. Either individually here, or with a high-level description of what we're trying to do in the implementation of next.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added docs to each field in the latest commit.


#[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
Expand All @@ -2760,21 +2762,34 @@ impl<T, F> Iterator for DrainFilter<'_, T, F>
type Item = T;

fn next(&mut self) -> Option<T> {
struct SetIdxOnDrop<'a> {
idx: &'a mut usize,
new_idx: usize,
}

impl<'a> Drop for SetIdxOnDrop<'a> {
fn drop(&mut self) {
*self.idx = self.new_idx;
}
}

unsafe {
while self.idx != self.old_len {
while self.idx < self.old_len {
let i = self.idx;
self.idx += 1;
let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
if (self.pred)(&mut v[i]) {
let mut set_idx = SetIdxOnDrop { new_idx: self.idx, idx: &mut self.idx };
self.panic_flag = true;
let drained = (self.pred)(&mut v[i]);
self.panic_flag = false;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a terrible compulsion to try to encode this state in some magic combination of old_len/idx/del, but this is probably clearest, and easiest for llvm to evaporate when it sees pred can't unwind.

Copy link
Contributor Author

@aloucks aloucks Jul 7, 2019 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I considered that route initially but came to the conclusion that even if it's possible, the simple flag would be much easier to understand and maintain.

set_idx.new_idx += 1;
if drained {
self.del += 1;
return Some(ptr::read(&v[i]));
} else if self.del > 0 {
}
else if self.del > 0 {
let del = self.del;
let src: *const T = &v[i];
let dst: *mut T = &mut v[i - del];
// This is safe because self.vec has length 0
// thus its elements will not have Drop::drop
// called on them in the event of a panic.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this comment no longer true? Why is it safe instead then?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe the intent was to leak instead of double drop, but it didn't quite work. I've done some minor refactoring and added additional comments.

It's safe now because there are additional checks in DrainFilter::drop that perform the cleanup in the event of a panic in the filter predicate.

ptr::copy_nonoverlapping(src, dst, 1);
}
}
Expand All @@ -2792,9 +2807,47 @@ impl<T, F> Drop for DrainFilter<'_, T, F>
where F: FnMut(&mut T) -> bool,
{
fn drop(&mut self) {
self.for_each(drop);
unsafe {
self.vec.set_len(self.old_len - self.del);
// If the predicate panics, we still need to backshift everything
// down after the last successfully drained element, but no additional
// elements are drained or checked.
struct BackshiftOnDrop<'a, 'b, T, F>
where
F: FnMut(&mut T) -> bool,
{
drain: &'b mut DrainFilter<'a, T, F>,
}

impl<'a, 'b, T, F> Drop for BackshiftOnDrop<'a, 'b, T, F>
where
F: FnMut(&mut T) -> bool
{
fn drop(&mut self) {
unsafe {
while self.drain.idx < self.drain.old_len {
let i = self.drain.idx;
self.drain.idx += 1;
let v = slice::from_raw_parts_mut(
self.drain.vec.as_mut_ptr(),
self.drain.old_len,
);
if self.drain.del > 0 {
let del = self.drain.del;
let src: *const T = &v[i];
let dst: *mut T = &mut v[i - del];
ptr::copy_nonoverlapping(src, dst, 1);
}
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why this while loop exists. Surely this should just be:

if self.drain.idx < self.drain.old_len {
  // It looks like `pred` panicked, so we didn't process all the elements.
  // This is a pretty messed up state, and there isn't really an obviously right
  // thing to do (and we don't want to keep trying to execute `pred`). So we
  // just backshift all the unprocessed elements and tell the vec that they still
  // exist, hoping that doesn't mess up anyone further along in the panic.
  let idx = self.drain.idx;
  let num_deleted = self.drain.del;
  let tail_len = self.drain.old_len - idx;
  let ptr = self.drain.vec.as_mut_ptr();
  if num_deleted > 0 {
    ptr.add(idx).copy_to(ptr.add(idx - num_deleted), tail_len);
  }
}

self.drain.vec.set_len(self.drain.old_len - self.drain.del);

(Here I modernized the code a bit to use the newer raw pointer APIs, and some clearer names. It would be a nice cleanup to this code if you also did that to the Iterator's fields and its next impl as well, although not a blocker.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I incorporated this in the latest commit. I also consolidated the the self.drain.del check into the parent if and made the src and dst pointers explicit while still using the more modern pointer APIs.

self.drain.vec.set_len(self.drain.old_len - self.drain.del);
}
}
}

let backshift = BackshiftOnDrop {
drain: self
};

if !backshift.drain.panic_flag {
backshift.drain.for_each(drop);
}
}
}