From 17d84eb9e279e97c542fad91ef7c5e9a6d242282 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 11:07:28 -0700 Subject: [PATCH 01/16] Create `AssetPresence` as a replacement to Option. This currently buys us nothing, but it's a nice stepping stone to making `AssetPresence` have a locked variant. --- crates/bevy_asset/src/assets.rs | 117 ++++++++++++++++++++++++-------- 1 file changed, 87 insertions(+), 30 deletions(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 3c33aa43faf91..f6055ceb83da6 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -107,7 +107,53 @@ 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), +} + +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), + } + } + + /// Gets a mutable borrow to the stored asset (if unlocked). + pub fn as_mut(&mut self) -> Option<&mut A> { + match self { + Self::None => None, + Self::Unlocked(asset) => Some(asset), + } + } + + /// 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(_)) + } } /// Stores [`Asset`] values in a Vec-like storage identified by [`AssetIndex`]. @@ -152,7 +198,7 @@ impl DenseAssetStorage { if !exists { self.len += 1; } - *value = Some(asset); + *value = AssetPresence::Unlocked(asset); Ok(exists) } else { Err(InvalidGenerationError { @@ -167,7 +213,7 @@ impl DenseAssetStorage { /// 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 +223,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 +231,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; } } }; @@ -236,13 +286,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,7 +334,10 @@ 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, queued_events: Vec>, /// Assets managed by the `Assets` struct with live strong `Handle`s @@ -353,8 +406,8 @@ 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() { self.queued_events .push(AssetEvent::Modified { id: uuid.into() }); @@ -362,7 +415,7 @@ impl Assets { self.queued_events .push(AssetEvent::Added { id: uuid.into() }); } - result + result.is_some() } pub(crate) fn insert_with_index( &mut self, @@ -416,7 +469,7 @@ 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()), } } @@ -427,7 +480,7 @@ impl Assets { 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 } => self.hash_map.get_mut(&uuid)?.as_mut(), }; if result.is_some() { self.queued_events.push(AssetEvent::Modified { id }); @@ -437,7 +490,7 @@ impl Assets { /// 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,12 +501,12 @@ 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); 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), } } @@ -512,11 +565,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. @@ -576,7 +631,7 @@ impl 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> { @@ -597,19 +652,21 @@ impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> { marker: PhantomData, }; self.queued_events.push(AssetEvent::Modified { id }); - if let Some(value) = value { + if let Some(value) = value.as_mut() { return Some((id, value)); } } } } - 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 Some(value) = value.as_mut() else { + continue; + }; + return Some((id, value)); } + None } } From d030b11b48553010cacd3ca511b0e9adc4771e16 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 11:14:43 -0700 Subject: [PATCH 02/16] Add a `Locked` variant to `AssetPresence`. This variant is still unused. --- crates/bevy_asset/src/assets.rs | 74 +++++++++++++++++++++++++-------- 1 file changed, 56 insertions(+), 18 deletions(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index f6055ceb83da6..d9aab63a19b2b 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -121,6 +121,8 @@ pub enum AssetPresence { 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 { @@ -134,14 +136,16 @@ impl AssetPresence { 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) -> Option<&mut A> { + pub fn as_mut(&mut self) -> Result<&mut A, MutableAssetError> { match self { - Self::None => None, - Self::Unlocked(asset) => Some(asset), + Self::None => Err(MutableAssetError::Missing), + Self::Unlocked(asset) => Ok(asset), + Self::Locked(_) => Err(MutableAssetError::Locked), } } @@ -152,7 +156,36 @@ impl AssetPresence { /// Returns whether the asset is present. pub fn is_some(&self) -> bool { - matches!(self, Self::Unlocked(_)) + matches!(self, Self::Unlocked(_) | Self::Locked(_)) + } +} + +/// 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, +} + +/// 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))), + } } } @@ -265,15 +298,17 @@ impl DenseAssetStorage { } } - pub(crate) fn get_mut(&mut self, index: AssetIndex) -> Option<&mut A> { - let entry = self.storage.get_mut(index.index as usize)?; + 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 => None, + Entry::None => Err(MutableAssetError::Missing), Entry::Some { value, generation } => { if *generation == index.generation { value.as_mut() } else { - None + Err(MutableAssetError::Missing) } } } @@ -390,12 +425,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`. @@ -476,13 +511,16 @@ impl Assets { /// 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`]. #[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)?.as_mut(), + 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 @@ -635,7 +673,7 @@ pub struct AssetsMutIterator<'a, A: Asset> { } 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 { @@ -652,16 +690,16 @@ impl<'a, A: Asset> Iterator for AssetsMutIterator<'a, A> { marker: PhantomData, }; self.queued_events.push(AssetEvent::Modified { id }); - if let Some(value) = value.as_mut() { - return Some((id, value)); - } + if let Ok(presence) = value.try_into() { + return Some((id, presence)); + }; } } } for (key, value) in &mut self.hash_map { let id = AssetId::Uuid { uuid: *key }; self.queued_events.push(AssetEvent::Modified { id }); - let Some(value) = value.as_mut() else { + let Ok(value) = value.try_into() else { continue; }; return Some((id, value)); From c56b6e25dc078050e5cda4cd2fa6ba79c27f9026 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Fri, 20 Sep 2024 22:51:30 -0700 Subject: [PATCH 03/16] Update ReflectAsset to match the AssetPresence API. --- crates/bevy_asset/src/reflect.rs | 77 ++++++++++++++++++++++++++++---- 1 file changed, 68 insertions(+), 9 deletions(-) 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); } } From cf8024cfe88193d0f324b6425b4379f5724c7698 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 14:19:40 -0700 Subject: [PATCH 04/16] Allow locking and unlocking assets. Now assets can be "locked" into an `Arc`. --- crates/bevy_asset/src/assets.rs | 110 ++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index d9aab63a19b2b..5fb38087397d7 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -158,6 +158,47 @@ impl AssetPresence { 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. @@ -169,6 +210,15 @@ pub enum MutableAssetError { 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. @@ -314,6 +364,40 @@ impl DenseAssetStorage { } } + /// 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.lock() + } else { + None + } + } + } + } + + /// 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 @@ -526,6 +610,32 @@ impl Assets { 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(); + 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(), + }, + } + } + + /// 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(); + 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(), + }, + } + } + /// 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>) -> AssetPresence { From 67a60c84851af08963e6b19e718a8bca05ed5207 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 18:00:14 -0700 Subject: [PATCH 05/16] Keep track of the set of currently locked assets. This will make it more performant to try unlocking all the locked assets in the "happy case" that few assets are locked. --- crates/bevy_asset/src/assets.rs | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 5fb38087397d7..2902409b510ca 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -8,7 +8,7 @@ use bevy_ecs::{ system::{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}; @@ -458,6 +458,9 @@ pub struct Assets { // `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`. @@ -473,6 +476,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(), } @@ -528,6 +532,8 @@ impl Assets { 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 { @@ -543,6 +549,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 { @@ -614,26 +625,34 @@ impl Assets { /// 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(); - match id { + 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(); - match id { + 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. @@ -652,6 +671,7 @@ impl Assets { 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).unwrap_or(AssetPresence::None), @@ -660,6 +680,7 @@ impl Assets { /// 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) => { @@ -693,6 +714,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)> { From 7450ba585edd5c4d8a3d629960cad0454ef346d3 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 18:21:45 -0700 Subject: [PATCH 06/16] Create a system to automatically unlock assets every frame. This way, if you lock an asset, you just have to drop its `Arc` and wait a frame for the asset to become mutable again (as opposed to requiring a system to figure out which assets can be unlocked and trying to unlock just those). The price is constant polling whenever an asset is locked. We have to repeatedly try unlocking an asset until it is actually ready to be unlocked. --- crates/bevy_asset/src/assets.rs | 23 ++++++++++++++++++++++- crates/bevy_asset/src/lib.rs | 14 +++++++++++++- 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 2902409b510ca..3e647c39d6e1c 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -4,8 +4,9 @@ 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, HashSet}; @@ -786,6 +787,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 diff --git a/crates/bevy_asset/src/lib.rs b/crates/bevy_asset/src/lib.rs index df5879609d676..46ad209718ef4 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 From 209465d1936f14ecf66e1968079a8deec5f18414 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 17:36:04 -0700 Subject: [PATCH 07/16] Add some convenience functions to AssetPresenceMut. This makes it easier to just unwrap a mutable borrow by just doing `.as_mut().unwrap()`. --- crates/bevy_asset/src/assets.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 3e647c39d6e1c..6316f7f1ea333 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -240,6 +240,25 @@ impl<'a, A> TryFrom<&'a mut AssetPresence> for AssetPresenceMut<'a, A> { } } +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`]. struct DenseAssetStorage { storage: Vec>, From e654299a84e4610904c3d1be19c43aa1579d7f8a Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 16:31:49 -0700 Subject: [PATCH 08/16] Flatten out the logic for `DenseAssetStorage::insert`. --- crates/bevy_asset/src/assets.rs | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 6316f7f1ea333..276ebe91ec8d2 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -295,23 +295,21 @@ 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 = AssetPresence::Unlocked(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). From ed806533b9c8d27f967ae2e38d76471fcbf52155 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 17:11:29 -0700 Subject: [PATCH 09/16] Handle locked assets when moving assets from the render world to the main world. For simplicity, we just clone the assets. In theory, we could have an enum to store the RenderAssets as either a locked Arc or an Unlocked asset, but I'd rather just avoid that branch. --- crates/bevy_render/src/render_asset.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) 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())); From 2574e624847468c9926d11d6ee10f978d094f3b3 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 17:20:10 -0700 Subject: [PATCH 10/16] Handle materials being locked when setting wireframe colors. --- crates/bevy_pbr/src/wireframe.rs | 14 ++++++++++++-- crates/bevy_sprite/src/mesh2d/wireframe2d.rs | 14 ++++++++++++-- 2 files changed, 24 insertions(+), 4 deletions(-) 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_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; } } From 13019fe5b22d53c41b65ba04807116c4114e3dbe Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 22:47:32 -0700 Subject: [PATCH 11/16] Handle meshlets possibly being locked. This is actually a totally fine case, we just need a borrow to be able to clone the internal Arc's. --- crates/bevy_pbr/src/meshlet/meshlet_mesh_manager.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 From c8ffbf147618d612e370770c6a709fc02c16d250 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 19:32:59 -0700 Subject: [PATCH 12/16] Update all the examples to handle asset locking. --- examples/3d/animated_material.rs | 2 +- examples/3d/parallax_mapping.rs | 10 +++++++--- examples/3d/tonemapping.rs | 2 +- examples/3d/transparency_3d.rs | 3 ++- examples/animation/animation_masks.rs | 2 +- examples/asset/alter_mesh.rs | 2 +- examples/asset/alter_sprite.rs | 2 +- examples/asset/asset_decompression.rs | 8 +++++--- examples/ui/ui_material.rs | 2 +- 9 files changed, 20 insertions(+), 13 deletions(-) diff --git a/examples/3d/animated_material.rs b/examples/3d/animated_material.rs index 131284568ddc5..52c1aa34c2f9d 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 Ok(material) = materials.get_mut(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..9a630f8417d1b 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 Ok(mat) = materials.get_mut(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/ui/ui_material.rs b/examples/ui/ui_material.rs index 17a8ca9b60511..a6ad16d8f3dc9 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 Ok(material) = materials.get_mut(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); From a8f4663babad692c6d0355fae3a086704038403d Mon Sep 17 00:00:00 2001 From: andriyDev Date: Sat, 21 Sep 2024 23:25:57 -0700 Subject: [PATCH 13/16] Add a test of asset locking and unlocking. --- crates/bevy_asset/src/lib.rs | 53 ++++++++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/crates/bevy_asset/src/lib.rs b/crates/bevy_asset/src/lib.rs index 46ad209718ef4..0a3539b96da9e 100644 --- a/crates/bevy_asset/src/lib.rs +++ b/crates/bevy_asset/src/lib.rs @@ -627,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}; @@ -1552,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 From 167aad4139208d31004c5a6ccca3bf420b4f7367 Mon Sep 17 00:00:00 2001 From: andriyDev Date: Thu, 26 Sep 2024 23:51:35 -0700 Subject: [PATCH 14/16] Create an example for the use of asset locking. --- Cargo.toml | 11 ++ examples/README.md | 1 + examples/asset/asset_locking.rs | 254 ++++++++++++++++++++++++++++ examples/asset/files/heights.li.ron | 6 + 4 files changed, 272 insertions(+) create mode 100644 examples/asset/asset_locking.rs create mode 100644 examples/asset/files/heights.li.ron 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/examples/README.md b/examples/README.md index 616e37648aed0..fdfac723dd221 100644 --- a/examples/README.md +++ b/examples/README.md @@ -225,6 +225,7 @@ Example | Description [Alter Sprite](../examples/asset/alter_sprite.rs) | Shows how to modify texture assets after spawning. [Asset Decompression](../examples/asset/asset_decompression.rs) | Demonstrates loading a compressed asset [Asset Loading](../examples/asset/asset_loading.rs) | Demonstrates various methods to load assets +[Asset Locking](../examples/asset/asset_locking.rs) | Demonstrates how to lock assets and use them in an async context [Asset Processing](../examples/asset/processing/asset_processing.rs) | Demonstrates how to process and load custom assets [Asset Settings](../examples/asset/asset_settings.rs) | Demonstrates various methods of applying settings when loading an asset [Custom Asset](../examples/asset/custom_asset.rs) | Implements a custom asset loader 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 From 816fb296c7910befc0ce642c4ef7858105f65a1f Mon Sep 17 00:00:00 2001 From: andriyDev Date: Fri, 27 Sep 2024 13:29:48 -0700 Subject: [PATCH 15/16] Implement `get_mut_or_clone`. This makes working with asset locking much more palatable, as users don't need to explicitly handle the locking case in most situations. --- crates/bevy_asset/src/assets.rs | 35 ++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 276ebe91ec8d2..926751b37f6eb 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -622,7 +622,10 @@ impl Assets { } /// 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>) -> Result<&mut A, MutableAssetError> { let id: AssetId = id.into(); @@ -840,6 +843,36 @@ 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>, From 2c873770db793fcdd2e440b759ffe9d1893888bf Mon Sep 17 00:00:00 2001 From: andriyDev Date: Fri, 27 Sep 2024 13:35:19 -0700 Subject: [PATCH 16/16] Use `get_mut_or_clone` in all the examples. This is likely the common case for users, so we should have our examples use this function. --- examples/3d/animated_material.rs | 2 +- examples/3d/tonemapping.rs | 2 +- examples/animation/animation_masks.rs | 3 ++- examples/asset/alter_mesh.rs | 2 +- examples/asset/alter_sprite.rs | 2 +- examples/ui/ui_material.rs | 2 +- 6 files changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/3d/animated_material.rs b/examples/3d/animated_material.rs index 52c1aa34c2f9d..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 Ok(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/tonemapping.rs b/examples/3d/tonemapping.rs index 9a630f8417d1b..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 Ok(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/animation/animation_masks.rs b/examples/animation/animation_masks.rs index ac5f9db18cf5f..9653b9c688e9d 100644 --- a/examples/animation/animation_masks.rs +++ b/examples/animation/animation_masks.rs @@ -341,7 +341,8 @@ fn handle_button_toggles( // iterate just for clarity's sake.) for (animation_graph_handle, animation_player) in animation_players.iter_mut() { // The animation graph needs to have loaded. - let Ok(animation_graph) = animation_graphs.get_mut(animation_graph_handle) else { + let Some(animation_graph) = animation_graphs.get_mut_or_clone(animation_graph_handle) + else { continue; }; diff --git a/examples/asset/alter_mesh.rs b/examples/asset/alter_mesh.rs index fe93686d01ecf..a9f9afe321beb 100644 --- a/examples/asset/alter_mesh.rs +++ b/examples/asset/alter_mesh.rs @@ -205,7 +205,7 @@ fn alter_mesh( }; // Obtain a mutable reference to the Mesh asset. - let Ok(mesh) = meshes.get_mut(handle) else { + let Some(mesh) = meshes.get_mut_or_clone(handle) else { return; }; diff --git a/examples/asset/alter_sprite.rs b/examples/asset/alter_sprite.rs index a74aa61e32e87..1f1eeb9fff226 100644 --- a/examples/asset/alter_sprite.rs +++ b/examples/asset/alter_sprite.rs @@ -150,7 +150,7 @@ fn alter_asset(mut images: ResMut>, left_bird: Query<&Handle