From 3cf10e2ef2d5b2e6c5e8ea558818a079974c68e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Grabarz?= <frizi09@gmail.com>
Date: Mon, 17 May 2021 20:01:25 +0000
Subject: [PATCH] prevent memory leak when dropping ParallelSystemContainer
 (#2176)

`ParallelSystemContainer`'s `system` pointer was extracted from box, but it was never deallocated. This change adds missing drop implementation that cleans up that memory.
---
 crates/bevy_ecs/src/schedule/system_container.rs | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/schedule/system_container.rs b/crates/bevy_ecs/src/schedule/system_container.rs
index 7040ae80ea..727c007e68 100644
--- a/crates/bevy_ecs/src/schedule/system_container.rs
+++ b/crates/bevy_ecs/src/schedule/system_container.rs
@@ -7,7 +7,7 @@ use crate::{
     },
     system::{ExclusiveSystem, System},
 };
-use std::{borrow::Cow, ptr::NonNull};
+use std::{borrow::Cow, cell::UnsafeCell};
 
 /// System metadata like its name, labels, order requirements and component access.
 pub trait SystemContainer: GraphNode<Label = BoxedSystemLabel> {
@@ -106,7 +106,7 @@ impl SystemContainer for ExclusiveSystemContainer {
 }
 
 pub struct ParallelSystemContainer {
-    system: NonNull<dyn System<In = (), Out = ()>>,
+    system: Box<UnsafeCell<dyn System<In = (), Out = ()>>>,
     pub(crate) run_criteria_index: Option<usize>,
     pub(crate) run_criteria_label: Option<BoxedRunCriteriaLabel>,
     pub(crate) should_run: bool,
@@ -123,7 +123,8 @@ unsafe impl Sync for ParallelSystemContainer {}
 impl ParallelSystemContainer {
     pub(crate) fn from_descriptor(descriptor: ParallelSystemDescriptor) -> Self {
         ParallelSystemContainer {
-            system: unsafe { NonNull::new_unchecked(Box::into_raw(descriptor.system)) },
+            // SAFE: it is fine to wrap inner value with UnsafeCell, as it is repr(transparent)
+            system: unsafe { Box::from_raw(Box::into_raw(descriptor.system) as *mut _) },
             should_run: false,
             run_criteria_index: None,
             run_criteria_label: None,
@@ -140,20 +141,19 @@ impl ParallelSystemContainer {
     }
 
     pub fn system(&self) -> &dyn System<In = (), Out = ()> {
-        // SAFE: statically enforced shared access.
-        unsafe { self.system.as_ref() }
+        // SAFE: statically enforced shared access
+        unsafe { self.system.get().as_ref().unwrap() }
     }
 
     pub fn system_mut(&mut self) -> &mut dyn System<In = (), Out = ()> {
-        // SAFE: statically enforced exclusive access.
-        unsafe { self.system.as_mut() }
+        self.system.get_mut()
     }
 
     /// # Safety
     /// Ensure no other borrows exist along with this one.
     #[allow(clippy::mut_from_ref)]
     pub unsafe fn system_mut_unsafe(&self) -> &mut dyn System<In = (), Out = ()> {
-        &mut *self.system.as_ptr()
+        self.system.get().as_mut().unwrap()
     }
 
     pub fn should_run(&self) -> bool {