Skip to main content

bevy_render/
render_asset.rs

1use crate::{
2    render_resource::AsBindGroupError, Extract, ExtractSchedule, MainWorld, Render, RenderApp,
3    RenderStartup, RenderSystems, Res,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_asset::{Asset, AssetEvent, AssetId, Assets, RenderAssetUsages};
7use bevy_ecs::{
8    prelude::{Commands, IntoScheduleConfigs, Local, MessageReader, ResMut, Resource},
9    schedule::{ScheduleConfigs, SystemSet},
10    system::{ScheduleSystem, StaticSystemParam, SystemParam, SystemParamItem, SystemState},
11    world::{FromWorld, Mut},
12};
13use bevy_log::{debug, error};
14use bevy_platform::collections::{HashMap, HashSet};
15use core::marker::PhantomData;
16use core::sync::atomic::{AtomicUsize, Ordering};
17use thiserror::Error;
18
19#[derive(Debug, Error)]
20pub enum PrepareAssetError<E: Send + Sync + 'static> {
21    #[error("Failed to prepare asset")]
22    RetryNextUpdate(E),
23    #[error("Failed to build bind group: {0}")]
24    AsBindGroupError(AsBindGroupError),
25}
26
27/// The system set during which we extract modified assets to the render world.
28#[derive(SystemSet, Clone, PartialEq, Eq, Debug, Hash)]
29pub struct AssetExtractionSystems;
30
31/// Error returned when an asset due for extraction has already been extracted
32#[derive(Debug, Error)]
33pub enum AssetExtractionError {
34    #[error("The asset has already been extracted")]
35    AlreadyExtracted,
36    #[error("The asset type does not support extraction. To clone the asset to the renderworld, use `RenderAssetUsages::default()`")]
37    NoExtractionImplementation,
38}
39
40/// Describes how an asset gets extracted and prepared for rendering.
41///
42/// In the [`ExtractSchedule`] step the [`RenderAsset::SourceAsset`] is transferred
43/// from the "main world" into the "render world".
44///
45/// After that in the [`RenderSystems::PrepareAssets`] step the extracted asset
46/// is transformed into its GPU-representation of type [`RenderAsset`].
47pub trait RenderAsset: Send + Sync + 'static + Sized {
48    /// The representation of the asset in the "main world".
49    type SourceAsset: Asset + Clone;
50
51    /// Specifies all ECS data required by [`RenderAsset::prepare_asset`].
52    ///
53    /// For convenience use the [`lifetimeless`](bevy_ecs::system::lifetimeless) [`SystemParam`].
54    type Param: SystemParam;
55
56    /// Whether or not to unload the asset after extracting it to the render world.
57    #[inline]
58    fn asset_usage(_source_asset: &Self::SourceAsset) -> RenderAssetUsages {
59        RenderAssetUsages::default()
60    }
61
62    /// Size of the data the asset will upload to the gpu. Specifying a return value
63    /// will allow the asset to be throttled via [`RenderAssetBytesPerFrame`].
64    #[inline]
65    #[expect(
66        unused_variables,
67        reason = "The parameters here are intentionally unused by the default implementation; however, putting underscores here will result in the underscores being copied by rust-analyzer's tab completion."
68    )]
69    fn byte_len(source_asset: &Self::SourceAsset) -> Option<usize> {
70        None
71    }
72
73    /// Prepares the [`RenderAsset::SourceAsset`] for the GPU by transforming it into a [`RenderAsset`].
74    ///
75    /// ECS data may be accessed via `param`.
76    fn prepare_asset(
77        source_asset: Self::SourceAsset,
78        asset_id: AssetId<Self::SourceAsset>,
79        param: &mut SystemParamItem<Self::Param>,
80        previous_asset: Option<&Self>,
81    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>>;
82
83    /// Called whenever the [`RenderAsset::SourceAsset`] has been removed.
84    ///
85    /// You can implement this method if you need to access ECS data (via
86    /// `_param`) in order to perform cleanup tasks when the asset is removed.
87    ///
88    /// The default implementation does nothing.
89    fn unload_asset(
90        _source_asset: AssetId<Self::SourceAsset>,
91        _param: &mut SystemParamItem<Self::Param>,
92    ) {
93    }
94
95    /// Make a copy of the asset to be moved to the `RenderWorld` / gpu. Heavy internal data (pixels, vertex attributes)
96    /// should be moved into the copy, leaving this asset with only metadata.
97    /// An error may be returned to indicate that the asset has already been extracted, and should not
98    /// have been modified on the CPU side (as it cannot be transferred to GPU again).
99    /// The previous GPU asset is also provided, which can be used to check if the modification is valid.
100    fn take_gpu_data(
101        _source: &mut Self::SourceAsset,
102        _previous_gpu_asset: Option<&Self>,
103    ) -> Result<Self::SourceAsset, AssetExtractionError> {
104        Err(AssetExtractionError::NoExtractionImplementation)
105    }
106}
107
108/// This plugin extracts the changed assets from the "app world" into the "render world"
109/// and prepares them for the GPU. They can then be accessed from the [`RenderAssets`] resource.
110///
111/// Therefore it sets up the [`ExtractSchedule`] and
112/// [`RenderSystems::PrepareAssets`] steps for the specified [`RenderAsset`].
113///
114/// The `AFTER` generic parameter can be used to specify that `A::prepare_asset` should not be run until
115/// `prepare_assets::<AFTER>` has completed. This allows the `prepare_asset` function to depend on another
116/// prepared [`RenderAsset`], for example `Mesh::prepare_asset` relies on `RenderAssets::<GpuImage>` for morph
117/// targets, so the plugin is created as `RenderAssetPlugin::<RenderMesh, GpuImage>::default()`.
118pub struct RenderAssetPlugin<A: RenderAsset, AFTER: RenderAssetDependency + 'static = ()> {
119    phantom: PhantomData<fn() -> (A, AFTER)>,
120}
121
122impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Default
123    for RenderAssetPlugin<A, AFTER>
124{
125    fn default() -> Self {
126        Self {
127            phantom: Default::default(),
128        }
129    }
130}
131
132impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Plugin
133    for RenderAssetPlugin<A, AFTER>
134{
135    fn build(&self, app: &mut App) {
136        app.init_resource::<CachedExtractRenderAssetSystemState<A>>();
137        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
138            render_app
139                .init_resource::<ExtractedAssets<A>>()
140                .init_resource::<RenderAssets<A>>()
141                .allow_ambiguous_resource::<RenderAssets<A>>()
142                .init_resource::<PrepareNextFrameAssets<A>>()
143                .add_systems(RenderStartup, collect_render_assets_to_reextract::<A>)
144                .add_systems(
145                    ExtractSchedule,
146                    extract_render_asset::<A>.in_set(AssetExtractionSystems),
147                );
148            AFTER::register_system(
149                render_app,
150                prepare_assets::<A>.in_set(RenderSystems::PrepareAssets),
151            );
152        }
153    }
154}
155
156// helper to allow specifying dependencies between render assets
157pub trait RenderAssetDependency {
158    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>);
159}
160
161impl RenderAssetDependency for () {
162    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
163        render_app.add_systems(Render, system);
164    }
165}
166
167impl<A: RenderAsset> RenderAssetDependency for A {
168    fn register_system(render_app: &mut SubApp, system: ScheduleConfigs<ScheduleSystem>) {
169        render_app.add_systems(Render, system.after(prepare_assets::<A>));
170    }
171}
172
173/// Temporarily stores the extracted and removed assets of the current frame.
174#[derive(Resource)]
175pub struct ExtractedAssets<A: RenderAsset> {
176    /// The assets extracted this frame.
177    ///
178    /// These are assets that were either added or modified this frame.
179    pub extracted: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
180
181    /// IDs of the assets that were removed this frame.
182    ///
183    /// These assets will not be present in [`ExtractedAssets::extracted`].
184    pub removed: HashSet<AssetId<A::SourceAsset>>,
185
186    /// IDs of the assets that were modified this frame.
187    pub modified: HashSet<AssetId<A::SourceAsset>>,
188
189    /// IDs of the assets that were added this frame.
190    pub added: HashSet<AssetId<A::SourceAsset>>,
191}
192
193impl<A: RenderAsset> Default for ExtractedAssets<A> {
194    fn default() -> Self {
195        Self {
196            extracted: Default::default(),
197            removed: Default::default(),
198            modified: Default::default(),
199            added: Default::default(),
200        }
201    }
202}
203
204/// Stores all GPU representations ([`RenderAsset`])
205/// of [`RenderAsset::SourceAsset`] as long as they exist.
206#[derive(Resource)]
207pub struct RenderAssets<A: RenderAsset>(HashMap<AssetId<A::SourceAsset>, A>);
208
209impl<A: RenderAsset> Default for RenderAssets<A> {
210    fn default() -> Self {
211        Self(Default::default())
212    }
213}
214
215impl<A: RenderAsset> RenderAssets<A> {
216    pub fn get(&self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&A> {
217        self.0.get(&id.into())
218    }
219
220    pub fn get_mut(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<&mut A> {
221        self.0.get_mut(&id.into())
222    }
223
224    pub fn insert(&mut self, id: impl Into<AssetId<A::SourceAsset>>, value: A) -> Option<A> {
225        self.0.insert(id.into(), value)
226    }
227
228    pub fn remove(&mut self, id: impl Into<AssetId<A::SourceAsset>>) -> Option<A> {
229        self.0.remove(&id.into())
230    }
231
232    pub fn iter(&self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &A)> {
233        self.0.iter().map(|(k, v)| (*k, v))
234    }
235
236    pub fn iter_mut(&mut self) -> impl Iterator<Item = (AssetId<A::SourceAsset>, &mut A)> {
237        self.0.iter_mut().map(|(k, v)| (*k, v))
238    }
239}
240
241#[derive(Resource)]
242struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
243    state: SystemState<(
244        MessageReader<'static, 'static, AssetEvent<A::SourceAsset>>,
245        ResMut<'static, Assets<A::SourceAsset>>,
246        Option<Res<'static, RenderAssets<A>>>,
247    )>,
248}
249
250impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
251    fn from_world(world: &mut bevy_ecs::world::World) -> Self {
252        Self {
253            state: SystemState::new(world),
254        }
255    }
256}
257
258/// Resource inserted during [`RenderStartup`] containing asset IDs that need
259/// re-extraction from the main world after device recovery.
260#[derive(Resource)]
261pub(crate) struct RenderAssetsToReExtract<A: RenderAsset> {
262    ids: Vec<AssetId<A::SourceAsset>>,
263}
264
265/// Drains all asset IDs from [`RenderAssets<A>`] to mark for re-extraction.
266fn collect_render_assets_to_reextract<A: RenderAsset>(
267    mut commands: Commands,
268    mut render_assets: ResMut<RenderAssets<A>>,
269    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
270) {
271    let ids: Vec<_> = render_assets.0.drain().map(|(id, _)| id).collect();
272    prepare_next_frame.assets.clear();
273    if !ids.is_empty() {
274        commands.insert_resource(RenderAssetsToReExtract::<A> { ids });
275    }
276}
277
278/// Extracts all created or modified assets of the corresponding [`RenderAsset::SourceAsset`] type
279/// into the "render world", including any assets invalidated by device recovery.
280pub(crate) fn extract_render_asset<A: RenderAsset>(
281    mut to_reextract: Option<ResMut<RenderAssetsToReExtract<A>>>,
282    mut extracted_assets: ResMut<ExtractedAssets<A>>,
283    mut main_world: ResMut<MainWorld>,
284    mut needs_extracting: Local<HashSet<AssetId<A::SourceAsset>>>,
285) {
286    extracted_assets.extracted.clear();
287    extracted_assets.removed.clear();
288    extracted_assets.modified.clear();
289    extracted_assets.added.clear();
290    needs_extracting.clear();
291
292    let reextract_ids = to_reextract
293        .as_mut()
294        .map(|r| core::mem::take(&mut r.ids))
295        .filter(|ids| !ids.is_empty());
296
297    main_world.resource_scope(
298        |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
299            let (mut events, mut assets, maybe_render_assets) = cached_state.state.get_mut(world).unwrap();
300
301            if let Some(reextract_ids) = reextract_ids {
302                needs_extracting.extend(reextract_ids);
303            }
304
305            for event in events.read() {
306                #[expect(
307                    clippy::match_same_arms,
308                    reason = "LoadedWithDependencies is marked as a TODO, so it's likely this will no longer lint soon."
309                )]
310                match event {
311                    AssetEvent::Added { id } => {
312                        needs_extracting.insert(*id);
313                    }
314                    AssetEvent::Modified { id } => {
315                        needs_extracting.insert(*id);
316                        extracted_assets.modified.insert(*id);
317                    }
318                    AssetEvent::Removed { .. } => {
319                        // We don't care that the asset was removed from Assets<T> in the main world.
320                        // An asset is only removed from RenderAssets<T> when its last handle is dropped (AssetEvent::Unused).
321                    }
322                    AssetEvent::Unused { id } => {
323                        needs_extracting.remove(id);
324                        extracted_assets.modified.remove(id);
325                        extracted_assets.removed.insert(*id);
326                    }
327                    AssetEvent::LoadedWithDependencies { .. } => {
328                        // TODO: handle this
329                    }
330                }
331            }
332
333            for id in needs_extracting.drain() {
334                if let Some(asset) = assets.get(id) {
335                    let asset_usage = A::asset_usage(asset);
336                    if asset_usage.contains(RenderAssetUsages::RENDER_WORLD) {
337                        if asset_usage == RenderAssetUsages::RENDER_WORLD {
338                            if let Some(asset) = assets.get_mut_untracked(id) {
339                                let previous_asset = maybe_render_assets.as_ref().and_then(|render_assets| render_assets.get(id));
340                                match A::take_gpu_data(asset, previous_asset) {
341                                    Ok(gpu_data_asset) => {
342                                        extracted_assets.extracted.push((id, gpu_data_asset));
343                                        extracted_assets.added.insert(id);
344                                    }
345                                    Err(e) => {
346                                        error!("{} with RenderAssetUsages == RENDER_WORLD cannot be extracted: {e}", core::any::type_name::<A>());
347                                    }
348                                };
349                            }
350                        } else {
351                            extracted_assets.extracted.push((id, asset.clone()));
352                            extracted_assets.added.insert(id);
353                        }
354                    }
355                }
356            }
357
358            cached_state.state.apply(world);
359        },
360    );
361}
362
363// TODO: consider storing inside system?
364/// All assets that should be prepared next frame.
365#[derive(Resource)]
366pub struct PrepareNextFrameAssets<A: RenderAsset> {
367    assets: Vec<(AssetId<A::SourceAsset>, A::SourceAsset)>,
368}
369
370impl<A: RenderAsset> Default for PrepareNextFrameAssets<A> {
371    fn default() -> Self {
372        Self {
373            assets: Default::default(),
374        }
375    }
376}
377
378/// This system prepares all assets of the corresponding [`RenderAsset::SourceAsset`] type
379/// which where extracted this frame for the GPU.
380pub fn prepare_assets<A: RenderAsset>(
381    mut extracted_assets: ResMut<ExtractedAssets<A>>,
382    mut render_assets: ResMut<RenderAssets<A>>,
383    mut prepare_next_frame: ResMut<PrepareNextFrameAssets<A>>,
384    param: StaticSystemParam<<A as RenderAsset>::Param>,
385    bpf: Res<RenderAssetBytesPerFrameLimiter>,
386) {
387    let mut wrote_asset_count = 0;
388
389    let mut param = param.into_inner();
390    let queued_assets = core::mem::take(&mut prepare_next_frame.assets);
391    for (id, extracted_asset) in queued_assets {
392        if extracted_assets.removed.contains(&id) || extracted_assets.added.contains(&id) {
393            // skip previous frame's assets that have been removed or updated
394            continue;
395        }
396
397        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
398            // we could check if available bytes > byte_len here, but we want to make some
399            // forward progress even if the asset is larger than the max bytes per frame.
400            // this way we always write at least one (sized) asset per frame.
401            // in future we could also consider partial asset uploads.
402            if bpf.exhausted() {
403                prepare_next_frame.assets.push((id, extracted_asset));
404                continue;
405            }
406            size
407        } else {
408            0
409        };
410
411        let previous_asset = render_assets.get(id);
412        match A::prepare_asset(extracted_asset, id, &mut param, previous_asset) {
413            Ok(prepared_asset) => {
414                render_assets.insert(id, prepared_asset);
415                bpf.write_bytes(write_bytes);
416                wrote_asset_count += 1;
417            }
418            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
419                prepare_next_frame.assets.push((id, extracted_asset));
420            }
421            Err(PrepareAssetError::AsBindGroupError(e)) => {
422                error!(
423                    "{} Bind group construction failed: {e}",
424                    core::any::type_name::<A>()
425                );
426            }
427        }
428    }
429
430    for removed in extracted_assets.removed.drain() {
431        render_assets.remove(removed);
432        A::unload_asset(removed, &mut param);
433    }
434
435    for (id, extracted_asset) in extracted_assets.extracted.drain(..) {
436        // we remove previous here to ensure that if we are updating the asset then
437        // any users will not see the old asset after a new asset is extracted,
438        // even if the new asset is not yet ready or we are out of bytes to write.
439        let previous_asset = render_assets.remove(id);
440
441        let write_bytes = if let Some(size) = A::byte_len(&extracted_asset) {
442            if bpf.exhausted() {
443                prepare_next_frame.assets.push((id, extracted_asset));
444                continue;
445            }
446            size
447        } else {
448            0
449        };
450
451        match A::prepare_asset(extracted_asset, id, &mut param, previous_asset.as_ref()) {
452            Ok(prepared_asset) => {
453                render_assets.insert(id, prepared_asset);
454                bpf.write_bytes(write_bytes);
455                wrote_asset_count += 1;
456            }
457            Err(PrepareAssetError::RetryNextUpdate(extracted_asset)) => {
458                prepare_next_frame.assets.push((id, extracted_asset));
459            }
460            Err(PrepareAssetError::AsBindGroupError(e)) => {
461                error!(
462                    "{} Bind group construction failed: {e}",
463                    core::any::type_name::<A>()
464                );
465            }
466        }
467    }
468
469    if bpf.exhausted() && !prepare_next_frame.assets.is_empty() {
470        debug!(
471            "{} write budget exhausted with {} assets remaining (wrote {})",
472            core::any::type_name::<A>(),
473            prepare_next_frame.assets.len(),
474            wrote_asset_count
475        );
476    }
477}
478
479pub fn reset_render_asset_bytes_per_frame(
480    mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
481) {
482    bpf_limiter.reset();
483}
484
485pub fn extract_render_asset_bytes_per_frame(
486    bpf: Extract<Res<RenderAssetBytesPerFrame>>,
487    mut bpf_limiter: ResMut<RenderAssetBytesPerFrameLimiter>,
488) {
489    bpf_limiter.max_bytes = bpf.max_bytes;
490}
491
492/// A resource that defines the amount of data allowed to be transferred from CPU to GPU
493/// each frame, preventing choppy frames at the cost of waiting longer for GPU assets
494/// to become available.
495#[derive(Resource, Default)]
496pub struct RenderAssetBytesPerFrame {
497    pub max_bytes: Option<usize>,
498}
499
500impl RenderAssetBytesPerFrame {
501    /// `max_bytes`: the number of bytes to write per frame.
502    ///
503    /// This is a soft limit: only full assets are written currently, uploading stops
504    /// after the first asset that exceeds the limit.
505    ///
506    /// To participate, assets should implement [`RenderAsset::byte_len`]. If the default
507    /// is not overridden, the assets are assumed to be small enough to upload without restriction.
508    pub fn new(max_bytes: usize) -> Self {
509        Self {
510            max_bytes: Some(max_bytes),
511        }
512    }
513}
514
515/// A render-world resource that facilitates limiting the data transferred from CPU to GPU
516/// each frame, preventing choppy frames at the cost of waiting longer for GPU assets
517/// to become available.
518#[derive(Resource, Default)]
519pub struct RenderAssetBytesPerFrameLimiter {
520    /// Populated by [`RenderAssetBytesPerFrame`] during extraction.
521    pub max_bytes: Option<usize>,
522    /// Bytes written this frame.
523    pub bytes_written: AtomicUsize,
524}
525
526impl RenderAssetBytesPerFrameLimiter {
527    /// Reset the available bytes. Called once per frame during extraction by [`crate::RenderPlugin`].
528    pub fn reset(&mut self) {
529        if self.max_bytes.is_none() {
530            return;
531        }
532        self.bytes_written.store(0, Ordering::Relaxed);
533    }
534
535    /// Check how many bytes are available for writing.
536    pub fn available_bytes(&self, required_bytes: usize) -> usize {
537        if let Some(max_bytes) = self.max_bytes {
538            let total_bytes = self
539                .bytes_written
540                .fetch_add(required_bytes, Ordering::Relaxed);
541
542            // The bytes available is the inverse of the amount we overshot max_bytes
543            if total_bytes >= max_bytes {
544                required_bytes.saturating_sub(total_bytes - max_bytes)
545            } else {
546                required_bytes
547            }
548        } else {
549            required_bytes
550        }
551    }
552
553    /// Decreases the available bytes for the current frame.
554    pub(crate) fn write_bytes(&self, bytes: usize) {
555        if self.max_bytes.is_some() && bytes > 0 {
556            self.bytes_written.fetch_add(bytes, Ordering::Relaxed);
557        }
558    }
559
560    /// Returns `true` if there are no remaining bytes available for writing this frame.
561    pub(crate) fn exhausted(&self) -> bool {
562        if let Some(max_bytes) = self.max_bytes {
563            let bytes_written = self.bytes_written.load(Ordering::Relaxed);
564            bytes_written >= max_bytes
565        } else {
566            false
567        }
568    }
569}