From 8c250919e3f8ae9d485150966c4c8b7f38d8e543 Mon Sep 17 00:00:00 2001 From: sapir Date: Sat, 18 Dec 2021 21:29:24 +0000 Subject: [PATCH] Fix double drop in BlobVec::replace_unchecked (#2597) (#2848) # Objective I thought I'd have a go a trying to fix #2597. Hopefully fixes #2597. ## Solution I reused the memory pointed to by the value parameter, that is already required by `insert` to not be dropped, to contain the extracted value while dropping it. --- crates/bevy_ecs/src/storage/blob_vec.rs | 21 +++- crates/bevy_ecs/src/storage/sparse_set.rs | 15 ++- crates/bevy_ecs/src/world/mod.rs | 133 ++++++++++++++++++++++ 3 files changed, 162 insertions(+), 7 deletions(-) diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs index 49c92d92d1eb9..f6b890e7be2e4 100644 --- a/crates/bevy_ecs/src/storage/blob_vec.rs +++ b/crates/bevy_ecs/src/storage/blob_vec.rs @@ -87,7 +87,12 @@ impl BlobVec { /// # Safety /// - index must be in bounds - /// - memory must be reserved and uninitialized + /// - the memory in the `BlobVec` starting at index `index`, of a size matching this `BlobVec`'s + /// `item_layout`, must have been previously allocated, but not initialized yet + /// - the memory at `*value` must be previously initialized with an item matching this + /// `BlobVec`'s `item_layout` + /// - the item that was stored in `*value` is left logically uninitialised/moved out of after + /// calling this function, and as such should not be used or dropped by the caller. #[inline] pub unsafe fn initialize_unchecked(&mut self, index: usize, value: *mut u8) { debug_assert!(index < self.len()); @@ -97,12 +102,24 @@ impl BlobVec { /// # Safety /// - index must be in-bounds - // - memory must be previously initialized + /// - the memory in the `BlobVec` starting at index `index`, of a size matching this `BlobVec`'s + /// `item_layout`, must have been previously initialized with an item matching this `BlobVec`'s + /// item_layout + /// - the memory at `*value` must also be previously initialized with an item matching this + /// `BlobVec`'s `item_layout` + /// - the item that was stored in `*value` is left logically uninitialised/moved out of after + /// calling this function, and as such should not be used or dropped by the caller. pub unsafe fn replace_unchecked(&mut self, index: usize, value: *mut u8) { debug_assert!(index < self.len()); let ptr = self.get_unchecked(index); + // If `drop` panics, then when the collection is dropped during stack unwinding, the + // collection's `Drop` impl will call `drop` again for the old value (which is still stored + // in the collection), so we get a double drop. To prevent that, we set len to 0 until we're + // done. + let old_len = std::mem::replace(&mut self.len, 0); (self.drop)(ptr); std::ptr::copy_nonoverlapping(value, ptr, self.item_layout.size()); + self.len = old_len; } /// increases the length by one (and grows the vec if needed) with uninitialized memory and diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 42357ec87c6a3..8ad76046e6cdc 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -123,13 +123,18 @@ impl ComponentSparseSet { self.dense.len() == 0 } - /// Inserts the `entity` key and component `value` pair into this sparse set. - /// The caller is responsible for ensuring the value is not dropped. This collection will drop - /// the value when needed. + /// Inserts the `entity` key and component `value` pair into this sparse + /// set. This collection takes ownership of the contents of `value`, and + /// will drop the value when needed. Also, it may overwrite the contents of + /// the `value` pointer if convenient. The caller is responsible for + /// ensuring it does not drop `*value` after calling `insert`. /// /// # Safety - /// The `value` pointer must point to a valid address that matches the `Layout` - /// inside the `ComponentInfo` given when constructing this sparse set. + /// * The `value` pointer must point to a valid address that matches the + /// `Layout` inside the `ComponentInfo` given when constructing this + /// sparse set. + /// * The caller is responsible for ensuring it does not drop `*value` after + /// calling `insert`. pub unsafe fn insert(&mut self, entity: Entity, value: *mut u8, change_tick: u32) { if let Some(&dense_index) = self.sparse.get(entity) { self.dense.replace_unchecked(dense_index, value); diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index e9f56a58a8ec9..739a1d9d30f31 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -1191,3 +1191,136 @@ impl Default for MainThreadValidator { } } } + +#[cfg(test)] +mod tests { + use super::World; + use bevy_ecs_macros::Component; + use std::{ + panic, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + }; + + // For bevy_ecs_macros + use crate as bevy_ecs; + + type ID = u8; + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum DropLogItem { + Create(ID), + Drop(ID), + } + + #[derive(Component)] + struct MayPanicInDrop { + drop_log: Arc>>, + expected_panic_flag: Arc, + should_panic: bool, + id: u8, + } + + impl MayPanicInDrop { + fn new( + drop_log: &Arc>>, + expected_panic_flag: &Arc, + should_panic: bool, + id: u8, + ) -> Self { + println!("creating component with id {}", id); + drop_log.lock().unwrap().push(DropLogItem::Create(id)); + + Self { + drop_log: Arc::clone(drop_log), + expected_panic_flag: Arc::clone(expected_panic_flag), + should_panic, + id, + } + } + } + + impl Drop for MayPanicInDrop { + fn drop(&mut self) { + println!("dropping component with id {}", self.id); + + { + let mut drop_log = self.drop_log.lock().unwrap(); + drop_log.push(DropLogItem::Drop(self.id)); + // Don't keep the mutex while panicking, or we'll poison it. + drop(drop_log); + } + + if self.should_panic { + self.expected_panic_flag.store(true, Ordering::SeqCst); + panic!("testing what happens on panic inside drop"); + } + } + } + + struct DropTestHelper { + drop_log: Arc>>, + /// Set to `true` right before we intentionally panic, so that if we get + /// a panic, we know if it was intended or not. + expected_panic_flag: Arc, + } + + impl DropTestHelper { + pub fn new() -> Self { + Self { + drop_log: Arc::new(Mutex::new(Vec::::new())), + expected_panic_flag: Arc::new(AtomicBool::new(false)), + } + } + + pub fn make_component(&self, should_panic: bool, id: ID) -> MayPanicInDrop { + MayPanicInDrop::new(&self.drop_log, &self.expected_panic_flag, should_panic, id) + } + + pub fn finish(self, panic_res: std::thread::Result<()>) -> Vec { + let drop_log = Arc::try_unwrap(self.drop_log) + .unwrap() + .into_inner() + .unwrap(); + let expected_panic_flag = self.expected_panic_flag.load(Ordering::SeqCst); + + if !expected_panic_flag { + match panic_res { + Ok(()) => panic!("Expected a panic but it didn't happen"), + Err(e) => panic::resume_unwind(e), + } + } + + drop_log + } + } + + #[test] + fn panic_while_overwriting_component() { + let helper = DropTestHelper::new(); + + let res = panic::catch_unwind(|| { + let mut world = World::new(); + world + .spawn() + .insert(helper.make_component(true, 0)) + .insert(helper.make_component(false, 1)); + + println!("Done inserting! Dropping world..."); + }); + + let drop_log = helper.finish(res); + + assert_eq!( + &*drop_log, + [ + DropLogItem::Create(0), + DropLogItem::Create(1), + DropLogItem::Drop(0), + DropLogItem::Drop(1) + ] + ); + } +}