Asset re-loading while it's being deleted (#2011)

fixes #824
fixes #1956 

* marked asset loading methods as `must_use`
* fixed asset re-loading while asset is still loading to work as comment is describing code
* introduced a 1 frame delay between unused asset marking and actual asset removal

crates/bevy_asset/src/asset_server.rs

@@ -10,7 +10,7 @@ use bevy_log::warn;
 use bevy_tasks::TaskPool;
 use bevy_utils::{HashMap, Uuid};
 use crossbeam_channel::TryRecvError;
-use parking_lot::RwLock;
+use parking_lot::{Mutex, RwLock};
 use std::{collections::hash_map::Entry, path::Path, sync::Arc};
 use thiserror::Error;
 
@@ -45,6 +45,7 @@ fn format_missing_asset_ext(exts: &[String]) -> String {
 pub(crate) struct AssetRefCounter {
     pub(crate) channel: Arc<RefChangeChannel>,
     pub(crate) ref_counts: Arc<RwLock<HashMap<HandleId, usize>>>,
+    pub(crate) mark_unused_assets: Arc<Mutex<Vec<HandleId>>>,
 }
 
 pub struct AssetServerInternal {
@@ -224,6 +225,7 @@ impl AssetServer {
     /// The name of the asset folder is set inside the
     /// [`AssetServerSettings`](crate::AssetServerSettings) resource. The default name is
     /// `"assets"`.
+    #[must_use = "not using the returned strong handle may result in the unexpected release of the asset"]
     pub fn load<'a, T: Asset, P: Into<AssetPath<'a>>>(&self, path: P) -> Handle<T> {
         self.load_untyped(path).typed()
     }
@@ -253,11 +255,12 @@ impl AssetServer {
                 }),
             };
 
-            // if asset is already loaded (or is loading), don't load again
+            // if asset is already loaded or is loading, don't load again
             if !force
-                && source_info
+                && (source_info
                     .committed_assets
                     .contains(&asset_path_id.label_id())
+                    || source_info.load_state == LoadState::Loading)
             {
                 return Ok(asset_path_id);
             }
@@ -325,7 +328,8 @@ impl AssetServer {
             let type_uuid = loaded_asset.value.as_ref().unwrap().type_uuid();
             source_info.asset_types.insert(label_id, type_uuid);
             for dependency in loaded_asset.dependencies.iter() {
-                self.load_untyped(dependency.clone());
+                // another handle already exists created from the asset path
+                let _ = self.load_untyped(dependency.clone());
             }
         }
 
@@ -337,6 +341,7 @@ impl AssetServer {
         Ok(asset_path_id)
     }
 
+    #[must_use = "not using the returned strong handle may result in the unexpected release of the asset"]
     pub fn load_untyped<'a, P: Into<AssetPath<'a>>>(&self, path: P) -> HandleUntyped {
         let handle_id = self.load_untracked(path.into(), false);
         self.get_handle_untyped(handle_id)
@@ -356,6 +361,7 @@ impl AssetServer {
         asset_path.into()
     }
 
+    #[must_use = "not using the returned strong handles may result in the unexpected release of the assets"]
     pub fn load_folder<P: AsRef<Path>>(
         &self,
         path: P,
@@ -385,33 +391,14 @@ impl AssetServer {
     }
 
     pub fn free_unused_assets(&self) {
-        let receiver = &self.server.asset_ref_counter.channel.receiver;
+        let mut potential_frees = self.server.asset_ref_counter.mark_unused_assets.lock();
-        let mut ref_counts = self.server.asset_ref_counter.ref_counts.write();
-        let asset_sources = self.server.asset_sources.read();
-        let mut potential_frees = Vec::new();
-        loop {
-            let ref_change = match receiver.try_recv() {
-                Ok(ref_change) => ref_change,
-                Err(TryRecvError::Empty) => break,
-                Err(TryRecvError::Disconnected) => panic!("RefChange channel disconnected."),
-            };
-            match ref_change {
-                RefChange::Increment(handle_id) => *ref_counts.entry(handle_id).or_insert(0) += 1,
-                RefChange::Decrement(handle_id) => {
-                    let entry = ref_counts.entry(handle_id).or_insert(0);
-                    *entry -= 1;
-                    if *entry == 0 {
-                        potential_frees.push(handle_id);
-                    }
-                }
-            }
-        }
 
         if !potential_frees.is_empty() {
+            let ref_counts = self.server.asset_ref_counter.ref_counts.read();
+            let asset_sources = self.server.asset_sources.read();
             let asset_lifecycles = self.server.asset_lifecycles.read();
-            for potential_free in potential_frees {
+            for potential_free in potential_frees.drain(..) {
-                if let Some(i) = ref_counts.get(&potential_free).cloned() {
+                if let Some(&0) = ref_counts.get(&potential_free) {
-                    if i == 0 {
                     let type_uuid = match potential_free {
                         HandleId::Id(type_uuid, _) => Some(type_uuid),
                         HandleId::AssetPathId(id) => asset_sources
@@ -428,6 +415,32 @@ impl AssetServer {
             }
         }
     }
+
+    pub fn mark_unused_assets(&self) {
+        let receiver = &self.server.asset_ref_counter.channel.receiver;
+        let mut ref_counts = self.server.asset_ref_counter.ref_counts.write();
+        let mut potential_frees = None;
+        loop {
+            let ref_change = match receiver.try_recv() {
+                Ok(ref_change) => ref_change,
+                Err(TryRecvError::Empty) => break,
+                Err(TryRecvError::Disconnected) => panic!("RefChange channel disconnected."),
+            };
+            match ref_change {
+                RefChange::Increment(handle_id) => *ref_counts.entry(handle_id).or_insert(0) += 1,
+                RefChange::Decrement(handle_id) => {
+                    let entry = ref_counts.entry(handle_id).or_insert(0);
+                    *entry -= 1;
+                    if *entry == 0 {
+                        potential_frees
+                            .get_or_insert_with(|| {
+                                self.server.asset_ref_counter.mark_unused_assets.lock()
+                            })
+                            .push(handle_id);
+                    }
+                }
+            }
+        }
     }
 
     fn create_assets_in_load_context(&self, load_context: &mut LoadContext) {
@@ -504,6 +517,7 @@ impl AssetServer {
 
 pub fn free_unused_assets_system(asset_server: Res<AssetServer>) {
     asset_server.free_unused_assets();
+    asset_server.mark_unused_assets();
 }
 
 #[cfg(test)]