Allow removing and reloading assets with live handles (#10785)

# Objective

Fixes #10444 

Currently manually removing an asset prevents it from being reloaded
while there are still active handles. Doing so will result in a panic,
because the storage entry has been marked as "empty / None" but the ID
is still assumed to be active by the asset server.

Patterns like `images.remove() -> asset_server.reload()` and
`images.remove() -> images.insert()` would fail if the handle was still
alive.

## Solution

Most of the groundwork for this was already laid in Bevy Asset V2. This
is largely just a matter of splitting out `remove` into two separate
operations:

* `remove_dropped`: remove the stored asset, invalidate the internal
Assets entry (preventing future insertions with the old id), and recycle
the id
* `remove_still_alive`: remove the stored asset, but leave the entry
otherwise untouched (and dont recycle the id).

`remove_still_alive` and `insert` can be called any number of times (in
any order) for an id until `remove_dropped` has been called, which will
invalidate the id.

From a user-facing perspective, there are no API changes and this is non
breaking. The public `Assets::remove` will internally call
`remove_still_alive`. `remove_dropped` can only be called by the
internal "handle management" system.

---

## Changelog

- Fix a bug preventing `Assets::remove` from blocking future inserts for
a specific `AssetIndex`.

crates/bevy_asset/src/assets.rs

@@ -1,4 +1,4 @@
-use crate as bevy_asset;
+use crate::{self as bevy_asset, LoadState};
 use crate::{Asset, AssetEvent, AssetHandleProvider, AssetId, AssetServer, Handle, UntypedHandle};
 use bevy_ecs::{
     prelude::EventWriter,
@@ -87,12 +87,11 @@ pub struct LoadedUntypedAsset {
 // PERF: do we actually need this to be an enum? Can we just use an "invalid" generation instead
 #[derive(Default)]
 enum Entry<A: Asset> {
+    /// None is an indicator that this entry does not have live handles.
     #[default]
     None,
-    Some {
+    /// Some is an indicator that there is a live handle active for the entry at this [`AssetIndex`]
-        value: Option<A>,
+    Some { value: Option<A>, generation: u32 },
-        generation: u32,
-    },
 }
 
 /// Stores [`Asset`] values in a Vec-like storage identified by [`AssetIndex`].
@@ -151,7 +150,26 @@ impl<A: Asset> DenseAssetStorage<A> {
     }
 
     /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
-    pub(crate) fn remove(&mut self, index: AssetIndex) -> Option<A> {
+    /// This will recycle the id and allow new entries to be inserted.
+    pub(crate) fn remove_dropped(&mut self, index: AssetIndex) -> Option<A> {
+        self.remove_internal(index, |dense_storage| {
+            dense_storage.storage[index.index as usize] = Entry::None;
+            dense_storage.allocator.recycle(index);
+        })
+    }
+
+    /// Removes the asset stored at the given `index` and returns it as [`Some`] (if the asset exists).
+    /// This will _not_ recycle the id. New values with the current ID can still be inserted. The ID will
+    /// not be reused until [`DenseAssetStorage::remove_dropped`] is called.
+    pub(crate) fn remove_still_alive(&mut self, index: AssetIndex) -> Option<A> {
+        self.remove_internal(index, |_| {})
+    }
+
+    fn remove_internal(
+        &mut self,
+        index: AssetIndex,
+        removed_action: impl FnOnce(&mut Self),
+    ) -> Option<A> {
         self.flush();
         let value = match &mut self.storage[index.index as usize] {
             Entry::None => return None,
@@ -166,8 +184,7 @@ impl<A: Asset> DenseAssetStorage<A> {
                 }
             }
         };
-        self.storage[index.index as usize] = Entry::None;
+        removed_action(self);
-        self.allocator.recycle(index);
         value
     }
 
@@ -393,11 +410,25 @@ impl<A: Asset> Assets<A> {
     pub fn remove_untracked(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
         let id: AssetId<A> = id.into();
         match id {
-            AssetId::Index { index, .. } => self.dense_storage.remove(index),
+            AssetId::Index { index, .. } => self.dense_storage.remove_still_alive(index),
             AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
         }
     }
 
+    /// Removes (and returns) the [`Asset`] with the given `id`, if its exists.
+    /// Note that this supports anything that implements `Into<AssetId<A>>`, which includes [`Handle`] and [`AssetId`].
+    pub(crate) fn remove_dropped(&mut self, id: impl Into<AssetId<A>>) -> Option<A> {
+        let id: AssetId<A> = id.into();
+        let result = match id {
+            AssetId::Index { index, .. } => self.dense_storage.remove_dropped(index),
+            AssetId::Uuid { uuid } => self.hash_map.remove(&uuid),
+        };
+        if result.is_some() {
+            self.queued_events.push(AssetEvent::Removed { id });
+        }
+        result
+    }
+
     /// Returns `true` if there are no assets in this collection.
     pub fn is_empty(&self) -> bool {
         self.dense_storage.is_empty() && self.hash_map.is_empty()
@@ -466,16 +497,21 @@ impl<A: Asset> Assets<A> {
         let mut not_ready = Vec::new();
         while let Ok(drop_event) = assets.handle_provider.drop_receiver.try_recv() {
             let id = drop_event.id;
-            if !assets.contains(id.typed()) {
-                not_ready.push(drop_event);
-                continue;
-            }
             if drop_event.asset_server_managed {
-                if infos.process_handle_drop(id.untyped(TypeId::of::<A>())) {
+                let untyped = id.untyped(TypeId::of::<A>());
-                    assets.remove(id.typed());
+                if let Some(info) = infos.get(untyped) {
+                    if info.load_state == LoadState::Loading
+                        || info.load_state == LoadState::NotLoaded
+                    {
+                        not_ready.push(drop_event);
+                        continue;
+                    }
+                }
+                if infos.process_handle_drop(untyped) {
+                    assets.remove_dropped(id.typed());
                 }
             } else {
-                assets.remove(id.typed());
+                assets.remove_dropped(id.typed());
             }
         }
         // TODO: this is _extremely_ inefficient find a better fix