despawn entities when switching worlds and some testbot fixes

2 files changed, 13 insertions, 0 deletions

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index c81f6f28..131aef16 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -555,6 +555,11 @@ impl Client {
         self.query::<Option<&T>>(&mut self.ecs.lock()).cloned()
     }
 
+    /// Get a resource from the ECS. This will clone the resource and return it.
+    pub fn resource<T: Resource + Clone>(&self) -> T {
+        self.ecs.lock().resource::<T>().clone()
+    }
+
     /// Get a required component for this client and call the given function.
     ///
     /// Similar to [`Self::component`], but doesn't clone the component since
diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
index a92423d7..c653e195 100644
--- a/azalea-client/src/disconnect.rs
+++ b/azalea-client/src/disconnect.rs
@@ -47,6 +47,7 @@ pub struct DisconnectEvent {
 pub fn remove_components_from_disconnected_players(
     mut commands: Commands,
     mut events: EventReader<DisconnectEvent>,
+    mut loaded_by_query: Query<&mut azalea_entity::LoadedBy>,
 ) {
     for DisconnectEvent { entity, .. } in events.read() {
         trace!("Got DisconnectEvent for {entity:?}");
@@ -62,6 +63,13 @@ pub fn remove_components_from_disconnected_players(
             .remove::<LocalPlayerEvents>();
         // note that we don't remove the client from the ECS, so if they decide
         // to reconnect they'll keep their state
+
+        // now we have to remove ourselves from the LoadedBy for every entity.
+        // in theory this could be inefficient if we have massive swarms... but in
+        // practice this is fine.
+        for mut loaded_by in &mut loaded_by_query.iter_mut() {
+            loaded_by.remove(entity);
+        }
     }
 }