Skip to content

Commit 6e871ab

Browse files
bushrat011899cart
andauthored
Implement Drop for CommandQueue (#10746)
# 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]>
1 parent fa903a1 commit 6e871ab

File tree

1 file changed

+48
-5
lines changed

1 file changed

+48
-5
lines changed

crates/bevy_ecs/src/system/commands/command_queue.rs

Lines changed: 48 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,11 @@ struct CommandMeta {
99
/// SAFETY: The `value` must point to a value of type `T: Command`,
1010
/// where `T` is some specific type that was used to produce this metadata.
1111
///
12+
/// `world` is optional to allow this one function pointer to perform double-duty as a drop.
13+
///
1214
/// Returns the size of `T` in bytes.
13-
apply_command_and_get_size: unsafe fn(value: OwningPtr<Unaligned>, world: &mut World) -> usize,
15+
consume_command_and_get_size:
16+
unsafe fn(value: OwningPtr<Unaligned>, world: &mut Option<&mut World>) -> usize,
1417
}
1518

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

5558
let meta = CommandMeta {
56-
apply_command_and_get_size: |command, world| {
57-
// SAFETY: According to the invariants of `CommandMeta.apply_command_and_get_size`,
59+
consume_command_and_get_size: |command, world| {
60+
// SAFETY: According to the invariants of `CommandMeta.consume_command_and_get_size`,
5861
// `command` must point to a value of type `C`.
5962
let command: C = unsafe { command.read_unaligned() };
60-
command.apply(world);
63+
match world {
64+
// Apply command to the provided world...
65+
Some(world) => command.apply(world),
66+
// ...or discard it.
67+
None => drop(command),
68+
}
6169
std::mem::size_of::<C>()
6270
},
6371
};
@@ -97,6 +105,14 @@ impl CommandQueue {
97105
// flush the previously queued entities
98106
world.flush();
99107

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

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

162+
impl Drop for CommandQueue {
163+
fn drop(&mut self) {
164+
self.apply_or_drop_queued(None);
165+
}
166+
}
167+
146168
#[cfg(test)]
147169
mod test {
148170
use super::*;
@@ -193,6 +215,27 @@ mod test {
193215
assert_eq!(drops_b.load(Ordering::Relaxed), 1);
194216
}
195217

218+
/// Asserts that inner [commands](`Command`) are dropped on early drop of [`CommandQueue`].
219+
/// Originally identified as an issue in [#10676](https://github.com/bevyengine/bevy/issues/10676)
220+
#[test]
221+
fn test_command_queue_inner_drop_early() {
222+
let mut queue = CommandQueue::default();
223+
224+
let (dropcheck_a, drops_a) = DropCheck::new();
225+
let (dropcheck_b, drops_b) = DropCheck::new();
226+
227+
queue.push(dropcheck_a);
228+
queue.push(dropcheck_b);
229+
230+
assert_eq!(drops_a.load(Ordering::Relaxed), 0);
231+
assert_eq!(drops_b.load(Ordering::Relaxed), 0);
232+
233+
drop(queue);
234+
235+
assert_eq!(drops_a.load(Ordering::Relaxed), 1);
236+
assert_eq!(drops_b.load(Ordering::Relaxed), 1);
237+
}
238+
196239
struct SpawnCommand;
197240

198241
impl Command for SpawnCommand {

0 commit comments

Comments
 (0)