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 all commits
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
109 changes: 109 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,115 @@ fn drain_filter_complex() {
}
}

#[test]
#[cfg(not(miri))] // Miri does not support catching panics
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]
#[cfg(not(miri))] // Miri does not support catching panics
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 drain_filter_unconsumed() {
let mut vec = vec![1, 2, 3, 4];
let drain = vec.drain_filter(|&mut x| x % 2 != 0);
drop(drain);
assert_eq!(vec, [2, 4]);
}

#[test]
fn test_reserve_exact() {
// This is all the same as test_reserve
Expand Down
69 changes: 60 additions & 9 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 @@ -2747,10 +2748,20 @@ pub struct DrainFilter<'a, T, F>
where F: FnMut(&mut T) -> bool,
{
vec: &'a mut Vec<T>,
/// The index of the item that will be inspected by the next call to `next`.
idx: usize,
/// The number of items that have been drained (removed) thus far.
del: usize,
/// The original length of `vec` prior to draining.
old_len: usize,
/// The filter test predicate.
pred: F,
/// A flag that indicates a panic has occured in the filter test prodicate.
/// This is used as a hint in the drop implmentation to prevent consumption
/// of the remainder of the `DrainFilter`. Any unprocessed items will be
/// backshifted in the `vec`, but no further items will be dropped or
/// tested by the filter predicate.
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 @@ -2761,20 +2772,23 @@ impl<T, F> Iterator for DrainFilter<'_, T, F>

fn next(&mut self) -> Option<T> {
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]) {
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.

// Update the index *after* the predicate is called. If the index
// is updated prior and the predicate panics, the element at this
// index would be leaked.
self.idx += 1;
if drained {
self.del += 1;
return Some(ptr::read(&v[i]));
} 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 +2806,46 @@ 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);
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 {
if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
// This is a pretty messed up state, and there isn't really an
// obviously right thing to do. 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. The backshift
// is required to prevent a double-drop of the last successfully
// drained item prior to a panic in the predicate.
let ptr = self.drain.vec.as_mut_ptr();
let src = ptr.add(self.drain.idx);
let dst = src.sub(self.drain.del);
let tail_len = self.drain.old_len - self.drain.idx;
src.copy_to(dst, tail_len);
}
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
};

// Attempt to consume any remaining elements if the filter predicate
// has not yet panicked. We'll backshift any remaining elements
// whether we've already panicked or if the consumption here panics.
if !backshift.drain.panic_flag {
backshift.drain.for_each(drop);
}
}
}