From 73547471cecf54376382edc1cf146666ba02b7ea Mon Sep 17 00:00:00 2001 From: JMS55 <47158642+JMS55@users.noreply.github.com> Date: Sat, 1 Aug 2026 23:01:28 -0400 Subject: [PATCH 1/2] Fix mesh re-allocation logic --- crates/bevy_render/Cargo.toml | 2 + crates/bevy_render/src/lib.rs | 2 + crates/bevy_render/src/mesh/allocator.rs | 181 +++++++++++++++++++- crates/bevy_render/src/slab_allocator.rs | 207 +++++++++++++++++++++-- crates/bevy_render/src/test_utils.rs | 43 +++++ 5 files changed, 422 insertions(+), 13 deletions(-) create mode 100644 crates/bevy_render/src/test_utils.rs diff --git a/crates/bevy_render/Cargo.toml b/crates/bevy_render/Cargo.toml index ab3228c638355..183e320a67a46 100644 --- a/crates/bevy_render/Cargo.toml +++ b/crates/bevy_render/Cargo.toml @@ -134,6 +134,8 @@ tracing = "0.1" [dev-dependencies] proptest = "1" proptest-derive = "0.8" +# noop backend used for tests via test_utils.rs +wgpu = { version = "30", default-features = false, features = ["noop"] } [target.'cfg(all(target_arch = "wasm32", target_feature = "atomics"))'.dependencies] send_wrapper = { version = "0.6.0" } diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs index 48c62ea0e80ca..3b7b22af3fced 100644 --- a/crates/bevy_render/src/lib.rs +++ b/crates/bevy_render/src/lib.rs @@ -63,6 +63,8 @@ pub mod slab_allocator; pub mod storage; pub mod sync_component; pub mod sync_world; +#[cfg(test)] +pub(crate) mod test_utils; pub mod texture; pub mod uniform; pub mod view; diff --git a/crates/bevy_render/src/mesh/allocator.rs b/crates/bevy_render/src/mesh/allocator.rs index b868f91e53c43..26a113f0fec03 100644 --- a/crates/bevy_render/src/mesh/allocator.rs +++ b/crates/bevy_render/src/mesh/allocator.rs @@ -624,16 +624,24 @@ impl MeshAllocator { ); } - /// Frees allocations for meshes that were removed or modified this frame. + /// Frees allocations for meshes that were removed, modified, or re-extracted + /// this frame. fn free_meshes(&mut self, extracted_meshes: &ExtractedAssets) { let mut deallocation_stage = self.slab_allocator.stage_deallocation(); // TODO: Consider explicitly reusing allocations for changed meshes of // the same size + + // Free every mesh that `allocate_meshes` is about to reallocate. Despite + // its name, `added` holds every mesh extracted this frame rather than only + // the new ones, so it's exactly that set. This catches a mesh + // removed from `Assets` and reinserted under the same ID, which arrives as + // `Removed` then `Added` and so never appears in `modified`. let meshes_to_free = extracted_meshes .removed .iter() - .chain(extracted_meshes.modified.iter()); + .chain(extracted_meshes.modified.iter()) + .chain(extracted_meshes.added.iter()); for mesh_id in meshes_to_free { deallocation_stage.free(&MeshAllocationKey::new(*mesh_id, ElementClass::Metadata)); @@ -738,3 +746,172 @@ impl ElementClass { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::create_dummy_device; + use bevy_asset::{uuid::Uuid, RenderAssetUsages}; + use bevy_mesh::PrimitiveTopology; + + fn test_mesh() -> Mesh { + let mut mesh = Mesh::new( + PrimitiveTopology::TriangleList, + RenderAssetUsages::default(), + ); + mesh.insert_attribute(Mesh::ATTRIBUTE_POSITION, vec![[0.0f32, 0.0, 0.0]; 64]); + mesh + } + + /// Builds the extraction output for a mesh that was extracted this frame. + /// + /// This mirrors what `extract_render_asset` produces for an + /// `AssetEvent::Added`: the mesh lands in `extracted` and `added`, and in + /// neither `removed` nor `modified`. + fn extracted_mesh(id: AssetId, mesh: Mesh) -> ExtractedAssets { + let mut extracted_meshes = ExtractedAssets::::default(); + extracted_meshes.extracted.push((id, mesh)); + extracted_meshes.added.insert(id); + extracted_meshes + } + + /// Builds the extraction output for a mesh modified in place this frame. + /// + /// This mirrors what `extract_render_asset` produces for an + /// `AssetEvent::Modified`, as caused by `Assets::get_mut`: the mesh is + /// re-extracted, so it lands in `modified` on top of `extracted` and `added`. + fn modified_mesh(id: AssetId, mesh: Mesh) -> ExtractedAssets { + let mut extracted_meshes = extracted_mesh(id, mesh); + extracted_meshes.modified.insert(id); + extracted_meshes + } + + /// `free_meshes` must release meshes that are merely being re-extracted, not + /// only those flagged `removed` or `modified`. + #[test] + fn free_meshes_releases_reextracted_meshes() { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = MeshAllocator { + slab_allocator: SlabAllocator::new(), + general_vertex_slabs_supported: true, + }; + + let mesh_id = AssetId::::Uuid { + uuid: Uuid::from_u128(1), + }; + let extracted_meshes = extracted_mesh(mesh_id, test_mesh()); + + mesh_allocator.allocate_meshes( + &settings, + &extracted_meshes, + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + assert!(mesh_allocator.mesh_vertex_slice(&mesh_id).is_some()); + + // Being present in `added` alone must be enough to release the previous + // allocation. + mesh_allocator.free_meshes(&extracted_meshes); + + assert!( + mesh_allocator.key_to_slab.is_empty(), + "a re-extracted mesh was not freed, so its old allocation would leak" + ); + assert_eq!(mesh_allocator.slab_count(), 0); + } + + /// A mesh flagged `modified` must be freed even when it isn't re-extracted. + /// + /// `added` covers meshes that come back around for reallocation, but a mesh + /// can be modified and then leave `Assets` without emitting `Unused`, in + /// which case `modified` is the only record we get of it. + #[test] + fn free_meshes_releases_modified_meshes_that_were_not_reextracted() { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = MeshAllocator { + slab_allocator: SlabAllocator::new(), + general_vertex_slabs_supported: true, + }; + + let mesh_id = AssetId::::Uuid { + uuid: Uuid::from_u128(1), + }; + mesh_allocator.allocate_meshes( + &settings, + &extracted_mesh(mesh_id, test_mesh()), + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + assert!(mesh_allocator.mesh_vertex_slice(&mesh_id).is_some()); + + let mut extracted_meshes = ExtractedAssets::::default(); + extracted_meshes.modified.insert(mesh_id); + mesh_allocator.free_meshes(&extracted_meshes); + + assert!( + mesh_allocator.key_to_slab.is_empty(), + "a modified mesh that wasn't re-extracted was not freed" + ); + assert_eq!(mesh_allocator.slab_count(), 0); + } + + /// Runs `rounds` frames of free-then-allocate over a single mesh ID and + /// asserts that slab memory reaches a steady state. + fn assert_steady_state( + build: fn(AssetId, Mesh) -> ExtractedAssets, + rounds: usize, + ) { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = MeshAllocator { + slab_allocator: SlabAllocator::new(), + general_vertex_slabs_supported: true, + }; + + let mesh_id = AssetId::::Uuid { + uuid: Uuid::from_u128(1), + }; + + let mut baseline = None; + for _ in 0..rounds { + let extracted_meshes = build(mesh_id, test_mesh()); + mesh_allocator.free_meshes(&extracted_meshes); + mesh_allocator.allocate_meshes( + &settings, + &extracted_meshes, + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + + let size = mesh_allocator.slabs_size(); + match baseline { + None => baseline = Some(size), + Some(baseline) => { + assert_eq!(size, baseline, "slab memory grew across frames"); + } + } + } + + assert!(mesh_allocator.mesh_vertex_slice(&mesh_id).is_some()); + } + + /// Re-extracting the same mesh ID every frame must reach a steady state. + #[test] + fn reextracting_the_same_mesh_does_not_grow_the_slabs() { + assert_steady_state(extracted_mesh, 32); + } + + /// Modifying the same mesh in place every frame must reach a steady state. + #[test] + fn modifying_a_mesh_in_place_does_not_grow_the_slabs() { + assert_steady_state(modified_mesh, 32); + } +} diff --git a/crates/bevy_render/src/slab_allocator.rs b/crates/bevy_render/src/slab_allocator.rs index 13684939f97a5..118537dae534b 100644 --- a/crates/bevy_render/src/slab_allocator.rs +++ b/crates/bevy_render/src/slab_allocator.rs @@ -407,6 +407,9 @@ where pub allocator: &'a mut SlabAllocator, /// The set of slabs that have grown and need to be reallocated. slabs_to_reallocate: HashMap, SlabToReallocate>, + /// IDs of slabs that became empty because everything in them was + /// reallocated elsewhere. + empty_slabs: HashSet>, } impl<'a, I> Drop for AllocationStage<'a, I> @@ -414,10 +417,10 @@ where I: SlabItem, { fn drop(&mut self) { - if !self.slabs_to_reallocate.is_empty() { + if !self.slabs_to_reallocate.is_empty() || !self.empty_slabs.is_empty() { error!( - "Dropping an `AllocationStage` with uncommitted reallocations. You should call \ - `AllocationStage::commit`." + "Dropping an `AllocationStage` with uncommitted reallocations or slab free \ + operations. You should call `AllocationStage::commit`." ); } } @@ -429,7 +432,10 @@ where { /// Allocates space for an object of the given size with the given key and layout. /// - /// The key must not correspond to any current allocation. + /// If the key already corresponds to a live allocation, that allocation is + /// freed first. Prefer freeing it through a [`DeallocationStage`] before the + /// allocation stage begins, so that the allocator has every one of the + /// frame's holes to choose from rather than just this object's. pub fn allocate( &mut self, key: &I::Key, @@ -437,6 +443,8 @@ where layout: I::Layout, settings: &SlabAllocatorSettings, ) { + self.allocator + .free_existing_allocation(key, &mut self.empty_slabs); self.allocator.allocate( key, data_byte_len, @@ -448,14 +456,26 @@ where /// Allocates an object into its own dedicated slab. /// - /// The key must not correspond to any current allocation. + /// As with [`Self::allocate`], a live allocation under the same key is freed + /// first. pub fn allocate_large(&mut self, key: &I::Key, layout: I::Layout) { + self.allocator + .free_existing_allocation(key, &mut self.empty_slabs); self.allocator.allocate_large(key, layout); } /// Completes the transaction, performing any queued resize operations. pub fn commit(mut self, render_device: &RenderDevice, render_queue: &RenderQueue) { + // Drop slabs that were emptied by their contents being reallocated + // elsewhere. Do this before growing anything, so that we never create a + // buffer for a slab we're about to throw away. + self.allocator.free_empty_slabs(self.empty_slabs.drain()); + for (slab_id, slab_to_grow) in self.slabs_to_reallocate.drain() { + // The slab may have been freed just above. + if !self.allocator.slabs.contains_key(&slab_id) { + continue; + } self.allocator .reallocate_slab(render_device, render_queue, slab_id, slab_to_grow); } @@ -499,13 +519,11 @@ where { /// Schedules a free operation for the allocation with the given key. /// - /// The key must correspond to a live allocation. An error will be emitted - /// to the log otherwise. + /// Freeing a key that holds no allocation is a no-op, so callers are free to + /// speculatively free keys that may never have been allocated. pub fn free(&mut self, key: &I::Key) { - if let Some(slab_id) = self.allocator.key_to_slab.remove(key) { - self.allocator - .free_allocation_in_slab(key, slab_id, &mut self.empty_slabs); - } + self.allocator + .free_existing_allocation(key, &mut self.empty_slabs); } /// Performs all the free operations. @@ -624,6 +642,7 @@ where AllocationStage { allocator: self, slabs_to_reallocate: HashMap::default(), + empty_slabs: HashSet::default(), } } @@ -778,6 +797,16 @@ where ); } + /// Frees whatever allocation the given key currently holds, if any. + /// + /// If this empties the allocation's slab, that slab is added to the + /// `empty_slabs` set for the caller to reclaim on commit. + fn free_existing_allocation(&mut self, key: &I::Key, empty_slabs: &mut HashSet>) { + if let Some(slab_id) = self.key_to_slab.remove(key) { + self.free_allocation_in_slab(key, slab_id, empty_slabs); + } + } + /// Given a slab and the key corresponding to an object within it, marks /// the allocation as free. /// @@ -928,6 +957,15 @@ where fn free_empty_slabs(&mut self, empty_slabs: impl Iterator>) { for empty_slab in empty_slabs { + // The slab may have been refilled since it was marked empty, because + // a reallocated object usually lands back in the hole it just left. + // Destroying the slab here would take live data with it, so skip. + if let Some(Slab::General(general_slab)) = self.slabs.get(&empty_slab) + && !general_slab.is_empty() + { + continue; + } + self.slab_layouts.values_mut().for_each(|slab_ids| { let idx = slab_ids.iter().position(|&slab_id| slab_id == empty_slab); if let Some(idx) = idx { @@ -1141,3 +1179,150 @@ fn buffer_usages_to_str(buffer_usages: BufferUsages) -> &'static str { "" } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::create_dummy_device; + + /// A [`SlabItem`] for tests, keyed by a plain integer. + struct TestItem; + + impl SlabItem for TestItem { + type Key = u32; + type Layout = TestLayout; + + fn label() -> Cow<'static, str> { + "test".into() + } + } + + /// A four-byte element, one element per slot. + #[derive(Clone, PartialEq, Eq, Hash)] + struct TestLayout; + + impl SlabItemLayout for TestLayout { + fn size(&self) -> u64 { + 4 + } + + fn elements_per_slot(&self) -> u32 { + 1 + } + + fn buffer_usages(&self) -> BufferUsages { + BufferUsages::VERTEX + } + } + + /// Small slabs, so that a leak shows up as slab growth within a few rounds. + fn test_settings() -> SlabAllocatorSettings { + SlabAllocatorSettings { + // 256 slots. + min_slab_size: 1024, + // 1024 slots. + max_slab_size: 4096, + large_threshold: 4096, + growth_factor: 1.5, + } + } + + /// Allocates `byte_len` bytes under each of `keys` and marks the results + /// resident, as a frame of [`MeshAllocator`](crate::mesh::MeshAllocator) + /// would. + fn allocate_round( + allocator: &mut SlabAllocator, + keys: impl Iterator + Clone, + byte_len: u64, + device: &RenderDevice, + queue: &RenderQueue, + ) { + let mut stage = allocator.stage_allocation(); + for key in keys.clone() { + stage.allocate(&key, byte_len, TestLayout, &test_settings()); + } + stage.commit(device, queue); + + for key in keys { + allocator.copy_element_data(&key, byte_len as usize, |_| {}, device, queue); + } + } + + /// Reallocating a key that's still live must free its previous allocation + /// rather than orphaning it. + #[test] + fn reallocating_a_live_key_frees_the_old_allocation() { + let (device, queue) = create_dummy_device(); + let mut allocator = SlabAllocator::::new(); + + allocate_round(&mut allocator, 0..8, 256, &device, &queue); + + let baseline_size = allocator.slabs_size(); + let baseline_slabs = allocator.slab_count(); + assert!(baseline_size > 0, "nothing was allocated"); + + // Reallocate the same keys, without ever freeing them, many times over. + for _ in 0..32 { + allocate_round(&mut allocator, 0..8, 256, &device, &queue); + } + + assert_eq!( + allocator.slabs_size(), + baseline_size, + "reallocating live keys grew the slabs, so old allocations leaked" + ); + assert_eq!(allocator.slab_count(), baseline_slabs); + assert_eq!(allocator.key_to_slab.len(), 8); + + // Every key must still be readable. + for key in 0..8u32 { + let slab_id = allocator.key_to_slab[&key]; + assert!(allocator.slab_allocation_slice(&key, slab_id).is_some()); + } + } + + /// A slab emptied by a reallocation that lands back in that same slab must + /// not be destroyed on commit. + #[test] + fn slab_refilled_during_allocation_is_not_freed() { + let (device, queue) = create_dummy_device(); + let mut allocator = SlabAllocator::::new(); + + allocate_round(&mut allocator, 0..1, 256, &device, &queue); + let slab_id = allocator.key_to_slab[&0]; + + // Freeing the sole occupant marks the slab empty mid-stage; the + // reallocation then drops straight back into it. + allocate_round(&mut allocator, 0..1, 256, &device, &queue); + + assert_eq!(allocator.slab_count(), 1, "the live slab was destroyed"); + assert_eq!(allocator.key_to_slab[&0], slab_id); + assert!(allocator.slab_allocation_slice(&0, slab_id).is_some()); + } + + /// A slab genuinely emptied during an allocation stage must be reclaimed, or + /// the fix for reallocation would just trade one leak for another. + #[test] + fn slab_emptied_during_allocation_is_freed() { + let (device, queue) = create_dummy_device(); + let mut allocator = SlabAllocator::::new(); + + allocate_round(&mut allocator, 0..1, 256, &device, &queue); + assert_eq!(allocator.slab_count(), 1); + + // Reallocate the sole occupant at a size that forces it into a slab of + // its own, leaving the general slab empty. + allocate_round(&mut allocator, 0..1, 8192, &device, &queue); + + assert_eq!( + allocator.slab_count(), + 1, + "the emptied general slab was not reclaimed" + ); + let slab_id = allocator.key_to_slab[&0]; + assert!(matches!( + allocator.slabs.get(&slab_id), + Some(Slab::LargeObject(_)) + )); + } +} diff --git a/crates/bevy_render/src/test_utils.rs b/crates/bevy_render/src/test_utils.rs new file mode 100644 index 0000000000000..a002ca61ed71e --- /dev/null +++ b/crates/bevy_render/src/test_utils.rs @@ -0,0 +1,43 @@ +//! Helpers for this crate's unit tests. + +use alloc::sync::Arc; +use bevy_platform::future::block_on; +use wgpu::{ + BackendOptions, Backends, DeviceDescriptor, Instance, InstanceDescriptor, InstanceFlags, + NoopBackendOptions, RequestAdapterOptions, +}; + +use crate::renderer::{RenderDevice, RenderQueue, WgpuWrapper}; + +/// Creates a dummy [`RenderDevice`] and [`RenderQueue`] on `wgpu`'s noop backend. +/// +/// This lets tests exercise real `wgpu` resource creation without requiring a +/// GPU adapter, so they can run in headless environments. +pub fn create_dummy_device() -> (RenderDevice, RenderQueue) { + let instance = Instance::new(InstanceDescriptor { + backends: Backends::NOOP, + flags: InstanceFlags::default(), + memory_budget_thresholds: Default::default(), + display: None, + backend_options: BackendOptions { + noop: NoopBackendOptions { + enable: true, + ..Default::default() + }, + ..Default::default() + }, + }); + + let adapter = block_on(instance.request_adapter(&RequestAdapterOptions::default())) + .expect("the noop backend should always produce an adapter"); + let (device, queue) = block_on(adapter.request_device(&DeviceDescriptor { + required_limits: adapter.limits(), + ..Default::default() + })) + .expect("the noop backend should always produce a device"); + + ( + RenderDevice::from(device), + RenderQueue(Arc::new(WgpuWrapper::new(queue))), + ) +} From 465a87f584f5bfe14f1c1e83bacac3e2f65fd9e7 Mon Sep 17 00:00:00 2001 From: JMS55 <47158642+JMS55@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:36:26 -0400 Subject: [PATCH 2/2] Add more tests --- crates/bevy_render/src/mesh/allocator.rs | 337 ++++++++++++++++++++--- crates/bevy_render/src/slab_allocator.rs | 38 +++ 2 files changed, 344 insertions(+), 31 deletions(-) diff --git a/crates/bevy_render/src/mesh/allocator.rs b/crates/bevy_render/src/mesh/allocator.rs index 26a113f0fec03..2e1e4ffad43a5 100644 --- a/crates/bevy_render/src/mesh/allocator.rs +++ b/crates/bevy_render/src/mesh/allocator.rs @@ -752,7 +752,9 @@ mod tests { use super::*; use crate::test_utils::create_dummy_device; use bevy_asset::{uuid::Uuid, RenderAssetUsages}; + use bevy_math::bounding::Aabb3d; use bevy_mesh::PrimitiveTopology; + use glam::{Vec2, Vec3}; fn test_mesh() -> Mesh { let mut mesh = Mesh::new( @@ -763,6 +765,68 @@ mod tests { mesh } + /// A mesh that exercises every [`ElementClass`] at once. + fn full_mesh() -> Mesh { + let mut mesh = test_mesh(); + mesh.insert_indices(Indices::U32((0..64).collect())); + mesh.final_aabb = Some(Aabb3d::new(Vec3::ZERO, Vec3::ONE)); + mesh.final_uv_ranges[0] = Some(Aabb2d::new(Vec2::ZERO, Vec2::ONE)); + #[cfg(feature = "morph")] + mesh.set_morph_targets(vec![MorphAttributes::default(); 64]); + mesh + } + + /// A mesh whose vertex layout differs from [`test_mesh`], so that it sorts + /// into a different general slab. + fn wider_vertex_mesh() -> Mesh { + let mut mesh = test_mesh(); + mesh.insert_attribute(Mesh::ATTRIBUTE_NORMAL, vec![[0.0f32, 0.0, 1.0]; 64]); + mesh + } + + /// Builds a [`MeshAllocator`] backed by an empty slab allocator. + /// + /// Clearing `general_vertex_slabs_supported` sends every vertex array into a + /// slab of its own. + fn mesh_allocator(general_vertex_slabs_supported: bool) -> MeshAllocator { + MeshAllocator { + slab_allocator: SlabAllocator::new(), + general_vertex_slabs_supported, + } + } + + /// Allocator tuning with a small large-object threshold, so that an + /// otherwise modest test mesh is big enough to demand a slab of its own. + fn small_slab_settings() -> MeshAllocatorSettings { + MeshAllocatorSettings { + slab_allocator_settings: SlabAllocatorSettings { + min_slab_size: 1024, + max_slab_size: 4096, + large_threshold: 512, + growth_factor: 1.5, + }, + extra_buffer_usages: BufferUsages::empty(), + } + } + + fn mesh_id(id: u128) -> AssetId { + AssetId::::Uuid { + uuid: Uuid::from_u128(id), + } + } + + /// Whether the allocator currently holds an allocation of the given class + /// for the given mesh. + fn has_allocation( + mesh_allocator: &MeshAllocator, + mesh_id: AssetId, + class: ElementClass, + ) -> bool { + mesh_allocator + .key_to_slab + .contains_key(&MeshAllocationKey::new(mesh_id, class)) + } + /// Builds the extraction output for a mesh that was extracted this frame. /// /// This mirrors what `extract_render_asset` produces for an @@ -793,14 +857,9 @@ mod tests { let (render_device, render_queue) = create_dummy_device(); let settings = MeshAllocatorSettings::default(); let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); - let mut mesh_allocator = MeshAllocator { - slab_allocator: SlabAllocator::new(), - general_vertex_slabs_supported: true, - }; + let mut mesh_allocator = mesh_allocator(true); - let mesh_id = AssetId::::Uuid { - uuid: Uuid::from_u128(1), - }; + let mesh_id = mesh_id(1); let extracted_meshes = extracted_mesh(mesh_id, test_mesh()); mesh_allocator.allocate_meshes( @@ -833,14 +892,9 @@ mod tests { let (render_device, render_queue) = create_dummy_device(); let settings = MeshAllocatorSettings::default(); let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); - let mut mesh_allocator = MeshAllocator { - slab_allocator: SlabAllocator::new(), - general_vertex_slabs_supported: true, - }; + let mut mesh_allocator = mesh_allocator(true); - let mesh_id = AssetId::::Uuid { - uuid: Uuid::from_u128(1), - }; + let mesh_id = mesh_id(1); mesh_allocator.allocate_meshes( &settings, &extracted_mesh(mesh_id, test_mesh()), @@ -861,27 +915,196 @@ mod tests { assert_eq!(mesh_allocator.slab_count(), 0); } + /// `free_meshes` must release every class of allocation a mesh holds, not + /// just its vertex data. + #[test] + fn free_meshes_releases_every_element_class() { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = mesh_allocator(true); + + let mesh_id = mesh_id(1); + let extracted_meshes = extracted_mesh(mesh_id, full_mesh()); + mesh_allocator.allocate_meshes( + &settings, + &extracted_meshes, + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::Vertex + )); + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::Index + )); + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::Metadata + )); + #[cfg(feature = "morph")] + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::MorphTarget + )); + + mesh_allocator.free_meshes(&extracted_meshes); + + assert!( + mesh_allocator.key_to_slab.is_empty(), + "at least one element class was left allocated, so it would leak" + ); + assert_eq!(mesh_allocator.slab_count(), 0); + } + + /// A mesh that loses part of its data on re-extraction must give up the + /// matching allocations, which nothing will reallocate. + #[test] + fn reextracting_a_mesh_that_drops_its_extra_data_frees_those_allocations() { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = mesh_allocator(true); + + let mesh_id = mesh_id(1); + mesh_allocator.allocate_meshes( + &settings, + &extracted_mesh(mesh_id, full_mesh()), + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::Index + )); + + // The same ID comes back as a bare vertex-only mesh. + let extracted_meshes = extracted_mesh(mesh_id, test_mesh()); + mesh_allocator.free_meshes(&extracted_meshes); + mesh_allocator.allocate_meshes( + &settings, + &extracted_meshes, + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + + assert!(has_allocation( + &mesh_allocator, + mesh_id, + ElementClass::Vertex + )); + assert!( + !has_allocation(&mesh_allocator, mesh_id, ElementClass::Index), + "index data was dropped by the mesh but its allocation survived" + ); + assert!( + !has_allocation(&mesh_allocator, mesh_id, ElementClass::Metadata), + "metadata was dropped by the mesh but its allocation survived" + ); + #[cfg(feature = "morph")] + assert!( + !has_allocation(&mesh_allocator, mesh_id, ElementClass::MorphTarget), + "morph targets were dropped by the mesh but their allocation survived" + ); + } + + /// Changing a mesh's vertex layout moves it to a different slab, and the + /// slab it leaves behind must be reclaimed. + #[test] + fn reextracting_a_mesh_with_a_new_vertex_layout_reclaims_the_old_slab() { + let (render_device, render_queue) = create_dummy_device(); + let settings = MeshAllocatorSettings::default(); + let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); + let mut mesh_allocator = mesh_allocator(true); + + let mesh_id = mesh_id(1); + mesh_allocator.allocate_meshes( + &settings, + &extracted_mesh(mesh_id, test_mesh()), + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + assert_eq!(mesh_allocator.slab_count(), 1); + let original_slab = + mesh_allocator.key_to_slab[&MeshAllocationKey::new(mesh_id, ElementClass::Vertex)]; + + // Adding a normal attribute widens the vertex, which needs a slab with a + // different element layout. + let extracted_meshes = extracted_mesh(mesh_id, wider_vertex_mesh()); + mesh_allocator.free_meshes(&extracted_meshes); + mesh_allocator.allocate_meshes( + &settings, + &extracted_meshes, + &mut mesh_vertex_buffer_layouts, + &render_device, + &render_queue, + ); + + let new_slab = + mesh_allocator.key_to_slab[&MeshAllocationKey::new(mesh_id, ElementClass::Vertex)]; + assert_ne!( + new_slab, original_slab, + "the wider vertex should have landed in a slab with a different layout" + ); + assert_eq!( + mesh_allocator.slab_count(), + 1, + "the slab the mesh moved out of was not reclaimed" + ); + assert!(mesh_allocator.mesh_vertex_slice(&mesh_id).is_some()); + } + + /// One frame-loop scenario for [`assert_steady_state`]. + struct SteadyStateCase { + /// Whether the mesh arrives as merely re-extracted or as modified. + build: fn(AssetId, Mesh) -> ExtractedAssets, + build_mesh: fn() -> Mesh, + settings: MeshAllocatorSettings, + general_vertex_slabs_supported: bool, + } + + impl Default for SteadyStateCase { + fn default() -> Self { + Self { + build: extracted_mesh, + build_mesh: test_mesh, + settings: MeshAllocatorSettings::default(), + general_vertex_slabs_supported: true, + } + } + } + /// Runs `rounds` frames of free-then-allocate over a single mesh ID and /// asserts that slab memory reaches a steady state. - fn assert_steady_state( - build: fn(AssetId, Mesh) -> ExtractedAssets, - rounds: usize, - ) { + fn assert_steady_state(case: SteadyStateCase, rounds: usize) { + let SteadyStateCase { + build, + build_mesh, + settings, + general_vertex_slabs_supported, + } = case; + let (render_device, render_queue) = create_dummy_device(); - let settings = MeshAllocatorSettings::default(); let mut mesh_vertex_buffer_layouts = MeshVertexBufferLayouts::default(); - let mut mesh_allocator = MeshAllocator { - slab_allocator: SlabAllocator::new(), - general_vertex_slabs_supported: true, - }; + let mut mesh_allocator = mesh_allocator(general_vertex_slabs_supported); - let mesh_id = AssetId::::Uuid { - uuid: Uuid::from_u128(1), - }; + let mesh_id = mesh_id(1); let mut baseline = None; for _ in 0..rounds { - let extracted_meshes = build(mesh_id, test_mesh()); + let extracted_meshes = build(mesh_id, build_mesh()); mesh_allocator.free_meshes(&extracted_meshes); mesh_allocator.allocate_meshes( &settings, @@ -892,10 +1115,15 @@ mod tests { ); let size = mesh_allocator.slabs_size(); + let slab_count = mesh_allocator.slab_count(); match baseline { - None => baseline = Some(size), + None => baseline = Some((size, slab_count)), Some(baseline) => { - assert_eq!(size, baseline, "slab memory grew across frames"); + assert_eq!( + (size, slab_count), + baseline, + "slab memory grew across frames" + ); } } } @@ -906,12 +1134,59 @@ mod tests { /// Re-extracting the same mesh ID every frame must reach a steady state. #[test] fn reextracting_the_same_mesh_does_not_grow_the_slabs() { - assert_steady_state(extracted_mesh, 32); + assert_steady_state(SteadyStateCase::default(), 32); } /// Modifying the same mesh in place every frame must reach a steady state. #[test] fn modifying_a_mesh_in_place_does_not_grow_the_slabs() { - assert_steady_state(modified_mesh, 32); + assert_steady_state( + SteadyStateCase { + build: modified_mesh, + ..SteadyStateCase::default() + }, + 32, + ); + } + + /// A mesh carrying every element class must reach a steady state too, so + /// that a leak confined to one class cannot hide behind the vertex data. + #[test] + fn reextracting_a_mesh_with_every_element_class_does_not_grow_the_slabs() { + assert_steady_state( + SteadyStateCase { + build_mesh: full_mesh, + ..SteadyStateCase::default() + }, + 32, + ); + } + + /// Data too big to share a general slab gets one of its own, so a missed + /// free leaks a whole slab rather than a slot inside one. + #[test] + fn reextracting_a_mesh_too_large_for_a_general_slab_does_not_grow_the_slabs() { + assert_steady_state( + SteadyStateCase { + build_mesh: full_mesh, + settings: small_slab_settings(), + ..SteadyStateCase::default() + }, + 32, + ); + } + + /// The other route to a dedicated slab, taken when the platform cannot + /// share vertex slabs at all. + #[test] + fn reextracting_a_mesh_without_general_vertex_slabs_does_not_grow_the_slabs() { + assert_steady_state( + SteadyStateCase { + build_mesh: full_mesh, + general_vertex_slabs_supported: false, + ..SteadyStateCase::default() + }, + 32, + ); } } diff --git a/crates/bevy_render/src/slab_allocator.rs b/crates/bevy_render/src/slab_allocator.rs index 118537dae534b..bc1943fbd85b1 100644 --- a/crates/bevy_render/src/slab_allocator.rs +++ b/crates/bevy_render/src/slab_allocator.rs @@ -1300,6 +1300,44 @@ mod tests { assert!(allocator.slab_allocation_slice(&0, slab_id).is_some()); } + /// Allocates a dedicated slab for each of `keys` through + /// [`AllocationStage::allocate_large`]. + fn allocate_large_round( + allocator: &mut SlabAllocator, + keys: impl Iterator, + device: &RenderDevice, + queue: &RenderQueue, + ) { + let mut stage = allocator.stage_allocation(); + for key in keys { + stage.allocate_large(&key, TestLayout); + } + stage.commit(device, queue); + } + + /// [`AllocationStage::allocate_large`] must free a live allocation under the + /// same key, just as [`AllocationStage::allocate`] does. Each one takes a + /// brand new slab, so a missed free leaks a whole slab per round. + #[test] + fn reallocating_a_live_key_with_allocate_large_frees_the_old_allocation() { + let (device, queue) = create_dummy_device(); + let mut allocator = SlabAllocator::::new(); + + allocate_large_round(&mut allocator, 0..4, &device, &queue); + assert_eq!(allocator.slab_count(), 4); + + for _ in 0..32 { + allocate_large_round(&mut allocator, 0..4, &device, &queue); + } + + assert_eq!( + allocator.slab_count(), + 4, + "reallocating live keys left the previous dedicated slabs behind" + ); + assert_eq!(allocator.key_to_slab.len(), 4); + } + /// A slab genuinely emptied during an allocation stage must be reclaimed, or /// the fix for reallocation would just trade one leak for another. #[test]