From b30ff2ab766e71e5ef7b07fbc5f94534235d7c67 Mon Sep 17 00:00:00 2001 From: robtfm <50659922+robtfm@users.noreply.github.com> Date: Mon, 14 Aug 2023 22:27:51 +0100 Subject: [PATCH] allow asset loader pre-registration (#9429) # Objective - When loading gltf files during app creation (for example using a FromWorld impl and adding that as a resource), no loader was found. - As the gltf loader can load compressed formats, it needs to know what the GPU supports so it's not available at app creation time. ## Solution alternative to #9426 - add functionality to preregister the loader. loading assets with matching extensions will block until a real loader is registered. - preregister "gltf" and "glb". - prereigster image formats. the way this is set up, if a set of extensions are all registered with a single preregistration call, then later a loader is added that matches some of the extensions, assets using the remaining extensions will then fail. i think that should work well for image formats that we don't know are supported until later. --- crates/bevy_asset/Cargo.toml | 1 + crates/bevy_asset/src/asset_server.rs | 134 +++++++++++++++--- crates/bevy_asset/src/assets.rs | 11 ++ crates/bevy_gltf/src/lib.rs | 3 +- .../src/texture/image_texture_loader.rs | 4 +- crates/bevy_render/src/texture/mod.rs | 11 ++ 6 files changed, 138 insertions(+), 26 deletions(-) diff --git a/crates/bevy_asset/Cargo.toml b/crates/bevy_asset/Cargo.toml index 59c2b48aa5866..6f7d0886f3489 100644 --- a/crates/bevy_asset/Cargo.toml +++ b/crates/bevy_asset/Cargo.toml @@ -32,6 +32,7 @@ downcast-rs = "1.2.0" fastrand = "1.7.0" notify = { version = "6.0.0", optional = true } parking_lot = "0.12.1" +async-channel = "1.4.2" [target.'cfg(target_os = "android")'.dependencies] bevy_winit = { path = "../bevy_winit", version = "0.12.0-dev" } diff --git a/crates/bevy_asset/src/asset_server.rs b/crates/bevy_asset/src/asset_server.rs index d97f206bbdc36..2217c80c993fd 100644 --- a/crates/bevy_asset/src/asset_server.rs +++ b/crates/bevy_asset/src/asset_server.rs @@ -62,6 +62,15 @@ pub(crate) struct AssetRefCounter { pub(crate) mark_unused_assets: Arc>>, } +#[derive(Clone)] +enum MaybeAssetLoader { + Ready(Arc), + Pending { + sender: async_channel::Sender<()>, + receiver: async_channel::Receiver<()>, + }, +} + /// Internal data for the asset server. /// /// [`AssetServer`] is the public API for interacting with the asset server. @@ -70,7 +79,7 @@ pub struct AssetServerInternal { pub(crate) asset_ref_counter: AssetRefCounter, pub(crate) asset_sources: Arc>>, pub(crate) asset_lifecycles: Arc>>>, - loaders: RwLock>>, + loaders: RwLock>, extension_to_loader_index: RwLock>, handle_to_path: Arc>>>, } @@ -157,6 +166,28 @@ impl AssetServer { Assets::new(self.server.asset_ref_counter.channel.sender.clone()) } + /// Pre-register a loader that will later be added. + /// + /// Assets loaded with matching extensions will be blocked until the + /// real loader is added. + pub fn preregister_loader(&self, extensions: &[&str]) { + let mut loaders = self.server.loaders.write(); + let loader_index = loaders.len(); + for extension in extensions { + if self + .server + .extension_to_loader_index + .write() + .insert(extension.to_string(), loader_index) + .is_some() + { + warn!("duplicate preregistration for `{extension}`, any assets loaded with the previous loader will never complete."); + } + } + let (sender, receiver) = async_channel::bounded(1); + loaders.push(MaybeAssetLoader::Pending { sender, receiver }); + } + /// Adds the provided asset loader to the server. /// /// If `loader` has one or more supported extensions in conflict with loaders that came before @@ -166,14 +197,50 @@ impl AssetServer { T: AssetLoader, { let mut loaders = self.server.loaders.write(); - let loader_index = loaders.len(); + let next_loader_index = loaders.len(); + let mut maybe_existing_loader_index = None; + let mut loader_map = self.server.extension_to_loader_index.write(); + let mut maybe_sender = None; + for extension in loader.extensions() { - self.server - .extension_to_loader_index - .write() - .insert(extension.to_string(), loader_index); + if let Some(&extension_index) = loader_map.get(*extension) { + // replacing an existing entry + match maybe_existing_loader_index { + None => { + match &loaders[extension_index] { + MaybeAssetLoader::Ready(_) => { + // replacing an existing loader, nothing special to do + } + MaybeAssetLoader::Pending { sender, .. } => { + // the loader was pre-registered, store the channel to notify pending assets + maybe_sender = Some(sender.clone()); + } + } + } + Some(index) => { + // ensure the loader extensions are consistent + if index != extension_index { + warn!("inconsistent extensions between loader preregister_loader and add_loader, \ + loading `{extension}` assets will never complete."); + } + } + } + + maybe_existing_loader_index = Some(extension_index); + } else { + loader_map.insert(extension.to_string(), next_loader_index); + } + } + + if let Some(existing_index) = maybe_existing_loader_index { + loaders[existing_index] = MaybeAssetLoader::Ready(Arc::new(loader)); + if let Some(sender) = maybe_sender { + // notify after replacing the loader + let _ = sender.send_blocking(()); + } + } else { + loaders.push(MaybeAssetLoader::Ready(Arc::new(loader))); } - loaders.push(Arc::new(loader)); } /// Gets a strong handle for an asset with the provided id. @@ -188,7 +255,7 @@ impl AssetServer { HandleUntyped::strong(id.into(), sender) } - fn get_asset_loader(&self, extension: &str) -> Result, AssetServerError> { + fn get_asset_loader(&self, extension: &str) -> Result { let index = { // scope map to drop lock as soon as possible let map = self.server.extension_to_loader_index.read(); @@ -204,7 +271,8 @@ impl AssetServer { fn get_path_asset_loader>( &self, path: P, - ) -> Result, AssetServerError> { + include_pending: bool, + ) -> Result { let s = path .as_ref() .file_name() @@ -223,7 +291,9 @@ impl AssetServer { ext = &ext[idx + 1..]; exts.push(ext); if let Ok(loader) = self.get_asset_loader(ext) { - return Ok(loader); + if include_pending || matches!(loader, MaybeAssetLoader::Ready(_)) { + return Ok(loader); + } } } Err(AssetServerError::MissingAssetLoader { @@ -354,12 +424,21 @@ impl AssetServer { }; // get the according asset loader - let asset_loader = match self.get_path_asset_loader(asset_path.path()) { - Ok(loader) => loader, + let mut maybe_asset_loader = self.get_path_asset_loader(asset_path.path(), true); + + // if it's still pending, block until notified and refetch the new asset loader + if let Ok(MaybeAssetLoader::Pending { receiver, .. }) = maybe_asset_loader { + let _ = receiver.recv().await; + maybe_asset_loader = self.get_path_asset_loader(asset_path.path(), false); + } + + let asset_loader = match maybe_asset_loader { + Ok(MaybeAssetLoader::Ready(loader)) => loader, Err(err) => { set_asset_failed(); return Err(err); } + Ok(MaybeAssetLoader::Pending { .. }) => unreachable!(), }; // load the asset bytes @@ -492,7 +571,7 @@ impl AssetServer { if self.asset_io().is_dir(&child_path) { handles.extend(self.load_folder(&child_path)?); } else { - if self.get_path_asset_loader(&child_path).is_err() { + if self.get_path_asset_loader(&child_path, true).is_err() { continue; } let handle = @@ -711,8 +790,11 @@ mod test { let asset_server = setup("."); asset_server.add_loader(FakePngLoader); - let t = asset_server.get_path_asset_loader("test.png"); - assert_eq!(t.unwrap().extensions()[0], "png"); + let Ok(MaybeAssetLoader::Ready(t)) = asset_server.get_path_asset_loader("test.png", true) else { + panic!(); + }; + + assert_eq!(t.extensions()[0], "png"); } #[test] @@ -720,14 +802,16 @@ mod test { let asset_server = setup("."); asset_server.add_loader(FakePngLoader); - let t = asset_server.get_path_asset_loader("test.PNG"); - assert_eq!(t.unwrap().extensions()[0], "png"); + let Ok(MaybeAssetLoader::Ready(t)) = asset_server.get_path_asset_loader("test.PNG", true) else { + panic!(); + }; + assert_eq!(t.extensions()[0], "png"); } #[test] fn no_loader() { let asset_server = setup("."); - let t = asset_server.get_path_asset_loader("test.pong"); + let t = asset_server.get_path_asset_loader("test.pong", true); assert!(t.is_err()); } @@ -736,7 +820,7 @@ mod test { let asset_server = setup("."); assert!( - match asset_server.get_path_asset_loader("test.v1.2.3.pong") { + match asset_server.get_path_asset_loader("test.v1.2.3.pong", true) { Err(AssetServerError::MissingAssetLoader { extensions }) => extensions == vec!["v1.2.3.pong", "2.3.pong", "3.pong", "pong"], _ => false, @@ -771,8 +855,10 @@ mod test { let asset_server = setup("."); asset_server.add_loader(FakePngLoader); - let t = asset_server.get_path_asset_loader("test-v1.2.3.png"); - assert_eq!(t.unwrap().extensions()[0], "png"); + let Ok(MaybeAssetLoader::Ready(t)) = asset_server.get_path_asset_loader("test-v1.2.3.png", true) else { + panic!(); + }; + assert_eq!(t.extensions()[0], "png"); } #[test] @@ -780,8 +866,10 @@ mod test { let asset_server = setup("."); asset_server.add_loader(FakeMultipleDotLoader); - let t = asset_server.get_path_asset_loader("test.test.png"); - assert_eq!(t.unwrap().extensions()[0], "test.png"); + let Ok(MaybeAssetLoader::Ready(t)) = asset_server.get_path_asset_loader("test.test.png", true) else { + panic!(); + }; + assert_eq!(t.extensions()[0], "test.png"); } fn create_dir_and_file(file: impl AsRef) -> tempfile::TempDir { diff --git a/crates/bevy_asset/src/assets.rs b/crates/bevy_asset/src/assets.rs index 6669d9bc59632..8a91fd701ad8b 100644 --- a/crates/bevy_asset/src/assets.rs +++ b/crates/bevy_asset/src/assets.rs @@ -317,6 +317,10 @@ pub trait AddAsset { fn add_asset_loader(&mut self, loader: T) -> &mut Self where T: AssetLoader; + + /// Preregisters a loader for the given extensions, that will block asset loads until a real loader + /// is registered. + fn preregister_asset_loader(&mut self, extensions: &[&str]) -> &mut Self; } impl AddAsset for App { @@ -404,6 +408,13 @@ impl AddAsset for App { self.world.resource_mut::().add_loader(loader); self } + + fn preregister_asset_loader(&mut self, extensions: &[&str]) -> &mut Self { + self.world + .resource_mut::() + .preregister_loader(extensions); + self + } } /// Loads an internal asset from a project source file. diff --git a/crates/bevy_gltf/src/lib.rs b/crates/bevy_gltf/src/lib.rs index 90f8443aabc36..bd77ed2deba02 100644 --- a/crates/bevy_gltf/src/lib.rs +++ b/crates/bevy_gltf/src/lib.rs @@ -44,7 +44,8 @@ impl Plugin for GltfPlugin { .add_asset::() .add_asset::() .add_asset::() - .add_asset::(); + .add_asset::() + .preregister_asset_loader(&["gltf", "glb"]); } fn finish(&self, app: &mut App) { diff --git a/crates/bevy_render/src/texture/image_texture_loader.rs b/crates/bevy_render/src/texture/image_texture_loader.rs index 247d65e5a279f..da5845159b09a 100644 --- a/crates/bevy_render/src/texture/image_texture_loader.rs +++ b/crates/bevy_render/src/texture/image_texture_loader.rs @@ -17,7 +17,7 @@ pub struct ImageTextureLoader { supported_compressed_formats: CompressedImageFormats, } -const FILE_EXTENSIONS: &[&str] = &[ +pub(crate) const IMG_FILE_EXTENSIONS: &[&str] = &[ #[cfg(feature = "basis-universal")] "basis", #[cfg(feature = "bmp")] @@ -73,7 +73,7 @@ impl AssetLoader for ImageTextureLoader { } fn extensions(&self) -> &[&str] { - FILE_EXTENSIONS + IMG_FILE_EXTENSIONS } } diff --git a/crates/bevy_render/src/texture/mod.rs b/crates/bevy_render/src/texture/mod.rs index 25a65e5198a6f..350e31338889f 100644 --- a/crates/bevy_render/src/texture/mod.rs +++ b/crates/bevy_render/src/texture/mod.rs @@ -96,6 +96,17 @@ impl Plugin for ImagePlugin { update_texture_cache_system.in_set(RenderSet::Cleanup), ); } + + #[cfg(any( + feature = "png", + feature = "dds", + feature = "tga", + feature = "jpeg", + feature = "bmp", + feature = "basis-universal", + feature = "ktx2", + ))] + app.preregister_asset_loader(IMG_FILE_EXTENSIONS); } fn finish(&self, app: &mut App) {