Throttle render assets

crates/bevy_render/src/lib.rs

@@ -59,7 +59,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::deterministic::DeterministicRenderingConfig;
@@ -335,6 +337,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>()
@@ -376,7 +381,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

@@ -1430,6 +1430,18 @@ impl RenderAsset for Mesh {
         self.asset_usage
     }
 
+    fn byte_len(&self) -> Option<usize> {
+        let mut vertex_size = 0;
+        for attribute_data in self.attributes.values() {
+            let vertex_format = attribute_data.attribute.format;
+            vertex_size += vertex_format.get_size() as usize;
+        }
+
+        let vertex_count = self.count_vertices();
+        let index_bytes = self.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(
         self,

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;
@@ -38,6 +39,12 @@ pub trait RenderAsset: Asset + Clone {
     /// Whether or not to unload the asset after extracting it to the render world.
     fn asset_usage(&self) -> RenderAssetUsages;
 
+    /// Size of the data the asset will upload to the gpu. Specifying a return value
+    /// will allow the asset to be throttled via [`RenderAssetBytesPerFrame`].
+    fn byte_len(&self) -> Option<usize> {
+        None
+    }
+
     /// Prepares the asset for the GPU by transforming it into a [`RenderAsset::PreparedAsset`].
     ///
     /// ECS data may be accessed via `param`.
@@ -158,13 +165,15 @@ impl<A: RenderAsset> RenderAssetDependency for A {
 pub struct ExtractedAssets<A: RenderAsset> {
     extracted: Vec<(AssetId<A>, A)>,
     removed: Vec<AssetId<A>>,
+    added: Vec<AssetId<A>>,
 }
 
 impl<A: RenderAsset> Default for ExtractedAssets<A> {
     fn default() -> Self {
         Self {
             extracted: Default::default(),
             removed: Default::default(),
+            added: Default::default(),
         }
     }
 }
@@ -254,16 +263,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 = asset.asset_usage();
                     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);
                         }
                     }
                 }
@@ -272,6 +284,7 @@ fn extract_render_asset<A: RenderAsset>(mut commands: Commands, mut main_world:
             commands.insert_resource(ExtractedAssets {
                 extracted: extracted_assets,
                 removed,
+                added,
             });
             cached_state.state.apply(world);
         },
@@ -300,17 +313,28 @@ 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 = 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) {
+    let mut queued_assets = std::mem::take(&mut prepare_next_frame.assets).into_iter();
+    for (id, extracted_asset) in queued_assets.by_ref() {
+        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
             continue;
         }
 
+        if let Some(size) = extracted_asset.byte_len() {
+            if bpf.write_bytes(size) == 0 {
+                prepare_next_frame.assets.push((id, extracted_asset));
+                continue;
+            }
+        }
+
         match extracted_asset.prepare_asset(&mut param) {
             Ok(prepared_asset) => {
                 render_assets.insert(id, prepared_asset);
+                wrote += 1;
             }
             Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
                 prepare_next_frame.assets.push((id, extracted_asset));
@@ -322,14 +346,75 @@ pub fn prepare_assets<A: RenderAsset>(
         render_assets.remove(removed);
     }
 
-    for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
+    let mut extracted_assets = extracted_assets.extracted.drain(..);
+
+    for (id, extracted_asset) in extracted_assets.by_ref() {
+        if let Some(size) = extracted_asset.byte_len() {
+            if bpf.write_bytes(size) == 0 {
+                prepare_next_frame.assets.push((id, extracted_asset));
+                render_assets.remove(id);
+                continue;
+            }
+        }
+
         match extracted_asset.prepare_asset(&mut param) {
             Ok(prepared_asset) => {
                 render_assets.insert(id, prepared_asset);
+                wrote += 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
+        );
+    }
+}
+
+/// 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,
+        }
+    }
+
+    pub fn reset(&mut self) {
+        self.available = self.max_bytes.unwrap_or(usize::MAX);
+    }
+
+    pub fn write_bytes(&mut self, bytes: usize) -> usize {
+        if self.max_bytes.is_none() {
+            return bytes;
+        }
+
+        let write_bytes = bytes.min(self.available);
+        self.available -= write_bytes;
+        write_bytes
+    }
+
+    pub fn exhausted(&self) -> bool {
+        self.available == 0
+    }
 }

crates/bevy_render/src/texture/image.rs

@@ -834,6 +834,10 @@ impl RenderAsset for Image {
         SRes<DefaultImageSampler>,
     );
 
+    fn byte_len(&self) -> Option<usize> {
+        Some(self.data.len())
+    }
+
     fn asset_usage(&self) -> RenderAssetUsages {
         self.asset_usage
     }