Create diagnostics for material allocations

crates/bevy_pbr/src/diagnostic.rs

@@ -0,0 +1,101 @@
+use core::{any::type_name, marker::PhantomData};
+
+use bevy_app::{Plugin, PreUpdate};
+use bevy_diagnostic::{Diagnostic, DiagnosticPath, Diagnostics, RegisterDiagnostic};
+use bevy_ecs::{resource::Resource, system::Res};
+use bevy_platform::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use bevy_render::{Extract, ExtractSchedule, RenderApp};
+
+use crate::{Material, MaterialBindGroupAllocator};
+
+#[derive(Default)]
+pub struct MaterialAllocatorDiagnosticPlugin<M: Material> {
+    _phantom: PhantomData<M>,
+}
+
+impl<M: Material> MaterialAllocatorDiagnosticPlugin<M> {
+    /// Get the [`DiagnosticPath`] for slab count
+    pub fn slabs_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_slabs", type_name::<M>()])
+    }
+    /// Get the [`DiagnosticPath`] for total slabs size
+    pub fn slabs_size_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_slabs_size", type_name::<M>()])
+    }
+    /// Get the [`DiagnosticPath`] for material allocations
+    pub fn allocations_diagnostic_path() -> DiagnosticPath {
+        DiagnosticPath::from_components(["material_allocator_allocations", type_name::<M>()])
+    }
+}
+
+impl<M: Material> Plugin for MaterialAllocatorDiagnosticPlugin<M> {
+    fn build(&self, app: &mut bevy_app::App) {
+        app.register_diagnostic(
+            Diagnostic::new(Self::slabs_diagnostic_path()).with_suffix(" slabs"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(Self::slabs_size_diagnostic_path()).with_suffix(" bytes"),
+        )
+        .register_diagnostic(
+            Diagnostic::new(Self::allocations_diagnostic_path()).with_suffix(" meshes"),
+        )
+        .init_resource::<MaterialAllocatorMeasurements<M>>()
+        .add_systems(PreUpdate, add_material_allocator_measurement::<M>);
+
+        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
+            render_app.add_systems(ExtractSchedule, measure_allocator::<M>);
+        }
+    }
+}
+
+#[derive(Debug, Resource)]
+struct MaterialAllocatorMeasurements<M: Material> {
+    slabs: AtomicUsize,
+    slabs_size: AtomicUsize,
+    allocations: AtomicU64,
+    _phantom: PhantomData<M>,
+}
+
+impl<M: Material> Default for MaterialAllocatorMeasurements<M> {
+    fn default() -> Self {
+        Self {
+            slabs: AtomicUsize::default(),
+            slabs_size: AtomicUsize::default(),
+            allocations: AtomicU64::default(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+fn add_material_allocator_measurement<M: Material>(
+    mut diagnostics: Diagnostics,
+    measurements: Res<MaterialAllocatorMeasurements<M>>,
+) {
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::slabs_diagnostic_path(),
+        || measurements.slabs.load(Ordering::Relaxed) as f64,
+    );
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::slabs_size_diagnostic_path(),
+        || measurements.slabs_size.load(Ordering::Relaxed) as f64,
+    );
+    diagnostics.add_measurement(
+        &MaterialAllocatorDiagnosticPlugin::<M>::allocations_diagnostic_path(),
+        || measurements.allocations.load(Ordering::Relaxed) as f64,
+    );
+}
+
+fn measure_allocator<M: Material>(
+    measurements: Extract<Res<MaterialAllocatorMeasurements<M>>>,
+    allocator: Res<MaterialBindGroupAllocator<M>>,
+) {
+    measurements
+        .slabs
+        .store(allocator.slab_count(), Ordering::Relaxed);
+    measurements
+        .slabs_size
+        .store(allocator.slabs_size(), Ordering::Relaxed);
+    measurements
+        .allocations
+        .store(allocator.allocations(), Ordering::Relaxed);
+}

crates/bevy_pbr/src/lib.rs

@@ -29,6 +29,7 @@ mod cluster;
 mod components;
 pub mod decal;
 pub mod deferred;
+pub mod diagnostic;
 mod extended_material;
 mod fog;
 mod light;

crates/bevy_pbr/src/material_bind_groups.rs

@@ -611,6 +611,45 @@ where
             }
         }
     }
+
+    /// Get number of allocated slabs for bindless material, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn slab_count(&self) -> usize {
+        match self {
+            Self::Bindless(bless) => bless.slabs.len(),
+            Self::NonBindless(_) => 0,
+        }
+    }
+
+    /// Get total size of slabs allocated for bindless material, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn slabs_size(&self) -> usize {
+        match self {
+            Self::Bindless(bless) => bless
+                .slabs
+                .iter()
+                .flat_map(|slab| {
+                    slab.data_buffers
+                        .iter()
+                        .map(|(_, buffer)| buffer.buffer.len())
+                })
+                .sum(),
+            Self::NonBindless(_) => 0,
+        }
+    }
+
+    /// Get number of bindless material allocations in slabs, returns 0 if it is
+    /// [`Self::NonBindless`].
+    pub fn allocations(&self) -> u64 {
+        match self {
+            Self::Bindless(bless) => bless
+                .slabs
+                .iter()
+                .map(|slab| u64::from(slab.allocated_resource_count))
+                .sum(),
+            Self::NonBindless(_) => 0,
+        }
+    }
 }
 
 impl<M> MaterialBindlessIndexTable<M>

tests/get_mut_leak.rs

@@ -11,7 +11,10 @@ use bevy::{
     diagnostic::{DiagnosticsStore, LogDiagnosticsPlugin},
     ecs::system::{Commands, Local, Res, ResMut},
     math::primitives::Sphere,
-    pbr::{MeshMaterial3d, StandardMaterial},
+    pbr::{
+        diagnostic::MaterialAllocatorDiagnosticPlugin, Material, MeshMaterial3d, PreparedMaterial,
+        StandardMaterial,
+    },
     render::{
         diagnostic::{MeshAllocatorDiagnosticPlugin, RenderAssetDiagnosticPlugin},
         mesh::{Mesh, Mesh3d, Meshable, RenderMesh},
@@ -51,6 +54,39 @@ fn check_mesh_leak() {
     }
 }
 
+#[test]
+fn check_standard_material_leak() {
+    let mut app = App::new();
+    app.add_plugins((
+        DefaultPlugins
+            .build()
+            .disable::<AudioPlugin>()
+            .disable::<WinitPlugin>()
+            .disable::<WindowPlugin>(),
+        LogDiagnosticsPlugin {
+            wait_duration: Duration::ZERO,
+            ..Default::default()
+        },
+        RenderAssetDiagnosticPlugin::<PreparedMaterial<StandardMaterial>>::new(" materials"),
+        MaterialAllocatorDiagnosticPlugin::<StandardMaterial>::default(),
+    ))
+    .add_systems(Startup, mesh_setup)
+    .add_systems(
+        Update,
+        (
+            touch_mutably::<Mesh>,
+            crash_on_material_leak_detection::<StandardMaterial>,
+        ),
+    );
+
+    app.finish();
+    app.cleanup();
+
+    for _ in 0..100 {
+        app.update();
+    }
+}
+
 fn mesh_setup(
     mut commands: Commands,
     mut meshes: ResMut<Assets<Mesh>>,
@@ -93,3 +129,18 @@ fn crash_on_mesh_leak_detection(diagnostic_store: Res<DiagnosticsStore>) {
         );
     }
 }
+
+fn crash_on_material_leak_detection<M: Material>(diagnostic_store: Res<DiagnosticsStore>) {
+    if let (Some(materials), Some(allocations)) = (
+        diagnostic_store
+            .get_measurement(
+                &RenderAssetDiagnosticPlugin::<PreparedMaterial<M>>::render_asset_diagnostic_path(),
+            )
+            .filter(|diag| diag.value > 0.),
+        diagnostic_store
+            .get_measurement(&MaterialAllocatorDiagnosticPlugin::<M>::allocations_diagnostic_path())
+            .filter(|diag| diag.value > 0.),
+    ) {
+        assert!(materials.value < allocations.value * 10., "Detected leak");
+    }
+}