properly remove from the EntityIdIndex component on despawn

7 files changed, 224 insertions, 52 deletions

diff --git a/azalea-client/src/test_simulation.rs b/azalea-client/src/test_simulation.rs
index 67dcb818..564db0cf 100644
--- a/azalea-client/src/test_simulation.rs
+++ b/azalea-client/src/test_simulation.rs
@@ -2,29 +2,32 @@ use std::{fmt::Debug, sync::Arc, time::Duration};
 
 use azalea_auth::game_profile::GameProfile;
 use azalea_buf::AzaleaWrite;
+use azalea_core::delta::PositionDelta8;
 use azalea_core::game_type::{GameMode, OptionalGameType};
-use azalea_core::position::ChunkPos;
+use azalea_core::position::{ChunkPos, Vec3};
 use azalea_core::resource_location::ResourceLocation;
 use azalea_core::tick::GameTick;
 use azalea_entity::metadata::PlayerMetadataBundle;
 use azalea_protocol::packets::common::CommonPlayerSpawnInfo;
+use azalea_protocol::packets::config::{ClientboundFinishConfiguration, ClientboundRegistryData};
 use azalea_protocol::packets::game::c_level_chunk_with_light::ClientboundLevelChunkPacketData;
 use azalea_protocol::packets::game::c_light_update::ClientboundLightUpdatePacketData;
 use azalea_protocol::packets::game::{
-    ClientboundLevelChunkWithLight, ClientboundLogin, ClientboundRespawn,
+    ClientboundAddEntity, ClientboundLevelChunkWithLight, ClientboundLogin, ClientboundRespawn,
 };
 use azalea_protocol::packets::{ConnectionProtocol, Packet, ProtocolPacket};
-use azalea_registry::DimensionType;
+use azalea_registry::{DimensionType, EntityKind};
 use azalea_world::palette::{PalettedContainer, PalettedContainerKind};
 use azalea_world::{Chunk, Instance, MinecraftEntityId, Section};
 use bevy_app::App;
 use bevy_ecs::{prelude::*, schedule::ExecutorKind};
 use parking_lot::{Mutex, RwLock};
-use simdnbt::owned::Nbt;
+use simdnbt::owned::{Nbt, NbtCompound, NbtTag};
 use tokio::task::JoinHandle;
 use tokio::{sync::mpsc, time::sleep};
 use uuid::Uuid;
 
+use crate::disconnect::DisconnectEvent;
 use crate::{
     ClientInformation, GameProfileComponent, InConfigState, InstanceHolder, LocalPlayerBundle,
     events::LocalPlayerEvents,
@@ -48,29 +51,49 @@ impl Simulation {
         let mut app = create_simulation_app();
         let mut entity = app.world_mut().spawn_empty();
         let (player, clear_outgoing_packets_receiver_task, incoming_packet_queue, rt) =
-            create_local_player_bundle(entity.id(), initial_connection_protocol);
+            create_local_player_bundle(entity.id(), ConnectionProtocol::Configuration);
         entity.insert(player);
 
         let entity = entity.id();
 
         tick_app(&mut app);
 
-        #[allow(clippy::single_match)]
-        match initial_connection_protocol {
-            ConnectionProtocol::Configuration => {
-                app.world_mut().entity_mut(entity).insert(InConfigState);
-                tick_app(&mut app);
-            }
-            _ => {}
-        }
+        // start in the config state
+        app.world_mut().entity_mut(entity).insert(InConfigState);
+        tick_app(&mut app);
 
-        Self {
+        let mut simulation = Self {
             app,
             entity,
             rt,
             incoming_packet_queue,
             clear_outgoing_packets_receiver_task,
+        };
+
+        #[allow(clippy::single_match)]
+        match initial_connection_protocol {
+            ConnectionProtocol::Configuration => {}
+            ConnectionProtocol::Game => {
+                simulation.receive_packet(ClientboundRegistryData {
+                    registry_id: ResourceLocation::new("minecraft:dimension_type"),
+                    entries: vec![(
+                        ResourceLocation::new("minecraft:overworld"),
+                        Some(NbtCompound::from_values(vec![
+                            ("height".into(), NbtTag::Int(384)),
+                            ("min_y".into(), NbtTag::Int(-64)),
+                        ])),
+                    )]
+                    .into_iter()
+                    .collect(),
+                });
+
+                simulation.receive_packet(ClientboundFinishConfiguration);
+                simulation.tick();
+            }
+            _ => unimplemented!("unsupported ConnectionProtocol {initial_connection_protocol:?}"),
         }
+
+        simulation
     }
 
     pub fn receive_packet<P: ProtocolPacket + Debug>(&mut self, packet: impl Packet<P>) {
@@ -98,6 +121,14 @@ impl Simulation {
             .chunks
             .get(&chunk_pos)
     }
+
+    pub fn disconnect(&mut self) {
+        // send DisconnectEvent
+        self.app.world_mut().send_event(DisconnectEvent {
+            entity: self.entity,
+            reason: None,
+        });
+    }
 }
 
 #[allow(clippy::type_complexity)]
@@ -270,3 +301,21 @@ pub fn make_basic_empty_chunk(
         light_data: ClientboundLightUpdatePacketData::default(),
     }
 }
+
+pub fn make_basic_add_entity(
+    entity_type: EntityKind,
+    id: i32,
+    position: impl Into<Vec3>,
+) -> ClientboundAddEntity {
+    ClientboundAddEntity {
+        id: id.into(),
+        uuid: Uuid::from_u128(1234),
+        entity_type,
+        position: position.into(),
+        x_rot: 0,
+        y_rot: 0,
+        y_head_rot: 0,
+        data: 0,
+        velocity: PositionDelta8::default(),
+    }
+}
diff --git a/azalea-client/tests/change_dimension_to_nether_and_back.rs b/azalea-client/tests/change_dimension_to_nether_and_back.rs
index 748ea713..969c5053 100644
--- a/azalea-client/tests/change_dimension_to_nether_and_back.rs
+++ b/azalea-client/tests/change_dimension_to_nether_and_back.rs
@@ -2,7 +2,7 @@ use azalea_client::{InConfigState, InGameState, test_simulation::*};
 use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
 use azalea_entity::LocalEntity;
 use azalea_protocol::packets::{
-    ConnectionProtocol,
+    ConnectionProtocol, Packet,
     config::{ClientboundFinishConfiguration, ClientboundRegistryData},
 };
 use azalea_registry::DimensionType;
@@ -12,6 +12,21 @@ use simdnbt::owned::{NbtCompound, NbtTag};
 
 #[test]
 fn test_change_dimension_to_nether_and_back() {
+    generic_test_change_dimension_to_nether_and_back(true);
+    generic_test_change_dimension_to_nether_and_back(false);
+}
+
+fn generic_test_change_dimension_to_nether_and_back(using_respawn: bool) {
+    let make_basic_login_or_respawn_packet = if using_respawn {
+        |dimension: DimensionType, instance_name: ResourceLocation| {
+            make_basic_respawn_packet(dimension, instance_name).into_variant()
+        }
+    } else {
+        |dimension: DimensionType, instance_name: ResourceLocation| {
+            make_basic_login_packet(dimension, instance_name).into_variant()
+        }
+    };
+
     let _ = tracing_subscriber::fmt::try_init();
 
     let mut simulation = Simulation::new(ConnectionProtocol::Configuration);
@@ -83,7 +98,7 @@ fn test_change_dimension_to_nether_and_back() {
     // NETHER
     //
 
-    simulation.receive_packet(make_basic_respawn_packet(
+    simulation.receive_packet(make_basic_login_or_respawn_packet(
         DimensionType::new_raw(2), // nether
         ResourceLocation::new("azalea:b"),
     ));
@@ -105,7 +120,7 @@ fn test_change_dimension_to_nether_and_back() {
     simulation
         .chunk(ChunkPos::new(0, 0))
         .expect("chunk should exist");
-    simulation.receive_packet(make_basic_respawn_packet(
+    simulation.receive_packet(make_basic_login_or_respawn_packet(
         DimensionType::new_raw(2), // nether
         ResourceLocation::new("minecraft:nether"),
     ));
diff --git a/azalea-client/tests/client_disconnect.rs b/azalea-client/tests/client_disconnect.rs
new file mode 100644
index 00000000..354ac788
--- /dev/null
+++ b/azalea-client/tests/client_disconnect.rs
@@ -0,0 +1,21 @@
+use azalea_client::test_simulation::*;
+use azalea_protocol::packets::ConnectionProtocol;
+use azalea_world::InstanceName;
+use bevy_log::tracing_subscriber;
+
+#[test]
+fn test_client_disconnect() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Game);
+
+    simulation.disconnect();
+    simulation.tick();
+
+    // make sure we're disconnected
+    let is_connected = simulation.has_component::<InstanceName>();
+    assert!(!is_connected);
+
+    // tick again to make sure nothing goes wrong
+    simulation.tick();
+}
diff --git a/azalea-client/tests/despawn_entities_when_changing_dimension.rs b/azalea-client/tests/despawn_entities_when_changing_dimension.rs
index 133506c9..3d7f2cb4 100644
--- a/azalea-client/tests/despawn_entities_when_changing_dimension.rs
+++ b/azalea-client/tests/despawn_entities_when_changing_dimension.rs
@@ -1,20 +1,14 @@
 use azalea_client::test_simulation::*;
-use azalea_core::{
-    delta::PositionDelta8,
-    position::{ChunkPos, Vec3},
-    resource_location::ResourceLocation,
-};
+use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
 use azalea_entity::metadata::Cow;
 use azalea_protocol::packets::{
     ConnectionProtocol,
     config::{ClientboundFinishConfiguration, ClientboundRegistryData},
-    game::ClientboundAddEntity,
 };
-use azalea_registry::DimensionType;
+use azalea_registry::{DimensionType, EntityKind};
 use bevy_ecs::query::With;
 use bevy_log::tracing_subscriber;
 use simdnbt::owned::{NbtCompound, NbtTag};
-use uuid::Uuid;
 
 #[test]
 fn test_despawn_entities_when_changing_dimension() {
@@ -59,17 +53,7 @@ fn test_despawn_entities_when_changing_dimension() {
     simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
     simulation.tick();
     // spawn a cow
-    simulation.receive_packet(ClientboundAddEntity {
-        id: 123.into(),
-        uuid: Uuid::from_u128(1234),
-        entity_type: azalea_registry::EntityKind::Cow,
-        position: Vec3::new(0., 64., 0.),
-        x_rot: 0,
-        y_rot: 0,
-        y_head_rot: 0,
-        data: 0,
-        velocity: PositionDelta8::default(),
-    });
+    simulation.receive_packet(make_basic_add_entity(EntityKind::Cow, 123, (0.5, 64., 0.5)));
     simulation.tick();
     // make sure it's spawned
     let mut cow_query = simulation.app.world_mut().query_filtered::<(), With<Cow>>();
diff --git a/azalea-client/tests/move_despawned_entity.rs b/azalea-client/tests/move_despawned_entity.rs
new file mode 100644
index 00000000..74f7e942
--- /dev/null
+++ b/azalea-client/tests/move_despawned_entity.rs
@@ -0,0 +1,51 @@
+use azalea_client::test_simulation::*;
+use azalea_core::{position::ChunkPos, resource_location::ResourceLocation};
+use azalea_entity::metadata::Cow;
+use azalea_protocol::packets::{ConnectionProtocol, game::ClientboundMoveEntityRot};
+use azalea_registry::{DimensionType, EntityKind};
+use azalea_world::MinecraftEntityId;
+use bevy_ecs::query::With;
+use bevy_log::tracing_subscriber;
+
+#[test]
+fn test_move_despawned_entity() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let mut simulation = Simulation::new(ConnectionProtocol::Game);
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(0),
+        ResourceLocation::new("azalea:overworld"),
+    ));
+
+    simulation.receive_packet(make_basic_empty_chunk(ChunkPos::new(0, 0), (384 + 64) / 16));
+    simulation.tick();
+    // spawn a cow
+    simulation.receive_packet(make_basic_add_entity(EntityKind::Cow, 123, (0.5, 64., 0.5)));
+    simulation.tick();
+
+    // make sure it's spawned
+    let mut cow_query = simulation.app.world_mut().query_filtered::<(), With<Cow>>();
+    let cow_iter = cow_query.iter(simulation.app.world());
+    assert_eq!(cow_iter.count(), 1, "cow should be despawned");
+
+    // despawn the cow by receiving a login packet
+    simulation.receive_packet(make_basic_login_packet(
+        DimensionType::new_raw(0),
+        ResourceLocation::new("azalea:overworld"),
+    ));
+    simulation.tick();
+
+    // make sure it's despawned
+    let mut cow_query = simulation.app.world_mut().query_filtered::<(), With<Cow>>();
+    let cow_iter = cow_query.iter(simulation.app.world());
+    assert_eq!(cow_iter.count(), 0, "cow should be despawned");
+
+    // send a move_entity_rot
+    simulation.receive_packet(ClientboundMoveEntityRot {
+        entity_id: MinecraftEntityId(123),
+        y_rot: 0,
+        x_rot: 0,
+        on_ground: false,
+    });
+    simulation.tick();
+}
diff --git a/azalea-entity/src/plugin/indexing.rs b/azalea-entity/src/plugin/indexing.rs
index fefecb06..4e0902e7 100644
--- a/azalea-entity/src/plugin/indexing.rs
+++ b/azalea-entity/src/plugin/indexing.rs
@@ -23,19 +23,6 @@ pub struct EntityUuidIndex {
     /// An index of entities by their UUIDs
     entity_by_uuid: HashMap<Uuid, Entity>,
 }
-
-/// An index of Minecraft entity IDs to Azalea ECS entities. This is a
-/// `Component` so local players can keep track of entity IDs independently from
-/// the instance.
-///
-/// If you need a per-instance instead of per-client version of this, you can
-/// use [`Instance::entity_by_id`].
-#[derive(Component, Default)]
-pub struct EntityIdIndex {
-    /// An index of entities by their MinecraftEntityId
-    entity_by_id: IntMap<MinecraftEntityId, Entity>,
-}
-
 impl EntityUuidIndex {
     pub fn new() -> Self {
         Self {
@@ -60,6 +47,19 @@ impl EntityUuidIndex {
     }
 }
 
+/// An index of Minecraft entity IDs to Azalea ECS entities. This is a
+/// `Component` so local players can keep track of entity IDs independently from
+/// the instance.
+///
+/// If you need a per-instance instead of per-client version of this, you can
+/// use [`Instance::entity_by_id`].
+#[derive(Component, Default)]
+pub struct EntityIdIndex {
+    /// An index of entities by their MinecraftEntityId
+    entity_by_id: IntMap<MinecraftEntityId, Entity>,
+    id_by_entity: HashMap<Entity, MinecraftEntityId>,
+}
+
 impl EntityIdIndex {
     pub fn get(&self, id: MinecraftEntityId) -> Option<Entity> {
         self.entity_by_id.get(&id).copied()
@@ -68,13 +68,31 @@ impl EntityIdIndex {
     pub fn contains_key(&self, id: MinecraftEntityId) -> bool {
         self.entity_by_id.contains_key(&id)
     }
+    pub fn contains_ecs_entity(&self, id: Entity) -> bool {
+        self.id_by_entity.contains_key(&id)
+    }
 
     pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
         self.entity_by_id.insert(id, entity);
+        self.id_by_entity.insert(entity, id);
     }
 
     pub fn remove(&mut self, id: MinecraftEntityId) -> Option<Entity> {
-        self.entity_by_id.remove(&id)
+        if let Some(entity) = self.entity_by_id.remove(&id) {
+            self.id_by_entity.remove(&entity);
+            Some(entity)
+        } else {
+            None
+        }
+    }
+
+    pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
+        if let Some(id) = self.id_by_entity.remove(&entity) {
+            self.entity_by_id.remove(&id);
+            Some(id)
+        } else {
+            None
+        }
     }
 }
 
@@ -154,6 +172,7 @@ pub fn remove_despawned_entities_from_indexes(
         ),
         (Changed<LoadedBy>, Without<LocalEntity>),
     >,
+    mut entity_id_index_query: Query<&mut EntityIdIndex>,
 ) {
     for (entity, uuid, minecraft_id, position, instance_name, loaded_by) in &query {
         let Some(instance_lock) = instance_container.get(instance_name) else {
@@ -207,6 +226,16 @@ pub fn remove_despawned_entities_from_indexes(
         if instance.entity_by_id.remove(minecraft_id).is_none() {
             warn!("Tried to remove entity {entity:?} from the id index but it was not there.");
         }
+
+        // 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."
+                );
+            }
+        }
+
         // and now remove the entity from the ecs
         commands.entity(entity).despawn();
         debug!("Despawned entity {entity:?} because it was not loaded by anything.");
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index 04bc94dc..c8546c2b 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -698,10 +698,33 @@ impl Swarm {
     ) {
         while let Some(event) = rx.recv().await {
             if rx.len() > 1_000 {
-                static WARNED: AtomicBool = AtomicBool::new(false);
-                if !WARNED.swap(true, atomic::Ordering::Relaxed) {
+                static WARNED_1_000: AtomicBool = AtomicBool::new(false);
+                if !WARNED_1_000.swap(true, atomic::Ordering::Relaxed) {
                     warn!("the client's Event channel has more than 1000 items!")
                 }
+
+                if rx.len() > 10_000 {
+                    static WARNED_10_000: AtomicBool = AtomicBool::new(false);
+                    if !WARNED_10_000.swap(true, atomic::Ordering::Relaxed) {
+                        warn!("the client's Event channel has more than 10,000 items!!")
+                    }
+
+                    if rx.len() > 100_000 {
+                        static WARNED_100_000: AtomicBool = AtomicBool::new(false);
+                        if !WARNED_100_000.swap(true, atomic::Ordering::Relaxed) {
+                            warn!("the client's Event channel has more than 100,000 items!!!")
+                        }
+
+                        if rx.len() > 1_000_000 {
+                            static WARNED_1_000_000: AtomicBool = AtomicBool::new(false);
+                            if !WARNED_1_000_000.swap(true, atomic::Ordering::Relaxed) {
+                                warn!(
+                                    "the client's Event channel has more than 1,000,000 items!!!! i sincerely hope no one ever sees this warning"
+                                )
+                            }
+                        }
+                    }
+                }
             }
 
             // we can't handle events here (since we can't copy the handler),