Allow AssetServer::load to acquire a guard item. (#13051)

# Objective

Supercedes #12881 . Added a simple implementation that allows the user
to react to multiple asset loads both synchronously and asynchronously.

## Solution

Added `load_acquire`, that holds an item and drops it when loading is
finished or failed.

When used synchronously 

Hold an `Arc<()>`, check for `Arc::strong_count() == 1` when all loading
completed.

When used asynchronously 

Hold a `SemaphoreGuard`, await on `acquire_all` for completion.

This implementation has more freedom than the original in my opinion.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Zachary Harrold <zac@harrold.com.au>

Cargo.toml

@@ -357,6 +357,7 @@ futures-lite = "2.0.1"
 crossbeam-channel = "0.5.0"
 argh = "0.1.12"
 thiserror = "1.0"
+event-listener = "5.3.0"
 
 [[example]]
 name = "hello_world"
@@ -1413,6 +1414,18 @@ description = "How to configure the texture to repeat instead of the default cla
 category = "Assets"
 wasm = true
 
+# Assets
+[[example]]
+name = "multi_asset_sync"
+path = "examples/asset/multi_asset_sync.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.multi_asset_sync]
+name = "Mult-asset synchronization"
+description = "Demonstrates how to wait for multiple assets to be loaded."
+category = "Assets"
+wasm = true
+
 # Async Tasks
 [[example]]
 name = "async_compute"

crates/bevy_asset/src/loader_builders.rs

@@ -117,7 +117,7 @@ impl<'ctx, 'builder> NestedLoader<'ctx, 'builder> {
         let handle = if self.load_context.should_load_dependencies {
             self.load_context
                 .asset_server
-                .load_with_meta_transform(path, self.meta_transform)
+                .load_with_meta_transform(path, self.meta_transform, ())
         } else {
             self.load_context
                 .asset_server

crates/bevy_asset/src/server/mod.rs

@@ -270,7 +270,31 @@ impl AssetServer {
     /// The asset load will fail and an error will be printed to the logs if the asset stored at `path` is not of type `A`.
     #[must_use = "not using the returned strong handle may result in the unexpected release of the asset"]
     pub fn load<'a, A: Asset>(&self, path: impl Into<AssetPath<'a>>) -> Handle<A> {
-        self.load_with_meta_transform(path, None)
+        self.load_with_meta_transform(path, None, ())
+    }
+
+    /// Begins loading an [`Asset`] of type `A` stored at `path` while holding a guard item.
+    /// The guard item is dropped when either the asset is loaded or loading has failed.
+    ///
+    /// This function returns a "strong" [`Handle`]. When the [`Asset`] is loaded (and enters [`LoadState::Loaded`]), it will be added to the
+    /// associated [`Assets`] resource.
+    ///
+    /// The guard item should notify the caller in its [`Drop`] implementation. See example `multi_asset_sync`.
+    /// Synchronously this can be a [`Arc<AtomicU32>`] that decrements its counter, asynchronously this can be a `Barrier`.
+    /// This function only guarantees the asset referenced by the [`Handle`] is loaded. If your asset is separated into
+    /// multiple files, sub-assets referenced by the main asset might still be loading, depend on the implementation of the [`AssetLoader`].
+    ///
+    /// Additionally, you can check the asset's load state by reading [`AssetEvent`] events, calling [`AssetServer::load_state`], or checking
+    /// the [`Assets`] storage to see if the [`Asset`] exists yet.
+    ///
+    /// The asset load will fail and an error will be printed to the logs if the asset stored at `path` is not of type `A`.
+    #[must_use = "not using the returned strong handle may result in the unexpected release of the asset"]
+    pub fn load_acquire<'a, A: Asset, G: Send + Sync + 'static>(
+        &self,
+        path: impl Into<AssetPath<'a>>,
+        guard: G,
+    ) -> Handle<A> {
+        self.load_with_meta_transform(path, None, guard)
     }
 
     /// Begins loading an [`Asset`] of type `A` stored at `path`. The given `settings` function will override the asset's
@@ -282,13 +306,33 @@ impl AssetServer {
         path: impl Into<AssetPath<'a>>,
         settings: impl Fn(&mut S) + Send + Sync + 'static,
     ) -> Handle<A> {
-        self.load_with_meta_transform(path, Some(loader_settings_meta_transform(settings)))
+        self.load_with_meta_transform(path, Some(loader_settings_meta_transform(settings)), ())
     }
 
-    pub(crate) fn load_with_meta_transform<'a, A: Asset>(
+    /// Begins loading an [`Asset`] of type `A` stored at `path` while holding a guard item.
+    /// The guard item is dropped when either the asset is loaded or loading has failed.
+    ///
+    /// This function only guarantees the asset referenced by the [`Handle`] is loaded. If your asset is separated into
+    /// multiple files, sub-assets referenced by the main asset might still be loading, depend on the implementation of the [`AssetLoader`].
+    ///
+    /// The given `settings` function will override the asset's
+    /// [`AssetLoader`] settings. The type `S` _must_ match the configured [`AssetLoader::Settings`] or `settings` changes
+    /// will be ignored and an error will be printed to the log.
+    #[must_use = "not using the returned strong handle may result in the unexpected release of the asset"]
+    pub fn load_acquire_with_settings<'a, A: Asset, S: Settings, G: Send + Sync + 'static>(
+        &self,
+        path: impl Into<AssetPath<'a>>,
+        settings: impl Fn(&mut S) + Send + Sync + 'static,
+        guard: G,
+    ) -> Handle<A> {
+        self.load_with_meta_transform(path, Some(loader_settings_meta_transform(settings)), guard)
+    }
+
+    pub(crate) fn load_with_meta_transform<'a, A: Asset, G: Send + Sync + 'static>(
         &self,
         path: impl Into<AssetPath<'a>>,
         meta_transform: Option<MetaTransform>,
+        guard: G,
     ) -> Handle<A> {
         let path = path.into().into_owned();
         let (handle, should_load) = self.data.infos.write().get_or_create_path_handle::<A>(
@@ -305,6 +349,7 @@ impl AssetServer {
                     if let Err(err) = server.load_internal(owned_handle, path, false, None).await {
                         error!("{}", err);
                     }
+                    drop(guard);
                 })
                 .detach();
         }

examples/README.md

@@ -214,6 +214,7 @@ Example | Description
 [Embedded Asset](../examples/asset/embedded_asset.rs) | Embed an asset in the application binary and load it
 [Extra asset source](../examples/asset/extra_source.rs) | Load an asset from a non-standard asset source
 [Hot Reloading of Assets](../examples/asset/hot_asset_reloading.rs) | Demonstrates automatic reloading of assets when modified on disk
+[Mult-asset synchronization](../examples/asset/multi_asset_sync.rs) | Demonstrates how to wait for multiple assets to be loaded.
 [Repeated texture configuration](../examples/asset/repeated_texture.rs) | How to configure the texture to repeat instead of the default clamp to edges
 
 ## Async Tasks

examples/asset/multi_asset_sync.rs

@@ -0,0 +1,305 @@
+//! This example illustrates how to wait for multiple assets to be loaded.
+
+use std::{
+    f32::consts::PI,
+    sync::{
+        atomic::{AtomicBool, AtomicU32, Ordering},
+        Arc,
+    },
+};
+
+use bevy::{gltf::Gltf, prelude::*, tasks::AsyncComputeTaskPool};
+use event_listener::Event;
+use futures_lite::Future;
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .init_state::<LoadingState>()
+        .insert_resource(AmbientLight {
+            color: Color::WHITE,
+            brightness: 2000.,
+        })
+        .add_systems(Startup, setup_assets)
+        .add_systems(Startup, setup_scene)
+        .add_systems(Startup, setup_ui)
+        // This showcases how to wait for assets using sync code.
+        // This approach polls a value in a system.
+        .add_systems(Update, wait_on_load.run_if(assets_loaded))
+        // This showcases how to wait for assets using async
+        // by spawning a `Future` in `AsyncComputeTaskPool`.
+        .add_systems(
+            Update,
+            get_async_loading_state.run_if(in_state(LoadingState::Loading)),
+        )
+        // This showcases how to react to asynchronous world mutation synchronously.
+        .add_systems(
+            OnExit(LoadingState::Loading),
+            despawn_loading_state_entities,
+        )
+        .run();
+}
+
+/// [`States`] of asset loading.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, States, Default)]
+pub enum LoadingState {
+    /// Is loading.
+    #[default]
+    Loading,
+    /// Loading completed.
+    Loaded,
+}
+
+/// Holds a bunch of [`Gltf`]s that takes time to load.
+#[derive(Debug, Resource)]
+pub struct OneHundredThings([Handle<Gltf>; 100]);
+
+/// This is required to support both sync and async.
+///
+/// For sync only the easiest implementation is
+/// [`Arc<()>`] and use [`Arc::strong_count`] for completion.
+/// [`Arc<Atomic*>`] is a more robust alternative.
+#[derive(Debug, Resource, Deref)]
+pub struct AssetBarrier(Arc<AssetBarrierInner>);
+
+/// This guard is to be acquired by [`AssetServer::load_acquire`]
+/// and dropped once finished.
+#[derive(Debug, Deref)]
+pub struct AssetBarrierGuard(Arc<AssetBarrierInner>);
+
+/// Tracks how many guards are remaining.
+#[derive(Debug, Resource)]
+pub struct AssetBarrierInner {
+    count: AtomicU32,
+    /// This can be omitted if async is not needed.
+    notify: Event,
+}
+
+/// State of loading asynchronously.
+#[derive(Debug, Resource)]
+pub struct AsyncLoadingState(Arc<AtomicBool>);
+
+/// Entities that are to be removed once loading finished
+#[derive(Debug, Component)]
+pub struct Loading;
+
+/// Marker for the "Loading..." Text component.
+#[derive(Debug, Component)]
+pub struct LoadingText;
+
+impl AssetBarrier {
+    /// Create an [`AssetBarrier`] with a [`AssetBarrierGuard`].
+    pub fn new() -> (AssetBarrier, AssetBarrierGuard) {
+        let inner = Arc::new(AssetBarrierInner {
+            count: AtomicU32::new(1),
+            notify: Event::new(),
+        });
+        (AssetBarrier(inner.clone()), AssetBarrierGuard(inner))
+    }
+
+    /// Returns true if all [`AssetBarrierGuard`] is dropped.
+    pub fn is_ready(&self) -> bool {
+        self.count.load(Ordering::Acquire) == 0
+    }
+
+    /// Wait for all [`AssetBarrierGuard`]s to be dropped asynchronously.
+    pub fn wait_async(&self) -> impl Future<Output = ()> + 'static {
+        let shared = self.0.clone();
+        async move {
+            loop {
+                // Acquire an event listener.
+                let listener = shared.notify.listen();
+                // If all barrier guards are dropped, return
+                if shared.count.load(Ordering::Acquire) == 0 {
+                    return;
+                }
+                // Wait for the last barrier guard to notify us
+                listener.await;
+            }
+        }
+    }
+}
+
+// Increment count on clone.
+impl Clone for AssetBarrierGuard {
+    fn clone(&self) -> Self {
+        self.count.fetch_add(1, Ordering::AcqRel);
+        AssetBarrierGuard(self.0.clone())
+    }
+}
+
+// Decrement count on drop.
+impl Drop for AssetBarrierGuard {
+    fn drop(&mut self) {
+        let prev = self.count.fetch_sub(1, Ordering::AcqRel);
+        if prev == 1 {
+            // Notify all listeners if count reaches 0.
+            self.notify.notify(usize::MAX);
+        }
+    }
+}
+
+fn setup_assets(mut commands: Commands, asset_server: Res<AssetServer>) {
+    let (barrier, guard) = AssetBarrier::new();
+    commands.insert_resource(OneHundredThings(std::array::from_fn(|i| match i % 5 {
+        0 => asset_server.load_acquire("models/GolfBall/GolfBall.glb", guard.clone()),
+        1 => asset_server.load_acquire("models/AlienCake/alien.glb", guard.clone()),
+        2 => asset_server.load_acquire("models/AlienCake/cakeBirthday.glb", guard.clone()),
+        3 => asset_server.load_acquire("models/FlightHelmet/FlightHelmet.gltf", guard.clone()),
+        4 => asset_server.load_acquire("models/torus/torus.gltf", guard.clone()),
+        _ => unreachable!(),
+    })));
+    let future = barrier.wait_async();
+    commands.insert_resource(barrier);
+
+    let loading_state = Arc::new(AtomicBool::new(false));
+    commands.insert_resource(AsyncLoadingState(loading_state.clone()));
+
+    // await the `AssetBarrierFuture`.
+    AsyncComputeTaskPool::get()
+        .spawn(async move {
+            future.await;
+            // Notify via `AsyncLoadingState`
+            loading_state.store(true, Ordering::Release);
+        })
+        .detach();
+}
+
+fn setup_ui(mut commands: Commands) {
+    // Display the result of async loading.
+    commands
+        .spawn(NodeBundle {
+            style: Style {
+                width: Val::Percent(100.),
+                height: Val::Percent(100.),
+                justify_content: JustifyContent::End,
+
+                ..default()
+            },
+            ..default()
+        })
+        .with_children(|b| {
+            b.spawn((
+                TextBundle {
+                    text: Text {
+                        sections: vec![TextSection {
+                            value: "Loading...".to_owned(),
+                            style: TextStyle {
+                                font_size: 64.0,
+                                color: Color::BLACK,
+                                ..Default::default()
+                            },
+                        }],
+                        justify: JustifyText::Right,
+                        ..Default::default()
+                    },
+                    ..Default::default()
+                },
+                LoadingText,
+            ));
+        });
+}
+
+fn setup_scene(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+) {
+    // Camera
+    commands.spawn(Camera3dBundle {
+        transform: Transform::from_xyz(10.0, 10.0, 15.0)
+            .looking_at(Vec3::new(0.0, 0.0, 0.0), Vec3::Y),
+        ..default()
+    });
+
+    // Light
+    commands.spawn(DirectionalLightBundle {
+        transform: Transform::from_rotation(Quat::from_euler(EulerRot::ZYX, 0.0, 1.0, -PI / 4.)),
+        directional_light: DirectionalLight {
+            shadows_enabled: true,
+            ..default()
+        },
+        ..default()
+    });
+
+    // Plane
+    commands.spawn((
+        PbrBundle {
+            mesh: meshes.add(Plane3d::default().mesh().size(50000.0, 50000.0)),
+            material: materials.add(Color::srgb(0.7, 0.2, 0.2)),
+            ..default()
+        },
+        Loading,
+    ));
+}
+
+// A run condition for all assets being loaded.
+fn assets_loaded(barrier: Option<Res<AssetBarrier>>) -> bool {
+    // If our barrier isn't ready, return early and wait another cycle
+    barrier.map(|b| b.is_ready()) == Some(true)
+}
+
+// This showcases how to wait for assets using sync code and systems.
+//
+// This function only runs if `assets_loaded` returns true.
+fn wait_on_load(
+    mut commands: Commands,
+    foxes: Res<OneHundredThings>,
+    gltfs: Res<Assets<Gltf>>,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+) {
+    // Change color of plane to green
+    commands.spawn((PbrBundle {
+        mesh: meshes.add(Plane3d::default().mesh().size(50000.0, 50000.0)),
+        material: materials.add(Color::srgb(0.3, 0.5, 0.3)),
+        transform: Transform::from_translation(Vec3::Z * -0.01),
+        ..default()
+    },));
+
+    // Spawn our scenes.
+    for i in 0..10 {
+        for j in 0..10 {
+            let index = i * 10 + j;
+            let position = Vec3::new(i as f32 - 5.0, 0.0, j as f32 - 5.0);
+            // All gltfs must exist because this is guarded by the `AssetBarrier`.
+            let gltf = gltfs.get(&foxes.0[index]).unwrap();
+            let scene = gltf.scenes.first().unwrap().clone();
+            commands.spawn(SceneBundle {
+                scene,
+                transform: Transform::from_translation(position),
+                ..Default::default()
+            });
+        }
+    }
+}
+
+// This showcases how to wait for assets using async.
+fn get_async_loading_state(
+    state: Res<AsyncLoadingState>,
+    mut next_loading_state: ResMut<NextState<LoadingState>>,
+    mut text: Query<&mut Text, With<LoadingText>>,
+) {
+    // Load the value written by the `Future`.
+    let is_loaded = state.0.load(Ordering::Acquire);
+
+    // If loaded, change the state.
+    if is_loaded {
+        next_loading_state.set(LoadingState::Loaded);
+        if let Ok(mut text) = text.get_single_mut() {
+            "Loaded!".clone_into(&mut text.sections[0].value);
+        }
+    }
+}
+
+// This showcases how to react to asynchronous world mutations synchronously.
+fn despawn_loading_state_entities(mut commands: Commands, loading: Query<Entity, With<Loading>>) {
+    // Despawn entities in the loading phase.
+    for entity in loading.iter() {
+        commands.entity(entity).despawn_recursive();
+    }
+
+    // Despawn resources used in the loading phase.
+    commands.remove_resource::<AssetBarrier>();
+    commands.remove_resource::<AsyncLoadingState>();
+}