Skip to content

Commit

Permalink
Implement Drop for CommandQueue (#10746)
Browse files Browse the repository at this point in the history
# Objective

- Fixes #10676, preventing a possible memory leak for commands which
owned resources.

## Solution

Implemented `Drop` for `CommandQueue`. This has been done entirely in
the private API of `CommandQueue`, ensuring no breaking changes. Also
added a unit test, `test_command_queue_inner_drop_early`, based on the
reproduction steps as outlined in #10676.

## Notes

I believe this can be applied to `0.12.1` as well, but I am uncertain of
the process to make that kind of change. Please let me know if there's
anything I can do to help with the back-porting of this change.

---------

Co-authored-by: Carter Anderson <[email protected]>
  • Loading branch information
bushrat011899 and cart committed Nov 30, 2023
1 parent f5ccf68 commit 86d0b8a
Showing 1 changed file with 48 additions and 5 deletions.
53 changes: 48 additions & 5 deletions crates/bevy_ecs/src/system/commands/command_queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ struct CommandMeta {
/// SAFETY: The `value` must point to a value of type `T: Command`,
/// where `T` is some specific type that was used to produce this metadata.
///
/// `world` is optional to allow this one function pointer to perform double-duty as a drop.
///
/// Returns the size of `T` in bytes.
apply_command_and_get_size: unsafe fn(value: OwningPtr<Unaligned>, world: &mut World) -> usize,
consume_command_and_get_size:
unsafe fn(value: OwningPtr<Unaligned>, world: &mut Option<&mut World>) -> usize,
}

/// Densely and efficiently stores a queue of heterogenous types implementing [`Command`].
Expand Down Expand Up @@ -53,11 +56,16 @@ impl CommandQueue {
}

let meta = CommandMeta {
apply_command_and_get_size: |command, world| {
// SAFETY: According to the invariants of `CommandMeta.apply_command_and_get_size`,
consume_command_and_get_size: |command, world| {
// SAFETY: According to the invariants of `CommandMeta.consume_command_and_get_size`,
// `command` must point to a value of type `C`.
let command: C = unsafe { command.read_unaligned() };
command.apply(world);
match world {
// Apply command to the provided world...
Some(world) => command.apply(world),
// ...or discard it.
None => drop(command),
}
std::mem::size_of::<C>()
},
};
Expand Down Expand Up @@ -97,6 +105,14 @@ impl CommandQueue {
// flush the previously queued entities
world.flush();

self.apply_or_drop_queued(Some(world));
}

/// If `world` is [`Some`], this will apply the queued [commands](`Command`).
/// If `world` is [`None`], this will drop the queued [commands](`Command`) (without applying them).
/// This clears the queue.
#[inline]
fn apply_or_drop_queued(&mut self, mut world: Option<&mut World>) {
// The range of pointers of the filled portion of `self.bytes`.
let bytes_range = self.bytes.as_mut_ptr_range();

Expand Down Expand Up @@ -127,7 +143,7 @@ impl CommandQueue {
// SAFETY: The data underneath the cursor must correspond to the type erased in metadata,
// since they were stored next to each other by `.push()`.
// For ZSTs, the type doesn't matter as long as the pointer is non-null.
let size = unsafe { (meta.apply_command_and_get_size)(cmd, world) };
let size = unsafe { (meta.consume_command_and_get_size)(cmd, &mut world) };
// Advance the cursor past the command. For ZSTs, the cursor will not move.
// At this point, it will either point to the next `CommandMeta`,
// or the cursor will be out of bounds and the loop will end.
Expand All @@ -138,6 +154,12 @@ impl CommandQueue {
}
}

impl Drop for CommandQueue {
fn drop(&mut self) {
self.apply_or_drop_queued(None);
}
}

#[cfg(test)]
mod test {
use super::*;
Expand Down Expand Up @@ -188,6 +210,27 @@ mod test {
assert_eq!(drops_b.load(Ordering::Relaxed), 1);
}

/// Asserts that inner [commands](`Command`) are dropped on early drop of [`CommandQueue`].
/// Originally identified as an issue in [#10676](https://github.com/bevyengine/bevy/issues/10676)
#[test]
fn test_command_queue_inner_drop_early() {
let mut queue = CommandQueue::default();

let (dropcheck_a, drops_a) = DropCheck::new();
let (dropcheck_b, drops_b) = DropCheck::new();

queue.push(dropcheck_a);
queue.push(dropcheck_b);

assert_eq!(drops_a.load(Ordering::Relaxed), 0);
assert_eq!(drops_b.load(Ordering::Relaxed), 0);

drop(queue);

assert_eq!(drops_a.load(Ordering::Relaxed), 1);
assert_eq!(drops_b.load(Ordering::Relaxed), 1);
}

struct SpawnCommand;

impl Command for SpawnCommand {
Expand Down

0 comments on commit 86d0b8a

Please sign in to comment.