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

@@ -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

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);

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<Mutex<Vec<DropLogItem>>>,
+        expected_panic_flag: Arc<AtomicBool>,
+        should_panic: bool,
+        id: u8,
+    }
+
+    impl MayPanicInDrop {
+        fn new(
+            drop_log: &Arc<Mutex<Vec<DropLogItem>>>,
+            expected_panic_flag: &Arc<AtomicBool>,
+            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<Mutex<Vec<DropLogItem>>>,
+        /// 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<AtomicBool>,
+    }
+
+    impl DropTestHelper {
+        pub fn new() -> Self {
+            Self {
+                drop_log: Arc::new(Mutex::new(Vec::<DropLogItem>::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<DropLogItem> {
+            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)
+            ]
+        );
+    }
+}