From a1a07945d63e4130d0ca3c069fbffbd447178f0d Mon Sep 17 00:00:00 2001
From: Boxy <supbscripter@gmail.com>
Date: Fri, 1 Jul 2022 21:54:28 +0000
Subject: [PATCH] fix some memory leaks detected by miri (#4959)

The first leak:
```rust
    #[test]
    fn blob_vec_drop_empty_capacity() {
        let item_layout = Layout::new::<Foo>();
        let drop = drop_ptr::<Foo>;
        let _ = unsafe { BlobVec::new(item_layout, Some(drop), 0) };
    }
```
this is because we allocate the swap scratch in blobvec regardless of what the capacity is, but we only deallocate if capacity is > 0

The second leak:
```rust
    #[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),
            ]
        );
    }
```
this is caused by us not running the drop impl on the to-be-inserted component if the drop impl of the overwritten component panics

---

managed to figure out where the leaks were by using this 10/10 command
```
cargo --quiet test --lib -- --list | sed 's/: test$//' | MIRIFLAGS="-Zmiri-disable-isolation" xargs -n1 cargo miri test --lib -- --exact
```
which runs every test one by one rather than all at once which let miri actually tell me which test had the leak :upside_down_face:
---
 .github/workflows/ci.yml                |  5 ++---
 crates/bevy_ecs/src/storage/blob_vec.rs | 18 +++++++++++++++++-
 crates/bevy_ecs/src/world/mod.rs        |  1 +
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 56e7b468e8..abb68fee29 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -99,11 +99,10 @@ jobs:
           RUSTFLAGS: -Zrandomize-layout
           # https://github.com/rust-lang/miri#miri--z-flags-and-environment-variables
           # -Zmiri-disable-isolation is needed because our executor uses `fastrand` which accesses system time.
-          # -Zmiri-ignore-leaks is needed because running bevy_ecs tests finds a memory leak but its impossible
-          # to track down because allocids are nondeterministic.
           # -Zmiri-permissive-provenance disables warnings against int2ptr casts (since those are used by once_cell)
           # -Zmiri-disable-weak-memory-emulation works around https://github.com/bevyengine/bevy/issues/5164.
-          MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-ignore-leaks -Zmiri-permissive-provenance -Zmiri-disable-weak-memory-emulation
+          # -Zmiri-ignore-leaks is necessary because a bunch of tests don't join all threads before finishing.
+          MIRIFLAGS: -Zmiri-ignore-leaks -Zmiri-disable-isolation -Zmiri-permissive-provenance -Zmiri-disable-weak-memory-emulation
 
   check-compiles:
     runs-on: ubuntu-latest
diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs
index 14d584b048..1a4404d9b2 100644
--- a/crates/bevy_ecs/src/storage/blob_vec.rs
+++ b/crates/bevy_ecs/src/storage/blob_vec.rs
@@ -151,8 +151,20 @@ impl BlobVec {
         let ptr = self.get_unchecked_mut(index).promote().as_ptr();
         self.len = 0;
         // Drop the old value, then write back, justifying the promotion
+        // If the drop impl for the old value panics then we run the drop impl for `value` too.
         if let Some(drop) = self.drop {
+            struct OnDrop<F: FnMut()>(F);
+            impl<F: FnMut()> Drop for OnDrop<F> {
+                fn drop(&mut self) {
+                    (self.0)();
+                }
+            }
+            let value = value.as_ptr();
+            let on_unwind = OnDrop(|| (drop)(OwningPtr::new(NonNull::new_unchecked(value))));
+
             (drop)(OwningPtr::new(NonNull::new_unchecked(ptr)));
+
+            core::mem::forget(on_unwind);
         }
         std::ptr::copy_nonoverlapping::<u8>(value.as_ptr(), ptr, self.item_layout.size());
         self.len = old_len;
@@ -304,12 +316,16 @@ impl BlobVec {
 impl Drop for BlobVec {
     fn drop(&mut self) {
         self.clear();
+        if self.item_layout.size() > 0 {
+            unsafe {
+                std::alloc::dealloc(self.swap_scratch.as_ptr(), self.item_layout);
+            }
+        }
         let array_layout =
             array_layout(&self.item_layout, self.capacity).expect("array layout should be valid");
         if array_layout.size() > 0 {
             unsafe {
                 std::alloc::dealloc(self.get_ptr_mut().as_ptr(), array_layout);
-                std::alloc::dealloc(self.swap_scratch.as_ptr(), self.item_layout);
             }
         }
     }
diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
index 5f4a952e62..d220495a82 100644
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@@ -1693,6 +1693,7 @@ mod tests {
                 DropLogItem::Create(0),
                 DropLogItem::Create(1),
                 DropLogItem::Drop(0),
+                DropLogItem::Drop(1),
             ]
         );
     }