Add support for tracking RenderAsset memory usage.

Author: BGR360

.github/workflows/tests.yml

@@ -1,3 +1,5 @@
+# TODO: build with all combinations of features.
+
 name: tests
 
 on:

Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 bevy = { version = "0.6", default-features = false }
 bytesize = "1"
 datasize = "0.2"
+parking_lot = "0.11"
 
 [dev-dependencies]
 maplit = "1"

examples/render_resources.rs

@@ -3,7 +3,10 @@
 //!
 //! Adapted from the official Bevy `many_cubes` example.
 
-use bevy::prelude::*;
+use bevy::{
+    prelude::*,
+    render::{mesh::GpuMesh, texture::GpuImage},
+};
 use bevy_datasize::prelude::*;
 
 fn main() {
@@ -48,12 +51,14 @@ fn setup(
 
 fn print_sizes(memory_usage: Res<MemoryUsage>) {
     let mesh_stats = memory_usage.get_stats::<Mesh>().unwrap();
+    let gpu_mesh_stats = memory_usage.get_stats::<GpuMesh>().unwrap();
     let image_stats = memory_usage.get_stats::<Image>().unwrap();
     let material_stats = memory_usage.get_stats::<StandardMaterial>().unwrap();
 
     println!();
     println!("Memory usage:");
     println!("Meshes: {mesh_stats}");
+    println!("GPU Meshes: {gpu_mesh_stats}");
     println!("Images: {image_stats}");
     println!("Materials: {material_stats}");
 }

src/app_ext.rs

@@ -1,15 +1,32 @@
+//! Extension traits for [`App`].
+
 use std::any::Any;
 
 use bevy::{
     app::App,
     asset::Asset,
-    ecs::{component::Component, schedule::IntoSystemDescriptor, system::Resource},
+    ecs::{
+        component::Component, schedule::IntoSystemDescriptor, schedule::StageLabel,
+        system::Resource,
+    },
+    prelude::CoreStage,
 };
 
-use crate::{estimator::ForwardingEstimator, systems, DataSize, MemoryUsage};
+#[cfg(feature = "bevy_render")]
+use bevy::render::{render_asset::RenderAsset, RenderApp, RenderStage};
+
+use crate::{
+    estimator::{ForwardingEstimator, FromConfig},
+    systems, DataSize, DataSizeEstimator, MemoryUsage,
+};
 
-/// [`App`] extension methods to register types for memory usage tracking.
-pub trait RegisterSizedTypes {
+/// [`App`] extension methods to register [`DataSize`] types for memory usage
+/// tracking.
+///
+/// For types that do not implement [`DataSize`], you will need to implement a
+/// [`DataSizeEstimator`] for them and register them using the methods in
+/// [`RegisterTypesWithEstimator`].
+pub trait RegisterSizedTypes: RegisterTypesWithEstimator {
     /// Registers the given [`Component`] type with the
     /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
     ///
@@ -19,9 +36,7 @@ pub trait RegisterSizedTypes {
     where
         T: Any + DataSize + Component,
     {
-        self.register_sized_type::<T, _, _>(
-            systems::update_stats_for_component::<T, ForwardingEstimator>,
-        )
+        self.register_component_with_estimator::<T, ForwardingEstimator>()
     }
 
     /// Registers the given [`Resource`] type with the
@@ -33,9 +48,7 @@ pub trait RegisterSizedTypes {
     where
         T: Any + DataSize + Resource,
     {
-        self.register_sized_type::<T, _, _>(
-            systems::update_stats_for_resource::<T, ForwardingEstimator>,
-        )
+        self.register_resource_with_estimator::<T, ForwardingEstimator>()
     }
 
     /// Registers the given [`Asset`] type with the
@@ -47,39 +60,145 @@ pub trait RegisterSizedTypes {
     where
         T: Any + DataSize + Asset,
     {
-        self.register_sized_type::<T, _, _>(
-            systems::update_stats_for_asset::<T, ForwardingEstimator>,
+        self.register_asset_with_estimator::<T, ForwardingEstimator>()
+    }
+
+    /// Registers the given [`Asset`] type with the
+    /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
+    ///
+    /// The following types will be available to query on the [`MemoryUsage`]
+    /// resource:
+    ///
+    /// * `T`
+    /// * `<T as RenderAsset>::ExtractedAsset`
+    /// * `<T as RenderAsset>::PreparedAsset`
+    #[cfg(feature = "bevy_render")]
+    fn register_sized_render_asset<T>(&mut self) -> &mut Self
+    where
+        T: Any + DataSize + RenderAsset,
+        <T as RenderAsset>::PreparedAsset: Any + DataSize,
+    {
+        self.register_render_asset_with_estimator::<T, ForwardingEstimator, ForwardingEstimator>()
+    }
+}
+
+impl RegisterSizedTypes for App {}
+
+/// [`App`] extension methods to register non-[`DataSize`] types for memory
+/// usage tracking.
+pub trait RegisterTypesWithEstimator: RegisterTypes {
+    /// Like [`RegisterSizedTypes::register_sized_component`], but uses the
+    /// given [`DataSizeEstimator`] type.
+    fn register_component_with_estimator<T, E>(&mut self) -> &mut Self
+    where
+        T: Any + Component,
+        E: DataSizeEstimator<T> + FromConfig + 'static,
+    {
+        self.register_type::<T, _, _, _>(
+            systems::update_stats_for_component::<T, E>,
+            CoreStage::Update,
         )
     }
 
+    /// Like [`RegisterSizedTypes::register_sized_resource`], but uses the
+    /// given [`DataSizeEstimator`] type.
+    fn register_resource_with_estimator<T, E>(&mut self) -> &mut Self
+    where
+        T: Any + Resource,
+        E: DataSizeEstimator<T> + FromConfig + 'static,
+    {
+        self.register_type::<T, _, _, _>(
+            systems::update_stats_for_resource::<T, E>,
+            CoreStage::Update,
+        )
+    }
+
+    /// Like [`RegisterSizedTypes::register_sized_asset`], but uses the
+    /// given [`DataSizeEstimator`] type.
+    fn register_asset_with_estimator<T, E>(&mut self) -> &mut Self
+    where
+        T: Any + Asset,
+        E: DataSizeEstimator<T> + FromConfig + 'static,
+    {
+        self.register_type::<T, _, _, _>(systems::update_stats_for_asset::<T, E>, CoreStage::Update)
+    }
+
+    /// Like [`RegisterSizedTypes::register_sized_asset`], but uses the given
+    /// [`DataSizeEstimator`] types to estimate the size of the [`RenderAsset`]
+    /// and its prepared format.
+    #[cfg(feature = "bevy_render")]
+    fn register_render_asset_with_estimator<T, E, F>(&mut self) -> &mut Self
+    where
+        T: Any + RenderAsset,
+        E: DataSizeEstimator<T> + FromConfig + 'static,
+        <T as RenderAsset>::PreparedAsset: Any,
+        F: DataSizeEstimator<<T as RenderAsset>::PreparedAsset> + FromConfig + 'static;
+}
+
+impl RegisterTypesWithEstimator for App {
+    #[cfg(feature = "bevy_render")]
+    fn register_render_asset_with_estimator<T, E, F>(&mut self) -> &mut Self
+    where
+        T: Any + RenderAsset,
+        E: DataSizeEstimator<T> + FromConfig + 'static,
+        <T as RenderAsset>::PreparedAsset: Any,
+        F: DataSizeEstimator<<T as RenderAsset>::PreparedAsset> + FromConfig + 'static,
+    {
+        RegisterTypes::register_type::<T, _, _, _>(
+            self,
+            systems::update_stats_for_asset::<T, E>,
+            CoreStage::Update,
+        );
+
+        if let Ok(render_app) = self.get_sub_app_mut(RenderApp) {
+            RegisterTypes::register_type::<<T as RenderAsset>::PreparedAsset, _, _, _>(
+                render_app,
+                systems::update_stats_for_render_asset::<T, F>,
+                // TODO: find the appropriate stage for this.
+                RenderStage::Render,
+            );
+
+            // Also register the type on the main app so there are entries for
+            // it in the hashmap.
+            register_type_on_app::<T>(self);
+        }
+
+        self
+    }
+}
+
+/// The lowest-level interface for registering types for memory usage tracking.
+pub trait RegisterTypes {
     /// Registers a type whose [`MemoryStats`] will be updated with the given
-    /// `system`.
+    /// `system`, which will run in the given `stage`.
     ///
     /// The given type `T` will be available to query on the [`MemoryUsage`]
     /// resource.
     ///
     /// [`MemoryStats`]: crate::MemoryStats
-    fn register_sized_type<T, S, Params>(&mut self, system: S) -> &mut Self
+    fn register_type<T, S, Params, L>(&mut self, system: S, stage: L) -> &mut Self
     where
         T: Any,
-        S: IntoSystemDescriptor<Params>;
+        S: IntoSystemDescriptor<Params>,
+        L: StageLabel;
 }
 
-impl RegisterSizedTypes for App {
-    fn register_sized_type<T, S, Params>(&mut self, system: S) -> &mut Self
+impl RegisterTypes for App {
+    fn register_type<T, S, Params, L>(&mut self, system: S, stage: L) -> &mut Self
     where
         T: Any,
         S: IntoSystemDescriptor<Params>,
+        L: StageLabel,
     {
-        register_type::<T>(self);
+        register_type_on_app::<T>(self);
 
-        self.add_system(system);
+        self.add_system_to_stage(stage, system);
 
         self
     }
 }
 
-fn register_type<T>(app: &mut App)
+fn register_type_on_app<T>(app: &mut App)
 where
     T: Any,
 {

src/builtins/mod.rs

@@ -1,14 +1,8 @@
 //! Support for Bevy's built-in types.
 
-use std::any::Any;
+use bevy::app::{PluginGroup, PluginGroupBuilder};
 
-use bevy::{
-    app::{App, PluginGroup, PluginGroupBuilder},
-    asset::Asset,
-    ecs::{component::Component, system::Resource},
-};
-
-use crate::{systems, DataSizeEstimator, MemoryUsagePlugin, RegisterSizedTypes};
+use crate::MemoryUsagePlugin;
 
 pub mod render;
 pub mod transform;
@@ -32,31 +26,3 @@ impl PluginGroup for DefaultMemoryUsagePlugins {
         group.add(TransformMemoryUsagePlugin);
     }
 }
-
-pub(crate) trait RegisterTypesWithEstimator: RegisterSizedTypes {
-    fn register_component_with_estimator<T, E>(&mut self) -> &mut Self
-    where
-        T: Any + Component,
-        E: DataSizeEstimator<T> + Default + 'static,
-    {
-        self.register_sized_type::<T, _, _>(systems::update_stats_for_component::<T, E>)
-    }
-
-    fn register_resource_with_estimator<T, E>(&mut self) -> &mut Self
-    where
-        T: Any + Resource,
-        E: DataSizeEstimator<T> + Default + 'static,
-    {
-        self.register_sized_type::<T, _, _>(systems::update_stats_for_resource::<T, E>)
-    }
-
-    fn register_asset_with_estimator<T, E>(&mut self) -> &mut Self
-    where
-        T: Any + Asset,
-        E: DataSizeEstimator<T> + Default + 'static,
-    {
-        self.register_sized_type::<T, _, _>(systems::update_stats_for_asset::<T, E>)
-    }
-}
-
-impl RegisterTypesWithEstimator for App {}

src/builtins/render/image.rs

@@ -1,7 +1,7 @@
 //! Memory usage tracking for Bevy's [`Image`] type.
 use bevy::{app::Plugin, render::texture::Image};
 
-use crate::{builtins::RegisterTypesWithEstimator, DataSize, DataSizeEstimator};
+use crate::{app_ext::RegisterTypesWithEstimator, DataSize, DataSizeEstimator};
 
 /// Adds memory tracking for [`Image`] assets.
 #[derive(Debug, Default)]

src/builtins/render/material.rs

@@ -1,7 +1,7 @@
 //! Memory usage tracking for Bevy's [`Image`] type.
 use bevy::{app::Plugin, pbr::StandardMaterial};
 
-use crate::{builtins::RegisterTypesWithEstimator, estimator::ZeroEstimator};
+use crate::{app_ext::RegisterTypesWithEstimator, estimator::ZeroEstimator};
 
 /// Adds memory tracking for [`StandardMaterial`] assets.
 #[derive(Debug, Default)]