fix entity deindexing happening at the wrong time

4 files changed, 59 insertions, 25 deletions

diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 6bb93178..ae1f3deb 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -703,7 +703,10 @@ impl GamePacketHandler<'_> {
                 let mut spawned =
                     commands.spawn((entity_id, LoadedBy(HashSet::from([self.player])), bundle));
                 let ecs_entity: Entity = spawned.id();
-                debug!("spawned entity {ecs_entity:?} with id {entity_id:?}");
+                debug!(
+                    "spawned entity {ecs_entity:?} with id {entity_id:?} at {pos:?}",
+                    pos = p.position
+                );
 
                 azalea_entity::indexing::add_entity_to_indexes(
                     entity_id,
@@ -918,11 +921,9 @@ impl GamePacketHandler<'_> {
 
                 debug!("Got move entity pos packet {p:?}");
 
-                let Some(entity) = entity_id_index.get(p.entity_id) else {
-                    debug!(
-                        "Got move entity pos packet for unknown entity id {}",
-                        p.entity_id
-                    );
+                let entity_id = p.entity_id;
+                let Some(entity) = entity_id_index.get(entity_id) else {
+                    debug!("Got move entity pos packet for unknown entity id {entity_id}");
                     return;
                 };
 
@@ -944,6 +945,11 @@ impl GamePacketHandler<'_> {
                         if new_pos != **position {
                             **position = new_pos;
                         }
+
+                        trace!(
+                            "Applied movement update for {entity_id} / {entity}",
+                            entity = entity_mut.id()
+                        );
                     },
                 ));
             },
@@ -1071,7 +1077,7 @@ impl GamePacketHandler<'_> {
                 for &id in &p.entity_ids {
                     let Some(entity) = entity_id_index.remove(id) else {
                         debug!(
-                            "Tried to remove entity with id {id} but it wasn't in the EntityIdIndex"
+                            "Tried to remove entity with id {id} but it wasn't in the EntityIdIndex. This may be expected on certain server setups (like if they're using VeryManyPlayers)."
                         );
                         continue;
                     };
@@ -1082,7 +1088,7 @@ impl GamePacketHandler<'_> {
                         continue;
                     };
 
-                    // the [`remove_despawned_entities_from_indexes`] system will despawn the entity
+                    // the `remove_despawned_entities_from_indexes` system will despawn the entity
                     // if it's not loaded by anything anymore
 
                     // also we can't just ecs.despawn because if we're in a swarm then the entity
diff --git a/azalea-client/src/plugins/packet/mod.rs b/azalea-client/src/plugins/packet/mod.rs
index 05f64e19..362154cc 100644
--- a/azalea-client/src/plugins/packet/mod.rs
+++ b/azalea-client/src/plugins/packet/mod.rs
@@ -1,4 +1,4 @@
-use azalea_entity::{EntityUpdateSet, metadata::Health};
+use azalea_entity::metadata::Health;
 use bevy_app::{App, First, Plugin, PreUpdate, Update};
 use bevy_ecs::{
     prelude::*,
@@ -46,9 +46,7 @@ impl Plugin for PacketPlugin {
         .add_systems(
             PreUpdate,
             (
-                game::process_packet_events
-                    // we want to index and deindex right after
-                    .before(EntityUpdateSet::Deindex),
+                game::process_packet_events,
                 config::process_packet_events,
                 login::handle_send_packet_event,
                 login::process_packet_events,
diff --git a/azalea-entity/src/plugin/indexing.rs b/azalea-entity/src/plugin/indexing.rs
index 6725197e..85f0f76c 100644
--- a/azalea-entity/src/plugin/indexing.rs
+++ b/azalea-entity/src/plugin/indexing.rs
@@ -1,6 +1,9 @@
 //! Stuff related to entity indexes and keeping track of entities in the world.
 
-use std::{collections::HashMap, fmt::Debug};
+use std::{
+    collections::{HashMap, HashSet},
+    fmt::Debug,
+};
 
 use azalea_core::position::ChunkPos;
 use azalea_world::{Instance, InstanceContainer, InstanceName, MinecraftEntityId};
@@ -12,7 +15,7 @@ use bevy_ecs::{
 };
 use derive_more::{Deref, DerefMut};
 use nohash_hasher::IntMap;
-use tracing::{debug, warn};
+use tracing::{debug, trace, warn};
 use uuid::Uuid;
 
 use super::LoadedBy;
@@ -75,22 +78,39 @@ impl EntityIdIndex {
     pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
         self.entity_by_id.insert(id, entity);
         self.id_by_entity.insert(entity, id);
+        trace!("Inserted {id} -> {entity:?} into a client's EntityIdIndex");
     }
 
     pub fn remove(&mut self, id: MinecraftEntityId) -> Option<Entity> {
         if let Some(entity) = self.entity_by_id.remove(&id) {
+            trace!(
+                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove)"
+            );
             self.id_by_entity.remove(&entity);
             Some(entity)
         } else {
+            trace!(
+                "Failed to remove {id} from a client's EntityIdIndex (using EntityIdIndex::remove)"
+            );
             None
         }
     }
 
     pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
         if let Some(id) = self.id_by_entity.remove(&entity) {
+            trace!(
+                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity)."
+            );
             self.entity_by_id.remove(&id);
             Some(id)
         } else {
+            // this is expected to happen when despawning entities if it was already
+            // despawned for another reason (like because the client received a
+            // remove_entities packet, or if we're in a shared instance where entity ids are
+            // different for each client)
+            trace!(
+                "Failed to remove {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity). This may be expected behavior."
+            );
             None
         }
     }
@@ -132,6 +152,7 @@ pub fn update_entity_chunk_positions(
                     .entry(new_chunk)
                     .or_default()
                     .insert(entity);
+                trace!("Entity {entity:?} moved from {old_chunk:?} to {new_chunk:?}");
             }
         }
     }
@@ -208,9 +229,22 @@ pub fn remove_despawned_entities_from_indexes(
                         instance.entities_by_chunk.remove(&chunk);
                     }
                 } else {
-                    warn!(
-                        "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there."
-                    );
+                    // search all the other chunks for it :(
+                    let mut found_in_other_chunks = HashSet::new();
+                    for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
+                        if entities_in_other_chunk.remove(&entity) {
+                            found_in_other_chunks.insert(other_chunk);
+                        }
+                    }
+                    if found_in_other_chunks.is_empty() {
+                        warn!(
+                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there or in any other chunks."
+                        );
+                    } else {
+                        warn!(
+                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there. Found in and removed from other chunk(s): {found_in_other_chunks:?}"
+                        );
+                    }
                 }
             }
             _ => {
@@ -225,17 +259,13 @@ pub fn remove_despawned_entities_from_indexes(
         }
         if instance.entity_by_id.remove(minecraft_id).is_none() {
             debug!(
-                "Tried to remove entity {entity:?} from the id index but it was not there. This may be expected if you're in a shared instance."
+                "Tried to remove entity {entity:?} from the per-instance entity id index but it was not there. This may be expected if you're in a shared instance."
             );
         }
 
         // remove it from every client's EntityIdIndex
         for mut entity_id_index in entity_id_index_query.iter_mut() {
-            if entity_id_index.remove_by_ecs_entity(entity).is_none() {
-                debug!(
-                    "Tried to remove entity {entity:?} from the id index but it was not there. This may be expected if you're in a shared instance."
-                );
-            }
+            entity_id_index.remove_by_ecs_entity(entity);
         }
 
         // and now remove the entity from the ecs
diff --git a/azalea-entity/src/plugin/mod.rs b/azalea-entity/src/plugin/mod.rs
index b6ae369f..115b25e4 100644
--- a/azalea-entity/src/plugin/mod.rs
+++ b/azalea-entity/src/plugin/mod.rs
@@ -9,7 +9,7 @@ use azalea_core::{
     tick::GameTick,
 };
 use azalea_world::{InstanceContainer, InstanceName, MinecraftEntityId};
-use bevy_app::{App, Plugin, PreUpdate, Update};
+use bevy_app::{App, Plugin, PostUpdate, Update};
 use bevy_ecs::prelude::*;
 use derive_more::{Deref, DerefMut};
 use indexing::EntityUuidIndex;
@@ -39,7 +39,7 @@ impl Plugin for EntityPlugin {
         // modified during update
         // despawned post-update (done by this plugin)
         app.add_systems(
-            PreUpdate,
+            PostUpdate,
             indexing::remove_despawned_entities_from_indexes.in_set(EntityUpdateSet::Deindex),
         )
         .add_systems(