Throttle render assets

crates/bevy_render/src/lib.rs

@@ -58,7 +58,9 @@ pub use extract_param::Extract;
 
 use bevy_hierarchy::ValidParentCheckPlugin;
 use bevy_window::{PrimaryWindow, RawHandleWrapper};
+use extract_resource::ExtractResourcePlugin;
 use globals::GlobalsPlugin;
+use render_asset::RenderAssetBytesPerFrame;
 use renderer::{RenderAdapter, RenderAdapterInfo, RenderDevice, RenderQueue};
 
 use crate::mesh::GpuMesh;
@@ -334,6 +336,9 @@ impl Plugin for RenderPlugin {
             MorphPlugin,
         ));
 
+        app.init_resource::<RenderAssetBytesPerFrame>()
+            .add_plugins(ExtractResourcePlugin::<RenderAssetBytesPerFrame>::default());
+
         app.register_type::<alpha::AlphaMode>()
             // These types cannot be registered in bevy_color, as it does not depend on the rest of Bevy
             .register_type::<bevy_color::Color>()
@@ -375,7 +380,14 @@ impl Plugin for RenderPlugin {
                 .insert_resource(device)
                 .insert_resource(queue)
                 .insert_resource(render_adapter)
-                .insert_resource(adapter_info);
+                .insert_resource(adapter_info)
+                .add_systems(
+                    Render,
+                    (|mut bpf: ResMut<RenderAssetBytesPerFrame>| {
+                        bpf.reset();
+                    })
+                    .in_set(RenderSet::Cleanup),
+                );
         }
     }
 }

crates/bevy_render/src/mesh/mesh/mod.rs

@@ -1477,6 +1477,18 @@ impl RenderAsset for GpuMesh {
         mesh.asset_usage
     }
 
+    fn byte_len(mesh: &Self::SourceAsset) -> Option<usize> {
+        let mut vertex_size = 0;
+        for attribute_data in mesh.attributes.values() {
+            let vertex_format = attribute_data.attribute.format;
+            vertex_size += vertex_format.get_size() as usize;
+        }
+
+        let vertex_count = mesh.count_vertices();
+        let index_bytes = mesh.get_index_buffer_bytes().map(<[_]>::len).unwrap_or(0);
+        Some(vertex_size * vertex_count + index_bytes)
+    }
+
     /// Converts the extracted mesh a into [`GpuMesh`].
     fn prepare_asset(
         mesh: Self::SourceAsset,

crates/bevy_render/src/render_asset.rs

@@ -8,7 +8,8 @@ use bevy_ecs::{
     world::{FromWorld, Mut},
 };
 use bevy_reflect::{Reflect, ReflectDeserialize, ReflectSerialize};
-use bevy_utils::{HashMap, HashSet};
+use bevy_render_macros::ExtractResource;
+use bevy_utils::{tracing::debug, HashMap, HashSet};
 use serde::{Deserialize, Serialize};
 use std::marker::PhantomData;
 use thiserror::Error;
@@ -41,6 +42,14 @@ pub trait RenderAsset: Send + Sync + 'static + Sized {
         RenderAssetUsages::default()
     }
 
+    /// Size of the data the asset will upload to the gpu. Specifying a return value
+    /// will allow the asset to be throttled via [`RenderAssetBytesPerFrame`].
+    #[inline]
+    #[allow(unused_variables)]
+    fn byte_len(source_asset: &Self::SourceAsset) -> Option<usize> {
+        None
+    }
+
     /// Prepares the [`RenderAsset::SourceAsset`] for the GPU by transforming it into a [`RenderAsset`].
     ///
     /// ECS data may be accessed via `param`.
@@ -161,13 +170,15 @@ impl<A: RenderAsset> RenderAssetDependency for A {
 pub struct ExtractedAssets<A: RenderAsset> {
     extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
     removed: Vec<AssetId<A::SourceAsset>>,
+    added: Vec<AssetId<A::SourceAsset>>,
 }
 
 impl<A: RenderAsset> Default for ExtractedAssets<A> {
     fn default() -> Self {
         Self {
             extracted: Default::default(),
             removed: Default::default(),
+            added: Default::default(),
         }
     }
 }
@@ -253,16 +264,19 @@ fn extract_render_asset<A: RenderAsset>(mut commands: Commands, mut main_world:
             }
 
             let mut extracted_assets = Vec::new();
+            let mut added = Vec::new();
             for id in changed_assets.drain() {
                 if let Some(asset) = assets.get(id) {
                     let asset_usage = A::asset_usage(asset);
                     if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
                         if asset_usage == RenderAssetUsages::RENDER_WORLD {
                             if let Some(asset) = assets.remove(id) {
                                 extracted_assets.push((id, asset));
+                                added.push(id);
                             }
                         } else {
                             extracted_assets.push((id, asset.clone()));
+                            added.push(id);
                         }
                     }
                 }
@@ -271,6 +285,7 @@ fn extract_render_asset<A: RenderAsset>(mut commands: Commands, mut main_world:
             commands.insert_resource(ExtractedAssets::<A> {
                 extracted: extracted_assets,
                 removed,
+                added,
             });
             cached_state.state.apply(world);
         },
@@ -299,17 +314,36 @@ pub fn prepare_assets<A: RenderAsset>(
     mut render_assets: ResMut<RenderAssets<A>>,
     mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
     param: StaticSystemParam<<A as RenderAsset>::Param>,
+    mut bpf: ResMut<RenderAssetBytesPerFrame>,
 ) {
+    let mut wrote_asset_count = 0;
+
     let mut param = param.into_inner();
     let queued_assets = std::mem::take(&mut prepare_next_frame.assets);
     for (id, extracted_asset) in queued_assets {
-        if extracted_assets.removed.contains(&id) {
+        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
             continue;
         }
 
+        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
+            // we could check if available bytes > byte_len here, but we want to make some
+            // forward progress even if the asset is larger than the max bytes per frame.
+            // this way we always write at least one (sized) asset per frame.
+            // in future we could also consider partial asset uploads.
+            if bpf.exhausted() {
+                prepare_next_frame.assets.push((id, extracted_asset));
+                continue;
+            }
+            size
+        } else {
+            0
+        };
+
         match A::prepare_asset(extracted_asset, &mut param) {
             Ok(prepared_asset) => {
                 render_assets.insert(id, prepared_asset);
+                bpf.write_bytes(write_bytes);
+                wrote_asset_count += 1;
             }
             Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
                 prepare_next_frame.assets.push((id, extracted_asset));
@@ -322,13 +356,91 @@ pub fn prepare_assets<A: RenderAsset>(
     }
 
     for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
+        // we remove previous here to ensure that if we are updating the asset then
+        // any users will not see the old asset after a new asset is extracted,
+        // even if the new asset is not yet ready or we are out of bytes to write.
+        render_assets.remove(id);
+
+        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
+            if bpf.exhausted() {
+                prepare_next_frame.assets.push((id, extracted_asset));
+                continue;
+            }
+            size
+        } else {
+            0
+        };
+
         match A::prepare_asset(extracted_asset, &mut param) {
             Ok(prepared_asset) => {
                 render_assets.insert(id, prepared_asset);
+                bpf.write_bytes(write_bytes);
+                wrote_asset_count += 1;
             }
             Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
                 prepare_next_frame.assets.push((id, extracted_asset));
             }
         }
     }
+
+    if bpf.exhausted() {
+        debug!(
+            "{} write budget exhausted with {} assets remaining (wrote {})",
+            std::any::type_name::<A>(),
+            prepare_next_frame.assets.len(),
+            wrote_asset_count
+        );
+    }
+}
+
+/// A resource that attempts to limit the amount of data transferred from cpu to gpu
+/// each frame, preventing choppy frames at the cost of waiting longer for gpu assets
+/// to become available
+#[derive(Resource, Default, Debug, Clone, Copy, ExtractResource)]
+pub struct RenderAssetBytesPerFrame {
+    pub max_bytes: Option<usize>,
+    pub available: usize,
+}
+
+impl RenderAssetBytesPerFrame {
+    /// `max_bytes`: the number of bytes to write per frame.
+    /// this is a soft limit: only full assets are written currently, uploading stops
+    /// after the first asset that exceeds the limit.
+    /// To participate, assets should implement [`RenderAsset::byte_len`]. If the default
+    /// is not overridden, the assets are assumed to be small enough to upload without restriction.
+    pub fn new(max_bytes: usize) -> Self {
+        Self {
+            max_bytes: Some(max_bytes),
+            available: 0,
+        }
+    }
+
+    /// Reset the available bytes. Called once per frame by the [`crate::RenderPlugin`].
+    pub fn reset(&mut self) {
+        self.available = self.max_bytes.unwrap_or(usize::MAX);
+    }
+
+    /// check how many bytes are available since the last reset
+    pub fn available_bytes(&self, required_bytes: usize) -> usize {
+        if self.max_bytes.is_none() {
+            return required_bytes;
+        }
+
+        required_bytes.min(self.available)
+    }
+
+    /// decrease the available bytes for the current frame
+    fn write_bytes(&mut self, bytes: usize) {
+        if self.max_bytes.is_none() {
+            return;
+        }
+
+        let write_bytes = bytes.min(self.available);
+        self.available -= write_bytes;
+    }
+
+    // check if any bytes remain available for writing this frame
+    fn exhausted(&self) -> bool {
+        self.available == 0
+    }
 }

crates/bevy_render/src/texture/image.rs

@@ -839,6 +839,11 @@ impl RenderAsset for GpuImage {
         image.asset_usage
     }
 
+    #[inline]
+    fn byte_len(image: &Self::SourceAsset) -> Option<usize> {
+        Some(image.data.len())
+    }
+
     /// Converts the extracted image into a [`GpuImage`].
     fn prepare_asset(
         image: Self::SourceAsset,