gltf: load textures asynchronously using io task pool (bevyengine#1767)

While trying to reduce load time of gltf files, I noticed most of the loading time is spent transforming bytes into an actual texture.

This PR add asynchronously loading for them using io task pool in gltf loader. It reduces loading of a large glb file from 15 seconds to 6~8 on my laptop

To allow asynchronous tasks in an asset loader, I added a reference to the task pool from the asset server in the load context, which I can use later in the loader.

Co-authored-by: Carter Anderson <[email protected]>

crates/bevy_asset/src/asset_server.rs

@@ -288,6 +288,7 @@ impl AssetServer {
             &self.server.asset_ref_counter.channel,
             &*self.server.asset_io,
             version,
+            &self.server.task_pool,
         );
         asset_loader
             .load(&bytes, &mut load_context)

crates/bevy_asset/src/loader.rs

@@ -8,6 +8,7 @@ use bevy_ecs::{
     system::{Res, ResMut},
 };
 use bevy_reflect::{TypeUuid, TypeUuidDynamic};
+use bevy_tasks::TaskPool;
 use bevy_utils::{BoxedFuture, HashMap};
 use crossbeam_channel::{Receiver, Sender};
 use downcast_rs::{impl_downcast, Downcast};
@@ -78,6 +79,7 @@ pub struct LoadContext<'a> {
     pub(crate) labeled_assets: HashMap<Option<String>, BoxedLoadedAsset>,
     pub(crate) path: &'a Path,
     pub(crate) version: usize,
+    pub(crate) task_pool: &'a TaskPool,
 }
 
 impl<'a> LoadContext<'a> {
@@ -86,13 +88,15 @@ impl<'a> LoadContext<'a> {
         ref_change_channel: &'a RefChangeChannel,
         asset_io: &'a dyn AssetIo,
         version: usize,
+        task_pool: &'a TaskPool,
     ) -> Self {
         Self {
             ref_change_channel,
             asset_io,
             labeled_assets: Default::default(),
             version,
             path,
+            task_pool,
         }
     }
 
@@ -134,6 +138,10 @@ impl<'a> LoadContext<'a> {
         }
         asset_metas
     }
+
+    pub fn task_pool(&self) -> &TaskPool {
+        self.task_pool
+    }
 }
 
 /// The result of loading an asset of type `T`

crates/bevy_gltf/src/loader.rs

@@ -235,49 +235,33 @@ async fn load_gltf<'a, 'b>(
         })
         .collect();
 
+    // TODO: use the threaded impl on wasm once wasm thread pool doesn't deadlock on it
+    #[cfg(target_arch = "wasm32")]
     for gltf_texture in gltf.textures() {
-        let mut texture = match gltf_texture.source().source() {
-            gltf::image::Source::View { view, mime_type } => {
-                let start = view.offset() as usize;
-                let end = (view.offset() + view.length()) as usize;
-                let buffer = &buffer_data[view.buffer().index()][start..end];
-                Texture::from_buffer(buffer, ImageType::MimeType(mime_type))?
-            }
-            gltf::image::Source::Uri { uri, mime_type } => {
-                let uri = percent_encoding::percent_decode_str(uri)
-                    .decode_utf8()
-                    .unwrap();
-                let uri = uri.as_ref();
-                let (bytes, image_type) = match DataUri::parse(uri) {
-                    Ok(data_uri) => (data_uri.decode()?, ImageType::MimeType(data_uri.mime_type)),
-                    Err(()) => {
-                        let parent = load_context.path().parent().unwrap();
-                        let image_path = parent.join(uri);
-                        let bytes = load_context.read_asset_bytes(image_path.clone()).await?;
-
-                        let extension = Path::new(uri).extension().unwrap().to_str().unwrap();
-                        let image_type = ImageType::Extension(extension);
-
-                        (bytes, image_type)
-                    }
-                };
-
-                Texture::from_buffer(
-                    &bytes,
-                    mime_type
-                        .map(|mt| ImageType::MimeType(mt))
-                        .unwrap_or(image_type),
-                )?
-            }
-        };
-        let texture_label = texture_label(&gltf_texture);
-        texture.sampler = texture_sampler(&gltf_texture);
-        if linear_textures.contains(&gltf_texture.index()) {
-            texture.format = TextureFormat::Rgba8Unorm;
-        }
-        load_context.set_labeled_asset::<Texture>(&texture_label, LoadedAsset::new(texture));
+        let (texture, label) =
+            load_texture(gltf_texture, &buffer_data, &linear_textures, &load_context).await?;
+        load_context.set_labeled_asset(&label, LoadedAsset::new(texture));
     }
 
+    #[cfg(not(target_arch = "wasm32"))]
+    load_context
+        .task_pool()
+        .scope(|scope| {
+            gltf.textures().for_each(|gltf_texture| {
+                let linear_textures = &linear_textures;
+                let load_context: &LoadContext = load_context;
+                let buffer_data = &buffer_data;
+                scope.spawn(async move {
+                    load_texture(gltf_texture, buffer_data, linear_textures, load_context).await
+                });
+            });
+        })
+        .into_iter()
+        .filter_map(|result| result.ok())
+        .for_each(|(texture, label)| {
+            load_context.set_labeled_asset(&label, LoadedAsset::new(texture));
+        });
+
     let mut scenes = vec![];
     let mut named_scenes = HashMap::new();
     for scene in gltf.scenes() {
@@ -325,6 +309,54 @@ async fn load_gltf<'a, 'b>(
     Ok(())
 }
 
+async fn load_texture<'a>(
+    gltf_texture: gltf::Texture<'a>,
+    buffer_data: &[Vec<u8>],
+    linear_textures: &HashSet<usize>,
+    load_context: &LoadContext<'a>,
+) -> Result<(Texture, String), GltfError> {
+    let mut texture = match gltf_texture.source().source() {
+        gltf::image::Source::View { view, mime_type } => {
+            let start = view.offset() as usize;
+            let end = (view.offset() + view.length()) as usize;
+            let buffer = &buffer_data[view.buffer().index()][start..end];
+            Texture::from_buffer(buffer, ImageType::MimeType(mime_type))?
+        }
+        gltf::image::Source::Uri { uri, mime_type } => {
+            let uri = percent_encoding::percent_decode_str(uri)
+                .decode_utf8()
+                .unwrap();
+            let uri = uri.as_ref();
+            let (bytes, image_type) = match DataUri::parse(uri) {
+                Ok(data_uri) => (data_uri.decode()?, ImageType::MimeType(data_uri.mime_type)),
+                Err(()) => {
+                    let parent = load_context.path().parent().unwrap();
+                    let image_path = parent.join(uri);
+                    let bytes = load_context.read_asset_bytes(image_path.clone()).await?;
+
+                    let extension = Path::new(uri).extension().unwrap().to_str().unwrap();
+                    let image_type = ImageType::Extension(extension);
+
+                    (bytes, image_type)
+                }
+            };
+
+            Texture::from_buffer(
+                &bytes,
+                mime_type
+                    .map(|mt| ImageType::MimeType(mt))
+                    .unwrap_or(image_type),
+            )?
+        }
+    };
+    texture.sampler = texture_sampler(&gltf_texture);
+    if (linear_textures).contains(&gltf_texture.index()) {
+        texture.format = TextureFormat::Rgba8Unorm;
+    }
+
+    Ok((texture, texture_label(&gltf_texture)))
+}
+
 fn load_material(material: &Material, load_context: &mut LoadContext) -> Handle<StandardMaterial> {
     let material_label = material_label(&material);