[Adopted] Add a method for asynchronously waiting for an asset to load (#15913)

# Objective

Currently, is is very painful to wait for an asset to load from the
context of an `async` task. While bevy's `AssetServer` is asynchronous
at its core, the public API is mainly focused on being used from
synchronous contexts such as bevy systems. Currently, the best way of
waiting for an asset handle to finish loading is to have a system that
runs every frame, and either listens for `AssetEvents` or manually polls
the asset server. While this is an acceptable interface for bevy
systems, it is extremely awkward to do this in a way that integrates
well with the `async` task system. At my work we had to create our own
(inefficient) abstraction that encapsulated the boilerplate of checking
an asset's load status and waking up a task when it's done.

## Solution

Add the method `AssetServer::wait_for_asset`, which returns a future
that suspends until the asset associated with a given `Handle` either
finishes loading or fails to load.

## Testing

- CI

## Notes

This is an adoption of #14431, the above description is directly from
that original PR.

---------

Co-authored-by: Joseph <21144246+JoJoJet@users.noreply.github.com>
Co-authored-by: andriyDev <andriydzikh@gmail.com>

Author: 3 people

crates/bevy_asset/src/server/info.rs

@@ -8,7 +8,7 @@ use alloc::sync::{Arc, Weak};
 use bevy_ecs::world::World;
 use bevy_tasks::Task;
 use bevy_utils::{tracing::warn, Entry, HashMap, HashSet, TypeIdMap};
-use core::any::TypeId;
+use core::{any::TypeId, task::Waker};
 use crossbeam_channel::Sender;
 use derive_more::derive::{Display, Error, From};
 use either::Either;
@@ -36,6 +36,8 @@ pub(crate) struct AssetInfo {
     /// The number of handle drops to skip for this asset.
     /// See usage (and comments) in `get_or_create_path_handle` for context.
     handle_drops_to_skip: usize,
+    /// List of tasks waiting for this asset to complete loading
+    pub(crate) waiting_tasks: Vec<Waker>,
 }
 
 impl AssetInfo {
@@ -54,6 +56,7 @@ impl AssetInfo {
             dependants_waiting_on_load: HashSet::default(),
             dependants_waiting_on_recursive_dep_load: HashSet::default(),
             handle_drops_to_skip: 0,
+            waiting_tasks: Vec::new(),
         }
     }
 }
@@ -616,6 +619,9 @@ impl AssetInfos {
             info.load_state = LoadState::Failed(error.clone());
             info.dep_load_state = DependencyLoadState::Failed(error.clone());
             info.rec_dep_load_state = RecursiveDependencyLoadState::Failed(error.clone());
+            for waker in info.waiting_tasks.drain(..) {
+                waker.wake();
+            }
             (
                 core::mem::take(&mut info.dependants_waiting_on_load),
                 core::mem::take(&mut info.dependants_waiting_on_recursive_dep_load),

crates/bevy_asset/src/server/mod.rs

@@ -25,7 +25,7 @@ use bevy_utils::{
     tracing::{error, info},
     HashSet,
 };
-use core::{any::TypeId, future::Future, panic::AssertUnwindSafe};
+use core::{any::TypeId, future::Future, panic::AssertUnwindSafe, task::Poll};
 use crossbeam_channel::{Receiver, Sender};
 use derive_more::derive::{Display, Error, From};
 use either::Either;
@@ -413,7 +413,7 @@ impl AssetServer {
         &self,
         handle: UntypedHandle,
         path: AssetPath<'static>,
-        mut infos: RwLockWriteGuard<AssetInfos>,
+        infos: RwLockWriteGuard<AssetInfos>,
         guard: G,
     ) {
         // drop the lock on `AssetInfos` before spawning a task that may block on it in single-threaded
@@ -433,7 +433,10 @@ impl AssetServer {
         });
 
         #[cfg(not(any(target_arch = "wasm32", not(feature = "multi_threaded"))))]
-        infos.pending_tasks.insert(handle.id(), task);
+        {
+            let mut infos = infos;
+            infos.pending_tasks.insert(handle.id(), task);
+        }
 
         #[cfg(any(target_arch = "wasm32", not(feature = "multi_threaded")))]
         task.detach();
@@ -1336,6 +1339,132 @@ impl AssetServer {
                 })
             })
     }
+
+    /// Returns a future that will suspend until the specified asset and its dependencies finish
+    /// loading.
+    ///
+    /// # Errors
+    ///
+    /// This will return an error if the asset or any of its dependencies fail to load,
+    /// or if the asset has not been queued up to be loaded.
+    pub async fn wait_for_asset<A: Asset>(
+        &self,
+        // NOTE: We take a reference to a handle so we know it will outlive the future,
+        // which ensures the handle won't be dropped while waiting for the asset.
+        handle: &Handle<A>,
+    ) -> Result<(), WaitForAssetError> {
+        self.wait_for_asset_id(handle.id().untyped()).await
+    }
+
+    /// Returns a future that will suspend until the specified asset and its dependencies finish
+    /// loading.
+    ///
+    /// # Errors
+    ///
+    /// This will return an error if the asset or any of its dependencies fail to load,
+    /// or if the asset has not been queued up to be loaded.
+    pub async fn wait_for_asset_untyped(
+        &self,
+        // NOTE: We take a reference to a handle so we know it will outlive the future,
+        // which ensures the handle won't be dropped while waiting for the asset.
+        handle: &UntypedHandle,
+    ) -> Result<(), WaitForAssetError> {
+        self.wait_for_asset_id(handle.id()).await
+    }
+
+    /// Returns a future that will suspend until the specified asset and its dependencies finish
+    /// loading.
+    ///
+    /// Note that since an asset ID does not count as a reference to the asset,
+    /// the future returned from this method will *not* keep the asset alive.
+    /// This may lead to the asset unexpectedly being dropped while you are waiting for it to
+    /// finish loading.
+    ///
+    /// When calling this method, make sure a strong handle is stored elsewhere to prevent the
+    /// asset from being dropped.
+    /// If you have access to an asset's strong [`Handle`], you should prefer to call
+    /// [`AssetServer::wait_for_asset`]
+    /// or [`wait_for_assest_untyped`](Self::wait_for_asset_untyped) to ensure the asset finishes
+    /// loading.
+    ///
+    /// # Errors
+    ///
+    /// This will return an error if the asset or any of its dependencies fail to load,
+    /// or if the asset has not been queued up to be loaded.
+    pub async fn wait_for_asset_id(
+        &self,
+        id: impl Into<UntypedAssetId>,
+    ) -> Result<(), WaitForAssetError> {
+        let id = id.into();
+        core::future::poll_fn(move |cx| self.wait_for_asset_id_poll_fn(cx, id)).await
+    }
+
+    /// Used by [`wait_for_asset_id`](AssetServer::wait_for_asset_id) in [`poll_fn`](core::future::poll_fn).
+    fn wait_for_asset_id_poll_fn(
+        &self,
+        cx: &mut core::task::Context<'_>,
+        id: UntypedAssetId,
+    ) -> Poll<Result<(), WaitForAssetError>> {
+        let infos = self.data.infos.read();
+
+        let Some(info) = infos.get(id) else {
+            return Poll::Ready(Err(WaitForAssetError::NotLoaded));
+        };
+
+        match (&info.load_state, &info.rec_dep_load_state) {
+            (LoadState::Loaded, RecursiveDependencyLoadState::Loaded) => Poll::Ready(Ok(())),
+            // Return an error immediately if the asset is not in the process of loading
+            (LoadState::NotLoaded, _) => Poll::Ready(Err(WaitForAssetError::NotLoaded)),
+            // If the asset is loading, leave our waker behind
+            (LoadState::Loading, _)
+            | (_, RecursiveDependencyLoadState::Loading)
+            | (LoadState::Loaded, RecursiveDependencyLoadState::NotLoaded) => {
+                // Check if our waker is already there
+                let has_waker = info
+                    .waiting_tasks
+                    .iter()
+                    .any(|waker| waker.will_wake(cx.waker()));
+
+                if has_waker {
+                    return Poll::Pending;
+                }
+
+                let mut infos = {
+                    // Must drop read-only guard to acquire write guard
+                    drop(infos);
+                    self.data.infos.write()
+                };
+
+                let Some(info) = infos.get_mut(id) else {
+                    return Poll::Ready(Err(WaitForAssetError::NotLoaded));
+                };
+
+                // If the load state changed while reacquiring the lock, immediately
+                // reawaken the task
+                let is_loading = matches!(
+                    (&info.load_state, &info.rec_dep_load_state),
+                    (LoadState::Loading, _)
+                        | (_, RecursiveDependencyLoadState::Loading)
+                        | (LoadState::Loaded, RecursiveDependencyLoadState::NotLoaded)
+                );
+
+                if !is_loading {
+                    cx.waker().wake_by_ref();
+                } else {
+                    // Leave our waker behind
+                    info.waiting_tasks.push(cx.waker().clone());
+                }
+
+                Poll::Pending
+            }
+            (LoadState::Failed(error), _) => {
+                Poll::Ready(Err(WaitForAssetError::Failed(error.clone())))
+            }
+            (_, RecursiveDependencyLoadState::Failed(error)) => {
+                Poll::Ready(Err(WaitForAssetError::DependencyFailed(error.clone())))
+            }
+        }
+    }
 }
 
 /// A system that manages internal [`AssetServer`] events, such as finalizing asset loads.
@@ -1359,6 +1488,11 @@ pub fn handle_internal_asset_events(world: &mut World) {
                         .get(&id.type_id())
                         .expect("Asset event sender should exist");
                     sender(world, id);
+                    if let Some(info) = infos.get_mut(id) {
+                        for waker in info.waiting_tasks.drain(..) {
+                            waker.wake();
+                        }
+                    }
                 }
                 InternalAssetEvent::Failed { id, path, error } => {
                     infos.process_asset_fail(id, error.clone());
@@ -1710,3 +1844,12 @@ impl core::fmt::Debug for AssetServer {
 /// This is appended to asset sources when loading a [`LoadedUntypedAsset`]. This provides a unique
 /// source for a given [`AssetPath`].
 const UNTYPED_SOURCE_SUFFIX: &str = "--untyped";
+
+/// An error when attempting to wait asynchronously for an [`Asset`] to load.
+#[derive(Error, Debug, Clone, Display)]
+pub enum WaitForAssetError {
+    #[display("tried to wait for an asset that is not being loaded")]
+    NotLoaded,
+    Failed(Arc<AssetLoadError>),
+    DependencyFailed(Arc<AssetLoadError>),
+}