diff --git a/Cargo.toml b/Cargo.toml index 6fdb2b37f4989..d17307edef797 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1458,6 +1458,17 @@ description = "Demonstrates various methods to load assets" category = "Assets" wasm = false +[[example]] +name = "asset_locking" +path = "examples/asset/asset_locking.rs" +doc-scrape-examples = true + +[package.metadata.example.asset_locking] +name = "Asset Locking" +description = "Demonstrates how to lock assets and use them in an async context" +category = "Assets" +wasm = false + [[example]] name = "asset_settings" path = "examples/asset/asset_settings.rs" diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 3c33aa43faf91..926751b37f6eb 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -4,11 +4,12 @@ use crate::{ }; use alloc::sync::Arc; use bevy_ecs::{ + change_detection::DetectChangesMut, prelude::EventWriter, - system::{Res, ResMut, Resource}, + system::{Local, Res, ResMut, Resource}, }; use bevy_reflect::{Reflect, TypePath}; -use bevy_utils::HashMap; +use bevy_utils::{HashMap, HashSet}; use core::{any::TypeId, iter::Enumerate, marker::PhantomData, sync::atomic::AtomicU32}; use crossbeam_channel::{Receiver, Sender}; use serde::{Deserialize, Serialize}; @@ -107,7 +108,155 @@ enum Entry { #[default] None, /// Some is an indicator that there is a live handle active for the entry at this [`AssetIndex`] - Some { value: Option, generation: u32 }, + Some { + value: AssetPresence, + generation: u32, + }, +} + +/// The current state of an asset's presence. +#[derive(Default)] +pub enum AssetPresence { + /// The asset is not present (either has not been added yet, or was removed). + #[default] + None, + /// The asset is present and unlocked. This asset may be mutated. + Unlocked(A), + /// The asset is present, but is locked. The asset cannot be mutated, but can be passed to threads. + Locked(Arc), +} + +impl AssetPresence { + /// Takes the current value out, and leaves [`AssetPresence::None`] in its place. + pub fn take(&mut self) -> Self { + core::mem::take(self) + } + + /// Gets a borrow to the stored asset, if there is an asset. + pub fn as_ref(&self) -> Option<&A> { + match self { + Self::None => None, + Self::Unlocked(asset) => Some(asset), + Self::Locked(asset_arc) => Some(asset_arc.as_ref()), + } + } + + /// Gets a mutable borrow to the stored asset (if unlocked). + pub fn as_mut(&mut self) -> Result<&mut A, MutableAssetError> { + match self { + Self::None => Err(MutableAssetError::Missing), + Self::Unlocked(asset) => Ok(asset), + Self::Locked(_) => Err(MutableAssetError::Locked), + } + } + + /// Returns whether the asset is missing. + pub fn is_none(&self) -> bool { + matches!(self, Self::None) + } + + /// Returns whether the asset is present. + pub fn is_some(&self) -> bool { + matches!(self, Self::Unlocked(_) | Self::Locked(_)) + } + + /// Locks the current asset, which prevents the asset from being modified + /// until it is unlocked. Returns [`None`] only if [`AssetPresence::None`]. + pub fn lock(&mut self) -> Option> { + match self { + Self::None => None, + Self::Locked(asset_arc) => Some(Arc::clone(asset_arc)), + Self::Unlocked(_) => { + let Self::Unlocked(asset) = self.take() else { + unreachable!() + }; + let asset = Arc::new(asset); + *self = Self::Locked(Arc::clone(&asset)); + Some(asset) + } + } + } + + /// Attempts to unlock the asset (returning it to a mutable state). Note an + /// already unlocked asset will return [`Ok`]. + pub fn try_unlock(&mut self) -> Result<(), UnlockAssetError> { + match self { + Self::None => Err(UnlockAssetError::Missing), + Self::Unlocked(_) => Ok(()), + Self::Locked(_) => { + let Self::Locked(asset_arc) = self.take() else { + unreachable!(); + }; + match Arc::try_unwrap(asset_arc) { + Ok(asset) => { + *self = Self::Unlocked(asset); + Ok(()) + } + Err(asset_arc) => { + *self = Self::Locked(asset_arc); + Err(UnlockAssetError::InUse) + } + } + } + } + } +} + +/// An error from trying to mutably access an asset. +#[derive(Clone, Copy, PartialEq, Eq, Error, Debug)] +pub enum MutableAssetError { + #[error("The asset is not present or has an invalid generation.")] + Missing, + #[error("The asset is currently locked and cannot be mutated.")] + Locked, +} + +/// An error from trying to unlock an asset. +#[derive(Clone, Copy, PartialEq, Eq, Error, Debug)] +pub enum UnlockAssetError { + #[error("The asset is not present or has an invalid generation.")] + Missing, + #[error("The asset is currently locked and is still being used.")] + InUse, +} + +/// A possibly mutable access to an asset. +pub enum AssetPresenceMut<'a, A> { + /// The asset is unlocked and can be mutated. + Unlocked(&'a mut A), + /// The asset is locked and cannot be mutated (but still exists). + Locked(Arc), +} + +impl<'a, A> TryFrom<&'a mut AssetPresence> for AssetPresenceMut<'a, A> { + type Error = (); + + fn try_from(value: &'a mut AssetPresence) -> Result { + match value { + AssetPresence::None => Err(()), + AssetPresence::Unlocked(value) => Ok(Self::Unlocked(value)), + AssetPresence::Locked(value_arc) => Ok(Self::Locked(Arc::clone(value_arc))), + } + } +} + +impl<'a, A> AsRef for AssetPresenceMut<'a, A> { + fn as_ref(&self) -> &A { + match self { + Self::Unlocked(asset) => asset, + Self::Locked(asset_arc) => asset_arc.as_ref(), + } + } +} + +impl<'a, A> AssetPresenceMut<'a, A> { + /// Gets a mutable borrow to the asset. Returns [`None`] if the asset is locked. + pub fn as_mut(&mut self) -> Option<&mut A> { + match self { + Self::Unlocked(asset) => Some(*asset), + Self::Locked(_) => None, + } + } } /// Stores [`Asset`] values in a Vec-like storage identified by [`AssetIndex`]. @@ -146,28 +295,26 @@ impl DenseAssetStorage { ) -> Result { self.flush(); let entry = &mut self.storage[index.index as usize]; - if let Entry::Some { value, generation } = entry { - if *generation == index.generation { - let exists = value.is_some(); - if !exists { - self.len += 1; - } - *value = Some(asset); - Ok(exists) - } else { - Err(InvalidGenerationError { - index, - current_generation: *generation, - }) - } - } else { + let Entry::Some { value, generation } = entry else { unreachable!("entries should always be valid after a flush"); + }; + if *generation != index.generation { + return Err(InvalidGenerationError { + index, + current_generation: *generation, + }); } + let exists = value.is_some(); + if !exists { + self.len += 1; + } + *value = AssetPresence::Unlocked(asset); + Ok(exists) } /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists). /// This will recycle the id and allow new entries to be inserted. - pub(crate) fn remove_dropped(&mut self, index: AssetIndex) -> Option { + pub(crate) fn remove_dropped(&mut self, index: AssetIndex) -> AssetPresence { self.remove_internal(index, |dense_storage| { dense_storage.storage[index.index as usize] = Entry::None; dense_storage.allocator.recycle(index); @@ -177,7 +324,7 @@ impl DenseAssetStorage { /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists). /// This will _not_ recycle the id. New values with the current ID can still be inserted. The ID will /// not be reused until [`DenseAssetStorage::remove_dropped`] is called. - pub(crate) fn remove_still_alive(&mut self, index: AssetIndex) -> Option { + pub(crate) fn remove_still_alive(&mut self, index: AssetIndex) -> AssetPresence { self.remove_internal(index, |_| {}) } @@ -185,15 +332,19 @@ impl DenseAssetStorage { &mut self, index: AssetIndex, removed_action: impl FnOnce(&mut Self), - ) -> Option { + ) -> AssetPresence { self.flush(); let value = match &mut self.storage[index.index as usize] { - Entry::None => return None, + Entry::None => return AssetPresence::None, Entry::Some { value, generation } => { if *generation == index.generation { - value.take().inspect(|_| self.len -= 1) + let value = value.take(); + if value.is_some() { + self.len -= 1; + } + value } else { - return None; + return AssetPresence::None; } } }; @@ -215,13 +366,31 @@ impl DenseAssetStorage { } } - pub(crate) fn get_mut(&mut self, index: AssetIndex) -> Option<&mut A> { + pub(crate) fn get_mut(&mut self, index: AssetIndex) -> Result<&mut A, MutableAssetError> { + let Some(entry) = self.storage.get_mut(index.index as usize) else { + return Err(MutableAssetError::Missing); + }; + match entry { + Entry::None => Err(MutableAssetError::Missing), + Entry::Some { value, generation } => { + if *generation == index.generation { + value.as_mut() + } else { + Err(MutableAssetError::Missing) + } + } + } + } + + /// Locks the asset at `index`, which prevents the asset from being modified + /// until it is unlocked. Returns [`None`] if the asset is missing. + pub(crate) fn lock(&mut self, index: AssetIndex) -> Option> { let entry = self.storage.get_mut(index.index as usize)?; match entry { Entry::None => None, Entry::Some { value, generation } => { if *generation == index.generation { - value.as_mut() + value.lock() } else { None } @@ -229,6 +398,24 @@ impl DenseAssetStorage { } } + /// Attempts to unlock the asset at `index` (returning it to a mutable + /// state). Note an already unlocked asset will return [`Ok`]. + pub(crate) fn try_unlock(&mut self, index: AssetIndex) -> Result<(), UnlockAssetError> { + let Some(entry) = self.storage.get_mut(index.index as usize) else { + return Err(UnlockAssetError::Missing); + }; + match entry { + Entry::None => Err(UnlockAssetError::Missing), + Entry::Some { value, generation } => { + if *generation == index.generation { + value.try_unlock() + } else { + Err(UnlockAssetError::Missing) + } + } + } + } + pub(crate) fn flush(&mut self) { // NOTE: this assumes the allocator index is monotonically increasing. let new_len = self @@ -236,13 +423,13 @@ impl DenseAssetStorage { .next_index .load(core::sync::atomic::Ordering::Relaxed); self.storage.resize_with(new_len as usize, || Entry::Some { - value: None, + value: AssetPresence::None, generation: 0, }); while let Ok(recycled) = self.allocator.recycled_receiver.try_recv() { let entry = &mut self.storage[recycled.index as usize]; *entry = Entry::Some { - value: None, + value: AssetPresence::None, generation: recycled.generation, }; } @@ -284,8 +471,14 @@ impl DenseAssetStorage { #[derive(Resource)] pub struct Assets { dense_storage: DenseAssetStorage, - hash_map: HashMap, + // Note the None variant of AssetPresence cannot be present inside the hashmap. + // We use `AssetPresence` here since locking assets requires taking the + // `AssetPresence` (and taking requires a default). + hash_map: HashMap>, handle_provider: AssetHandleProvider, + /// The set of currently locked assets. This just improves the speed of iterating through the + /// locked assets. + locked_assets: HashSet>, queued_events: Vec>, /// Assets managed by the `Assets` struct with live strong `Handle`s /// originating from `get_strong_handle`. @@ -301,6 +494,7 @@ impl Default for Assets { dense_storage, handle_provider, hash_map: Default::default(), + locked_assets: Default::default(), queued_events: Default::default(), duplicate_handles: Default::default(), } @@ -337,12 +531,12 @@ impl Assets { &mut self, id: impl Into>, insert_fn: impl FnOnce() -> A, - ) -> &mut A { + ) -> Result<&mut A, MutableAssetError> { let id: AssetId = id.into(); if self.get(id).is_none() { self.insert(id, insert_fn()); } - self.get_mut(id).unwrap() + self.get_mut(id) } /// Returns `true` if the `id` exists in this collection. Otherwise it returns `false`. @@ -353,16 +547,18 @@ impl Assets { } } - pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> Option { - let result = self.hash_map.insert(uuid, asset); + pub(crate) fn insert_with_uuid(&mut self, uuid: Uuid, asset: A) -> bool { + let result = self.hash_map.insert(uuid, AssetPresence::Unlocked(asset)); if result.is_some() { + // Make sure the asset isn't marked as locked anymore. + self.locked_assets.remove(&AssetId::Uuid { uuid }); self.queued_events .push(AssetEvent::Modified { id: uuid.into() }); } else { self.queued_events .push(AssetEvent::Added { id: uuid.into() }); } - result + result.is_some() } pub(crate) fn insert_with_index( &mut self, @@ -371,6 +567,11 @@ impl Assets { ) -> Result { let replaced = self.dense_storage.insert(index, asset)?; if replaced { + // Make sure the asset isn't marked as locked anymore. + self.locked_assets.remove(&AssetId::Index { + index, + marker: PhantomData, + }); self.queued_events .push(AssetEvent::Modified { id: index.into() }); } else { @@ -416,28 +617,68 @@ impl Assets { pub fn get(&self, id: impl Into>) -> Option<&A> { match id.into() { AssetId::Index { index, .. } => self.dense_storage.get(index), - AssetId::Uuid { uuid } => self.hash_map.get(&uuid), + AssetId::Uuid { uuid } => self.hash_map.get(&uuid).and_then(|a| a.as_ref()), } } /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if it exists. - /// Note that this supports anything that implements `Into>`, which includes [`Handle`] and [`AssetId`]. + /// + /// Consider using [`Self::get_mut_or_clone`] if the asset is [`Clone`]. Otherwise, you must + /// handle the case where the asset is lockd. Note that this supports anything that implements + /// `Into>`, which includes [`Handle`] and [`AssetId`]. #[inline] - pub fn get_mut(&mut self, id: impl Into>) -> Option<&mut A> { + pub fn get_mut(&mut self, id: impl Into>) -> Result<&mut A, MutableAssetError> { let id: AssetId = id.into(); let result = match id { AssetId::Index { index, .. } => self.dense_storage.get_mut(index), - AssetId::Uuid { uuid } => self.hash_map.get_mut(&uuid), + AssetId::Uuid { uuid } => match self.hash_map.get_mut(&uuid) { + None => Err(MutableAssetError::Missing), + Some(asset_presence) => asset_presence.as_mut(), + }, }; - if result.is_some() { + if result.is_ok() { self.queued_events.push(AssetEvent::Modified { id }); } result } + /// Locks the asset with the given `id`, which prevents the asset from being modified + /// until it is unlocked. Returns [`None`] if the asset is missing. + pub fn lock(&mut self, id: impl Into>) -> Option> { + let id: AssetId = id.into(); + let result = match id { + AssetId::Index { index, .. } => self.dense_storage.lock(index), + AssetId::Uuid { uuid } => match self.hash_map.get_mut(&uuid) { + None => None, + Some(asset_presence) => asset_presence.lock(), + }, + }; + if result.is_some() { + self.locked_assets.insert(id); + } + result + } + + /// Attempts to unlock the asset with the given `id` (returning it to a mutable + /// state). Note an already unlocked asset will return [`Ok`]. + pub fn try_unlock(&mut self, id: impl Into>) -> Result<(), UnlockAssetError> { + let id: AssetId = id.into(); + let result = match id { + AssetId::Index { index, .. } => self.dense_storage.try_unlock(index), + AssetId::Uuid { uuid } => match self.hash_map.get_mut(&uuid) { + None => Err(UnlockAssetError::Missing), + Some(asset_presence) => asset_presence.try_unlock(), + }, + }; + if result.is_ok() { + self.locked_assets.remove(&id); + } + result + } + /// Removes (and returns) the [`Asset`] with the given `id`, if it exists. /// Note that this supports anything that implements `Into>`, which includes [`Handle`] and [`AssetId`]. - pub fn remove(&mut self, id: impl Into>) -> Option { + pub fn remove(&mut self, id: impl Into>) -> AssetPresence { let id: AssetId = id.into(); let result = self.remove_untracked(id); if result.is_some() { @@ -448,17 +689,19 @@ impl Assets { /// Removes (and returns) the [`Asset`] with the given `id`, if it exists. This skips emitting [`AssetEvent::Removed`]. /// Note that this supports anything that implements `Into>`, which includes [`Handle`] and [`AssetId`]. - pub fn remove_untracked(&mut self, id: impl Into>) -> Option { + pub fn remove_untracked(&mut self, id: impl Into>) -> AssetPresence { let id: AssetId = id.into(); self.duplicate_handles.remove(&id); + self.locked_assets.remove(&id); match id { AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index), - AssetId::Uuid { uuid } => self.hash_map.remove(&uuid), + AssetId::Uuid { uuid } => self.hash_map.remove(&uuid).unwrap_or(AssetPresence::None), } } /// Removes the [`Asset`] with the given `id`. pub(crate) fn remove_dropped(&mut self, id: AssetId) { + self.locked_assets.remove(&id); match self.duplicate_handles.get_mut(&id) { None | Some(0) => {} Some(value) => { @@ -492,6 +735,12 @@ impl Assets { .chain(self.hash_map.keys().map(|uuid| AssetId::from(*uuid))) } + /// Returns an iterator over the [`AssetId`] of every locked [`Asset`] stored in this + /// collection. + pub fn locked_ids(&self) -> impl Iterator> + '_ { + self.locked_assets.iter().copied() + } + /// Returns an iterator over the [`AssetId`] and [`Asset`] ref of every asset in this collection. // PERF: this could be accelerated if we implement a skip list. Consider the cost/benefits pub fn iter(&self) -> impl Iterator, &A)> { @@ -512,11 +761,13 @@ impl Assets { (id, v) }), }) - .chain( - self.hash_map - .iter() - .map(|(i, v)| (AssetId::Uuid { uuid: *i }, v)), - ) + .chain(self.hash_map.iter().map(|(i, v)| { + ( + AssetId::Uuid { uuid: *i }, + v.as_ref() + .expect("hash_map never stores AssetPresence::None"), + ) + })) } /// Returns an iterator over the [`AssetId`] and mutable [`Asset`] ref of every asset in this collection. @@ -556,6 +807,26 @@ impl Assets { } } + /// A system that attempts to unlock any currently locked assets. + pub fn try_unlocking_locked_assets( + mut assets: ResMut, + // Use a local so that we avoid reallocating every frame. + mut locked_assets: Local>>, + ) { + locked_assets.extend(assets.locked_ids()); + for id in locked_assets.drain() { + if let Ok(()) = assets.bypass_change_detection().try_unlock(id) { + // Only mark the assets as changed if we were successful in unlocking. + assets.set_changed(); + } + } + } + + /// A run condition to check whether there are any locked assets. + pub(crate) fn has_locked_assets_condition(assets: Res) -> bool { + !assets.locked_assets.is_empty() + } + /// A system that applies accumulated asset change events to the [`Events`] resource. /// /// [`Events`]: bevy_ecs::event::Events @@ -572,15 +843,45 @@ impl Assets { } } +impl Assets { + /// Retrieves a mutable reference to the [`Asset`] with the given `id`, if it exists. + /// + /// If the asset is "unlocked", a mutable borrow is returned. If the asset is "locked", the + /// asset is cloned and replaces the existing asset (the previously locked asset is still + /// readable by holders of its [`Arc`]). + #[inline] + pub fn get_mut_or_clone(&mut self, id: impl Into>) -> Option<&mut A> { + let id = id.into(); + // PERF: This requires checking every asset if it is locked (which we redo when we + // `get_mut`). Ideally, we should only check to do the clone if we failed to `get_mut` the + // asset due to locking. Sadly, the borrow checker doesn't like it. I suspect polonius will + // fix this. + if self.locked_assets.contains(&id) { + let asset = self + .get(id) + .expect("the asset is locked, which means it's still contained in self"); + let cloned = asset.clone(); + self.insert(id, cloned); + } + match self.get_mut(id) { + Ok(borrow) => Some(borrow), + Err(MutableAssetError::Missing) => None, + Err(MutableAssetError::Locked) => { + unreachable!("the asset was already unlocked, or we explicitly cloned/replaced it"); + } + } + } +} + /// A mutable iterator over [`Assets`]. pub struct AssetsMutIterator<'a, A: Asset> { queued_events: &'a mut Vec>, dense_storage: Enumerate>>, - hash_map: bevy_utils::hashbrown::hash_map::IterMut<'a, Uuid, A>, + hash_map: bevy_utils::hashbrown::hash_map::IterMut<'a, Uuid, AssetPresence>, } impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> { - type Item = (AssetId, &'a mut A); + type Item = (AssetId, AssetPresenceMut<'a, A>); fn next(&mut self) -> Option { for (i, entry) in &mut self.dense_storage { @@ -597,19 +898,21 @@ impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> { marker: PhantomData, }; self.queued_events.push(AssetEvent::Modified { id }); - if let Some(value) = value { - return Some((id, value)); - } + if let Ok(presence) = value.try_into() { + return Some((id, presence)); + }; } } } - if let Some((key, value)) = self.hash_map.next() { + for (key, value) in &mut self.hash_map { let id = AssetId::Uuid { uuid: *key }; self.queued_events.push(AssetEvent::Modified { id }); - Some((id, value)) - } else { - None + let Ok(value) = value.try_into() else { + continue; + }; + return Some((id, value)); } + None } } diff --git a/crates/bevy_asset/src/lib.rs b/crates/bevy_asset/src/lib.rs index df5879609d676..0a3539b96da9e 100644 --- a/crates/bevy_asset/src/lib.rs +++ b/crates/bevy_asset/src/lib.rs @@ -564,7 +564,15 @@ impl AssetApp for App { .run_if(Assets::::asset_events_condition) .in_set(AssetEvents), ) - .add_systems(PreUpdate, Assets::::track_assets.in_set(TrackAssets)) + .add_systems( + PreUpdate, + ( + Assets::::track_assets.in_set(TrackAssets), + Assets::::try_unlocking_locked_assets + .run_if(Assets::::has_locked_assets_condition) + .in_set(TryUnlockAssets), + ), + ) } fn register_asset_reflect(&mut self) -> &mut Self @@ -596,6 +604,10 @@ impl AssetApp for App { #[derive(SystemSet, Hash, Debug, PartialEq, Eq, Clone)] pub struct TrackAssets; +/// A system set that holds all [`Assets::try_unlocking_locked_assets`] operations. +#[derive(SystemSet, Hash, Debug, PartialEq, Eq, Clone)] +pub struct TryUnlockAssets; + /// A system set where events accumulated in [`Assets`] are applied to the [`AssetEvent`] [`Events`] resource. /// /// [`Events`]: bevy_ecs::event::Events @@ -615,8 +627,8 @@ mod tests { }, loader::{AssetLoader, LoadContext}, Asset, AssetApp, AssetEvent, AssetId, AssetLoadError, AssetLoadFailedEvent, AssetPath, - AssetPlugin, AssetServer, Assets, DependencyLoadState, LoadState, - RecursiveDependencyLoadState, + AssetPlugin, AssetServer, Assets, DependencyLoadState, LoadState, MutableAssetError, + RecursiveDependencyLoadState, UnlockAssetError, }; use alloc::sync::Arc; use bevy_app::{App, Update}; @@ -1540,6 +1552,55 @@ mod tests { assert_eq!(events, expected_events); } + #[test] + fn asset_locking() { + let mut app = App::new(); + app.add_plugins(AssetPlugin::default()) + .init_asset::(); + + let mut assets = app.world_mut().resource_mut::>(); + let a = assets.add(SubText { text: "a".into() }); + let b = assets.add(SubText { text: "b".into() }); + let c = assets.add(SubText { text: "c".into() }); + + // Let everything propagate. + app.update(); + + let mut assets = app.world_mut().resource_mut::>(); + let a_arc = assets.lock(&a).expect("The asset exists"); + let c_arc = assets.lock(&c).expect("The asset exists"); + + assert!(matches!(assets.get_mut(&a), Err(MutableAssetError::Locked))); + assert!(assets.get_mut(&b).is_ok()); + assert!(matches!(assets.get_mut(&c), Err(MutableAssetError::Locked))); + + // Aropping the Arc should allow us to unlock the asset. + drop(a_arc); + assert!(assets.try_unlock(&a).is_ok()); + + // The asset with the living Arc still cannot be unlocked or mutably accessed. + assert!(matches!( + assets.try_unlock(&c), + Err(UnlockAssetError::InUse), + )); + assert!(matches!(assets.get_mut(&c), Err(MutableAssetError::Locked))); + + // Dropping the Arc should allow the asset to be automatically unlocked. + drop(c_arc); + app.update(); + + let mut assets = app.world_mut().resource_mut::>(); + // The asset was automatically unlocked! + assert!(assets.get_mut(&c).is_ok()); + + let b_arc = assets.lock(&b).expect("The asset exists"); + + assets.insert(&b, SubText { text: "d".into() }); + + assert_eq!(b_arc.text, "b"); + assert_eq!(assets.get_mut(&b).expect("The asset exists").text, "d"); + } + #[test] fn load_folder() { // The particular usage of GatedReader in this test will cause deadlocking if running single-threaded diff --git a/crates/bevy_asset/src/reflect.rs b/crates/bevy_asset/src/reflect.rs index 3aaa1580bb110..7ac6119e18bd7 100644 --- a/crates/bevy_asset/src/reflect.rs +++ b/crates/bevy_asset/src/reflect.rs @@ -1,9 +1,10 @@ +use alloc::sync::Arc; use core::any::{Any, TypeId}; use bevy_ecs::world::{unsafe_world_cell::UnsafeWorldCell, World}; use bevy_reflect::{FromReflect, FromType, PartialReflect, Reflect}; -use crate::{Asset, AssetId, Assets, Handle, UntypedAssetId, UntypedHandle}; +use crate::{Asset, AssetId, Assets, Handle, MutableAssetError, UntypedAssetId, UntypedHandle}; /// Type data for the [`TypeRegistry`](bevy_reflect::TypeRegistry) used to operate on reflected [`Asset`]s. /// @@ -21,12 +22,15 @@ pub struct ReflectAsset { // SAFETY: // - may only be called with an [`UnsafeWorldCell`] which can be used to access the corresponding `Assets` resource mutably // - may only be used to access **at most one** access at once - get_unchecked_mut: unsafe fn(UnsafeWorldCell<'_>, UntypedHandle) -> Option<&mut dyn Reflect>, + get_unchecked_mut: unsafe fn( + UnsafeWorldCell<'_>, + UntypedHandle, + ) -> Result<&mut dyn Reflect, MutableAssetError>, add: fn(&mut World, &dyn PartialReflect) -> UntypedHandle, insert: fn(&mut World, UntypedHandle, &dyn PartialReflect), len: fn(&World) -> usize, ids: for<'w> fn(&'w World) -> Box + 'w>, - remove: fn(&mut World, UntypedHandle) -> Option>, + remove: fn(&mut World, UntypedHandle) -> ReflectedAssetPresence, } impl ReflectAsset { @@ -50,7 +54,7 @@ impl ReflectAsset { &self, world: &'w mut World, handle: UntypedHandle, - ) -> Option<&'w mut dyn Reflect> { + ) -> Result<&'w mut dyn Reflect, MutableAssetError> { // SAFETY: unique world access #[expect( unsafe_code, @@ -96,7 +100,7 @@ impl ReflectAsset { &self, world: UnsafeWorldCell<'w>, handle: UntypedHandle, - ) -> Option<&'w mut dyn Reflect> { + ) -> Result<&'w mut dyn Reflect, MutableAssetError> { // SAFETY: requirements are deferred to the caller unsafe { (self.get_unchecked_mut)(world, handle) } } @@ -111,7 +115,7 @@ impl ReflectAsset { } /// Equivalent of [`Assets::remove`] - pub fn remove(&self, world: &mut World, handle: UntypedHandle) -> Option> { + pub fn remove(&self, world: &mut World, handle: UntypedHandle) -> ReflectedAssetPresence { (self.remove)(world, handle) } @@ -172,12 +176,67 @@ impl FromType for ReflectAsset { remove: |world, handle| { let mut assets = world.resource_mut::>(); let value = assets.remove(&handle.typed_debug_checked()); - value.map(|value| Box::new(value) as Box) + match value { + crate::AssetPresence::None => ReflectedAssetPresence::None, + crate::AssetPresence::Unlocked(asset) => { + ReflectedAssetPresence::Unlocked(Box::new(asset) as Box) + } + crate::AssetPresence::Locked(asset_arc) => { + ReflectedAssetPresence::Locked(asset_arc as Arc) + } + } }, } } } +/// The reflected version of [`crate::AssetPresence`]. +#[derive(Default)] +pub enum ReflectedAssetPresence { + /// The asset is not present (either has not been added yet, or was removed). + #[default] + None, + /// The asset is present and unlocked. This asset may be mutated. + Unlocked(Box), + /// The asset is present, but is locked. The asset cannot be mutated, but can be passed to threads. + Locked(Arc), +} + +impl ReflectedAssetPresence { + /// Takes the current value out, and leaves [`crate::AssetPresence::None`] in its place. + pub fn take(&mut self) -> Self { + core::mem::take(self) + } + + /// Gets a borrow to the stored asset, if there is an asset. + pub fn as_ref(&self) -> Option<&dyn Reflect> { + match self { + Self::None => None, + Self::Unlocked(box_reflect) => Some(box_reflect.as_ref()), + Self::Locked(arc_reflect) => Some(arc_reflect.as_ref()), + } + } + + /// Gets a mutable borrow to the stored asset (if unlocked). + pub fn as_mut(&mut self) -> Result<&mut dyn Reflect, MutableAssetError> { + match self { + Self::None => Err(MutableAssetError::Missing), + Self::Unlocked(asset) => Ok(asset.as_mut()), + Self::Locked(_) => Err(MutableAssetError::Locked), + } + } + + /// Returns whether the asset is missing. + pub fn is_none(&self) -> bool { + matches!(self, Self::None) + } + + /// Returns whether the asset is present. + pub fn is_some(&self) -> bool { + matches!(self, Self::Unlocked(_) | Self::Locked(_)) + } +} + /// Reflect type data struct relating a [`Handle`] back to the `T` asset type. /// /// Say you want to look up the asset values of a list of handles when you have access to their `&dyn Reflect` form. @@ -300,9 +359,9 @@ mod tests { .unwrap(); assert_eq!(asset.downcast_ref::().unwrap().field, "edited"); - reflect_asset + assert!(reflect_asset .remove(app.world_mut(), fetched_handle) - .unwrap(); + .is_some()); assert_eq!(reflect_asset.len(app.world()), 0); } } diff --git a/crates/bevy_pbr/src/meshlet/meshlet_mesh_manager.rs b/crates/bevy_pbr/src/meshlet/meshlet_mesh_manager.rs index 09e065b019f45..9c2da0fbb2b10 100644 --- a/crates/bevy_pbr/src/meshlet/meshlet_mesh_manager.rs +++ b/crates/bevy_pbr/src/meshlet/meshlet_mesh_manager.rs @@ -51,9 +51,8 @@ impl MeshletMeshManager { assets: &mut Assets, ) -> Range { let queue_meshlet_mesh = |asset_id: &AssetId| { - let meshlet_mesh = assets.remove_untracked(*asset_id).expect( - "MeshletMesh asset was already unloaded but is not registered with MeshletMeshManager", - ); + let meshlet_mesh = assets.remove_untracked(*asset_id); + let meshlet_mesh = meshlet_mesh.as_ref().expect("MeshletMesh asset was already unloaded but is not registered with MeshletMeshManager"); let vertex_data_slice = self .vertex_data diff --git a/crates/bevy_pbr/src/wireframe.rs b/crates/bevy_pbr/src/wireframe.rs index 86e8bd49e914e..401f40cd80317 100644 --- a/crates/bevy_pbr/src/wireframe.rs +++ b/crates/bevy_pbr/src/wireframe.rs @@ -41,7 +41,7 @@ impl Plugin for WireframePlugin { .add_systems( Update, ( - global_color_changed.run_if(resource_changed::), + global_color_changed, wireframe_color_changed, // Run `apply_global_wireframe_material` after `apply_wireframe_material` so that the global // wireframe setting is applied to a mesh on the same frame its wireframe marker component is removed. @@ -115,12 +115,22 @@ fn setup_global_wireframe_material( /// Updates the wireframe material of all entities without a [`WireframeColor`] or without a [`Wireframe`] component fn global_color_changed( + mut desynced: Local, config: Res, mut materials: ResMut>, global_material: Res, ) { - if let Some(global_material) = materials.get_mut(&global_material.handle) { + if config.is_changed() { + *desynced = true; + } + if !*desynced { + // There's been no change, so bail out early. + return; + } + if let Ok(global_material) = materials.get_mut(&global_material.handle) { global_material.color = config.default_color.into(); + // Clear the desynced flag to mark that we don't need to rerun this system. + *desynced = false; } } diff --git a/crates/bevy_render/src/render_asset.rs b/crates/bevy_render/src/render_asset.rs index 46ee40fd20ef8..4281be32ad7a7 100644 --- a/crates/bevy_render/src/render_asset.rs +++ b/crates/bevy_render/src/render_asset.rs @@ -2,7 +2,7 @@ use crate::{ render_resource::AsBindGroupError, ExtractSchedule, MainWorld, Render, RenderApp, RenderSet, }; use bevy_app::{App, Plugin, SubApp}; -use bevy_asset::{Asset, AssetEvent, AssetId, Assets}; +use bevy_asset::{Asset, AssetEvent, AssetId, AssetPresence, Assets}; use bevy_ecs::{ prelude::{Commands, EventReader, IntoSystemConfigs, ResMut, Resource}, schedule::SystemConfigs, @@ -288,9 +288,20 @@ pub(crate) fn extract_render_asset( 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.insert(id); + match assets.remove(id) { + AssetPresence::None => {} + AssetPresence::Unlocked(asset) => { + extracted_assets.push((id, asset)); + added.insert(id); + } + AssetPresence::Locked(asset) => { + // In theory, an asset can be added to the MAIN_WORLD and locked + // in the same frame, so clone the asset just to be safe. In + // practice, assets are almost always loaded and then left + // unlocked, so we skip the clone then. + extracted_assets.push((id, A::SourceAsset::clone(&asset))); + added.insert(id); + } } } else { extracted_assets.push((id, asset.clone())); diff --git a/crates/bevy_sprite/src/mesh2d/wireframe2d.rs b/crates/bevy_sprite/src/mesh2d/wireframe2d.rs index a8822c800e5ec..1e3d8e1712ae8 100644 --- a/crates/bevy_sprite/src/mesh2d/wireframe2d.rs +++ b/crates/bevy_sprite/src/mesh2d/wireframe2d.rs @@ -41,7 +41,7 @@ impl Plugin for Wireframe2dPlugin { .add_systems( Update, ( - global_color_changed.run_if(resource_changed::), + global_color_changed, wireframe_color_changed, // Run `apply_global_wireframe_material` after `apply_wireframe_material` so that the global // wireframe setting is applied to a mesh on the same frame its wireframe marker component is removed. @@ -112,12 +112,22 @@ fn setup_global_wireframe_material( /// Updates the wireframe material of all entities without a [`Wireframe2dColor`] or without a [`Wireframe2d`] component fn global_color_changed( + mut desynced: Local, config: Res, mut materials: ResMut>, global_material: Res, ) { - if let Some(global_material) = materials.get_mut(&global_material.handle) { + if config.is_changed() { + *desynced = true; + } + if !*desynced { + // There's been no change, so bail out early. + return; + } + if let Ok(global_material) = materials.get_mut(&global_material.handle) { global_material.color = config.default_color.into(); + // Clear the desynced flag to mark that we don't need to rerun this system. + *desynced = false; } } diff --git a/examples/3d/animated_material.rs b/examples/3d/animated_material.rs index 131284568ddc5..e89a08c6b37d6 100644 --- a/examples/3d/animated_material.rs +++ b/examples/3d/animated_material.rs @@ -54,7 +54,7 @@ fn animate_materials( mut materials: ResMut>, ) { for material_handle in material_handles.iter() { - if let Some(material) = materials.get_mut(material_handle) { + if let Some(material) = materials.get_mut_or_clone(material_handle) { if let Color::Hsla(ref mut hsla) = material.base_color { *hsla = hsla.rotate_hue(time.delta_seconds() * 100.0); } diff --git a/examples/3d/parallax_mapping.rs b/examples/3d/parallax_mapping.rs index f999c33670522..a8f5bcfee9c22 100644 --- a/examples/3d/parallax_mapping.rs +++ b/examples/3d/parallax_mapping.rs @@ -94,7 +94,9 @@ fn update_parallax_depth_scale( } if *depth_update { let mut text = text.single_mut(); - for (_, mat) in materials.iter_mut() { + for (_, mut mat) in materials.iter_mut() { + let mat = mat.as_mut().expect("The material is unlocked"); + let current_depth = mat.parallax_depth_scale; let new_depth = current_depth.lerp(target_depth.0, DEPTH_CHANGE_RATE); mat.parallax_depth_scale = new_depth; @@ -120,7 +122,8 @@ fn switch_method( let mut text = text.single_mut(); text.sections[2].value = format!("Method: {}\n", *current); - for (_, mat) in materials.iter_mut() { + for (_, mut mat) in materials.iter_mut() { + let mat = mat.as_mut().expect("The material is unlocked"); mat.parallax_mapping_method = current.0; } } @@ -143,7 +146,8 @@ fn update_parallax_layers( let mut text = text.single_mut(); text.sections[1].value = format!("Layers: {layer_count:.0}\n"); - for (_, mat) in materials.iter_mut() { + for (_, mut mat) in materials.iter_mut() { + let mat = mat.as_mut().expect("The material is unlocked"); mat.max_parallax_layer_count = layer_count; } } diff --git a/examples/3d/tonemapping.rs b/examples/3d/tonemapping.rs index 58c2a76caba71..92d7ee6eb2397 100644 --- a/examples/3d/tonemapping.rs +++ b/examples/3d/tonemapping.rs @@ -230,7 +230,7 @@ fn drag_drop_image( }; for mat_h in &image_mat { - if let Some(mat) = materials.get_mut(mat_h) { + if let Some(mat) = materials.get_mut_or_clone(mat_h) { mat.base_color_texture = Some(new_image.clone()); // Despawn the image viewer instructions diff --git a/examples/3d/transparency_3d.rs b/examples/3d/transparency_3d.rs index 5db6d7b693a1e..e9d45054a8f4d 100644 --- a/examples/3d/transparency_3d.rs +++ b/examples/3d/transparency_3d.rs @@ -116,7 +116,8 @@ fn setup( /// 6/8 opaque, then 5/8, etc. pub fn fade_transparency(time: Res)>, ) { for (entity, Compressed { compressed, .. }) in query.iter() { - let Some(GzAsset { uncompressed }) = compressed_assets.remove(compressed) else { - continue; + let GzAsset { uncompressed } = match compressed_assets.remove(compressed) { + AssetPresence::None => continue, + AssetPresence::Unlocked(asset) => asset, + AssetPresence::Locked(_) => panic!("this example does not lock assets"), }; let uncompressed = uncompressed.take::().unwrap(); diff --git a/examples/asset/asset_locking.rs b/examples/asset/asset_locking.rs new file mode 100644 index 0000000000000..6c13d33505fc0 --- /dev/null +++ b/examples/asset/asset_locking.rs @@ -0,0 +1,254 @@ +//! This example illustrates how to use assets in an async context (through locking). + +use std::sync::Arc; + +use bevy::{ + asset::AssetLoader, + color::palettes::tailwind, + math::FloatOrd, + prelude::*, + render::{mesh::Indices, render_asset::RenderAssetUsages}, + tasks::AsyncComputeTaskPool, +}; +use rand::{Rng, SeedableRng}; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +fn main() { + App::new() + .add_plugins( + // This just tells the asset server to look in the right examples folder + DefaultPlugins.set(AssetPlugin { + file_path: "examples/asset/files".to_string(), + ..Default::default() + }), + ) + .init_asset::() + .register_asset_loader(LinearInterpolationLoader) + .add_systems(Startup, setup) + .add_systems(Update, (start_mesh_generation, finish_mesh_generation)) + .run(); +} + +fn setup( + asset_server: Res, + mut materials: ResMut>, + meshes: Res>, + mut commands: Commands, +) { + // Spawn a camera. + commands.spawn(Camera3dBundle { + transform: Transform::from_translation(Vec3::new(15.0, 15.0, 15.0)) + .looking_at(Vec3::ZERO, Vec3::Y), + ..Default::default() + }); + + // Spawn a light. + commands.spawn(DirectionalLightBundle { + transform: Transform::default() + .looking_to(Dir3::from_xyz(1.0, -1.0, 0.0).unwrap(), Dir3::Y), + ..Default::default() + }); + + // Spawn the mesh. Reserve the handle so we can generate it later. + let mesh = meshes.reserve_handle(); + commands.spawn(PbrBundle { + mesh: mesh.clone(), + material: materials.add(StandardMaterial { + base_color: tailwind::SLATE_100.into(), + ..Default::default() + }), + ..Default::default() + }); + + // Create the parameters for mesh generation. + commands.insert_resource(MeshGeneration { + height_interpolation: asset_server.load("heights.li.ron"), + mesh, + size: UVec2::new(30, 30), + }); + + // Create the channel we will communicate across. + let (sender, receiver) = crossbeam_channel::bounded(1); + commands.insert_resource(MeshGenerationChannel { sender, receiver }); +} + +#[derive(Resource)] +struct MeshGeneration { + height_interpolation: Handle, + mesh: Handle, + size: UVec2, +} + +#[derive(Resource)] +struct MeshGenerationChannel { + sender: crossbeam_channel::Sender, + receiver: crossbeam_channel::Receiver, +} + +/// Starts a mesh generation task whenever the height interpolation asset is updated. +fn start_mesh_generation( + mut asset_events: EventReader>, + mut linear_interpolations: ResMut>, + mesh_generation: Res, + channel: Res, +) { + // Only recompute if the height interpolation asset has changed. + let regenerate_id = mesh_generation.height_interpolation.id(); + let mut recompute = false; + for asset_event in asset_events.read() { + match asset_event { + AssetEvent::Added { id } | AssetEvent::Modified { id } if *id == regenerate_id => { + recompute = true; + } + _ => {} + } + } + + if !recompute { + return; + } + + let task_pool = AsyncComputeTaskPool::get(); + let size = mesh_generation.size; + // Lock the height interpolation asset so we have an `Arc` to pass to the spawned task. + let height_interpolation = linear_interpolations + .lock(&mesh_generation.height_interpolation) + .expect("The asset is loaded"); + let channel = channel.sender.clone(); + // Spawn a task to generate the mesh, then send the resulting mesh acros the channel. + task_pool + .spawn(async move { + let mesh = generate_mesh(size, height_interpolation); + channel.send(mesh).expect("The channel never closes"); + }) + .detach(); +} + +/// Reads from the mesh generation channel and inserts the mesh asset. +fn finish_mesh_generation( + mesh_generation: Res, + channel: Res, + mut meshes: ResMut>, +) { + let Ok(mesh) = channel.receiver.try_recv() else { + return; + }; + meshes.insert(&mesh_generation.mesh, mesh); +} + +/// A basic linear interpolation curve implementation. +#[derive(Asset, TypePath, Serialize, Deserialize)] +struct LinearInterpolation(Vec<(f32, f32)>); + +impl LinearInterpolation { + /// Samples the linear interpolation at `value`. + fn sample(&self, value: f32) -> f32 { + match self.0.iter().position(|(x, _)| value < *x) { + None => self.0.last().expect("The interpolation is non-empty").1, + Some(0) => self.0.first().expect("The interpolation is non-empty").1, + Some(next) => { + let previous = next - 1; + + let (next_x, next_y) = self.0[next]; + let (previous_x, previous_y) = self.0[previous]; + + let alpha = (value - previous_x) / (next_x - previous_x); + + alpha * (next_y - previous_y) + previous_y + } + } + } +} + +#[derive(Default)] +struct LinearInterpolationLoader; + +#[derive(Debug, Error)] +enum LinearInterpolationLoaderError { + #[error(transparent)] + Io(#[from] std::io::Error), + #[error(transparent)] + RonSpannedError(#[from] ron::error::SpannedError), + #[error("The loaded interpolation is empty.")] + Empty, + #[error("The loaded interpolation contains duplicate X values")] + DuplicateXValues, +} + +impl AssetLoader for LinearInterpolationLoader { + type Asset = LinearInterpolation; + type Settings = (); + type Error = LinearInterpolationLoaderError; + + async fn load<'a>( + &'a self, + reader: &'a mut dyn bevy::asset::io::Reader, + _settings: &'a Self::Settings, + _load_context: &'a mut bevy::asset::LoadContext<'_>, + ) -> Result { + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await?; + let mut interpolation: LinearInterpolation = ron::de::from_bytes(&bytes)?; + if interpolation.0.is_empty() { + return Err(Self::Error::Empty); + } + interpolation.0.sort_by_key(|(key, _)| FloatOrd(*key)); + if interpolation + .0 + .windows(2) + .any(|window| window[0].0 == window[1].0) + { + return Err(Self::Error::DuplicateXValues); + } + Ok(interpolation) + } + + fn extensions(&self) -> &[&str] { + &["li.ron"] + } +} + +/// Generates the mesh given the interpolation curve and the size of the mesh. +fn generate_mesh(size: UVec2, interpolation: Arc) -> Mesh { + let mut rng = rand_chacha::ChaChaRng::seed_from_u64(12345); + + let center = Vec3::new((size.x as f32) / 2.0, 0.0, (size.y as f32) / -2.0); + + let mut vertices = Vec::with_capacity(((size.x + 1) * (size.y + 1)) as usize); + let mut uvs = Vec::with_capacity(((size.x + 1) * (size.y + 1)) as usize); + for y in 0..size.y + 1 { + for x in 0..size.x + 1 { + let height = interpolation.sample(rng.gen()); + vertices.push(Vec3::new(x as f32, height, -(y as f32)) - center); + uvs.push(Vec2::new(x as f32, -(y as f32))); + } + } + + let y_stride = size.x + 1; + let mut indices = Vec::with_capacity((size.x * size.y * 6) as usize); + for y in 0..size.y { + for x in 0..size.x { + indices.push(x + y * y_stride); + indices.push(x + 1 + y * y_stride); + indices.push(x + 1 + (y + 1) * y_stride); + indices.push(x + y * y_stride); + indices.push(x + 1 + (y + 1) * y_stride); + indices.push(x + (y + 1) * y_stride); + } + } + + let mut mesh = Mesh::new( + bevy_render::mesh::PrimitiveTopology::TriangleList, + RenderAssetUsages::RENDER_WORLD, + ) + .with_inserted_attribute(Mesh::ATTRIBUTE_POSITION, vertices) + .with_inserted_attribute(Mesh::ATTRIBUTE_UV_0, uvs) + .with_inserted_indices(Indices::U32(indices)); + + mesh.compute_normals(); + mesh.generate_tangents() + .expect("The tangents are well formed"); + + mesh +} diff --git a/examples/asset/files/heights.li.ron b/examples/asset/files/heights.li.ron new file mode 100644 index 0000000000000..8d4e03f71a21e --- /dev/null +++ b/examples/asset/files/heights.li.ron @@ -0,0 +1,6 @@ +([ + (0.0, 0.0), + (0.5, 0.1), + (0.7, 1.0), + (1.0, 2.0), +]) \ No newline at end of file diff --git a/examples/ui/ui_material.rs b/examples/ui/ui_material.rs index 17a8ca9b60511..fbf61566ee96f 100644 --- a/examples/ui/ui_material.rs +++ b/examples/ui/ui_material.rs @@ -87,7 +87,7 @@ fn animate( ) { let duration = 2.0; for handle in &q { - if let Some(material) = materials.get_mut(handle) { + if let Some(material) = materials.get_mut_or_clone(handle) { // rainbow color effect let new_color = Color::hsl((time.elapsed_seconds() * 60.0) % 360.0, 1., 0.5); let border_color = Color::hsl((time.elapsed_seconds() * 60.0) % 360.0, 0.75, 0.75);