optimizations at high entity counts

3 files changed, 279 insertions, 76 deletions

diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 9a1c022b..9f16c204 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -9,7 +9,7 @@ use azalea_core::{
 };
 use azalea_entity::{
     ActiveEffects, Dead, EntityBundle, EntityKindComponent, HasClientLoaded, LoadedBy, LocalEntity,
-    LookDirection, Physics, PlayerAbilities, Position, RelativeEntityUpdate,
+    LookDirection, Physics, PlayerAbilities, Position,
     indexing::{EntityIdIndex, EntityUuidIndex},
     inventory::Inventory,
     metadata::{Health, apply_metadata},
@@ -36,7 +36,10 @@ use crate::{
     inventory::{ClientsideCloseContainerEvent, MenuOpenedEvent, SetContainerContentEvent},
     local_player::{Experience, Hunger, LocalGameMode, TabList, WorldHolder},
     movement::{KnockbackData, KnockbackEvent},
-    packet::{as_system, declare_packet_handlers},
+    packet::{
+        as_system, declare_packet_handlers,
+        relative_updates::{EntityUpdateQuery, RelativeEntityUpdate, should_apply_entity_update},
+    },
     player::{GameProfileComponent, PlayerInfo},
     tick_counter::TicksConnected,
 };
@@ -716,36 +719,37 @@ impl GamePacketHandler<'_> {
         // vanilla servers use this packet for knockback, but note that the Explode
         // packet is also sometimes used by servers for knockback
 
-        as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
-            self.ecs,
-            |(mut commands, query)| {
-                let (entity_id_index, world_holder) = query.get(self.player).unwrap();
-
-                let Some(entity) = entity_id_index.get_by_minecraft_entity(p.id) else {
-                    // note that this log (and some other ones like the one in RemoveEntities)
-                    // sometimes happens when killing mobs. it seems to be a vanilla bug, which is
-                    // why it's a debug log instead of a warning
-                    debug!(
-                        "Got set entity motion packet for unknown entity id {}",
-                        p.id
-                    );
-                    return;
-                };
-
-                // this is to make sure the same entity velocity update doesn't get sent
-                // multiple times when in swarms
+        as_system::<(
+            Commands,
+            Query<(&EntityIdIndex, &WorldHolder)>,
+            EntityUpdateQuery,
+        )>(self.ecs, |(mut commands, query, entity_update_query)| {
+            let (entity_id_index, world_holder) = query.get(self.player).unwrap();
 
-                let data = KnockbackData::Set(p.delta.to_vec3());
+            let Some(entity) = entity_id_index.get_by_minecraft_entity(p.id) else {
+                // note that this log (and some other ones like the one in RemoveEntities)
+                // sometimes happens when killing mobs. it seems to be a vanilla bug, which is
+                // why it's a debug log instead of a warning
+                debug!(
+                    "Got set entity motion packet for unknown entity id {}",
+                    p.id
+                );
+                return;
+            };
 
-                commands.entity(entity).queue(RelativeEntityUpdate::new(
-                    world_holder.partial.clone(),
-                    move |entity_mut| {
-                        entity_mut
-                            .world_scope(|world| world.trigger(KnockbackEvent { entity, data }));
-                    },
-                ));
-            },
-        );
+            let data = KnockbackData::Set(p.delta.to_vec3());
+
+            // this is to make sure the same entity velocity update doesn't get sent
+            // multiple times when in swarms
+            if should_apply_entity_update(
+                &mut commands,
+                &mut world_holder.partial.write(),
+                entity,
+                entity_update_query,
+            ) {
+                commands.trigger(KnockbackEvent { entity, data });
+            }
+        });
     }
 
     pub fn set_entity_link(&mut self, p: &ClientboundSetEntityLink) {
@@ -795,8 +799,8 @@ impl GamePacketHandler<'_> {
 
         as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
             self.ecs,
-            |(mut commands, mut query)| {
-                let (entity_id_index, world_holder) = query.get_mut(self.player).unwrap();
+            |(mut commands, query)| {
+                let (entity_id_index, world_holder) = query.get(self.player).unwrap();
 
                 let Some(entity) = entity_id_index.get_by_minecraft_entity(p.id) else {
                     warn!("Got teleport entity packet for unknown entity id {}", p.id);
@@ -840,8 +844,8 @@ impl GamePacketHandler<'_> {
     pub fn move_entity_pos(&mut self, p: &ClientboundMoveEntityPos) {
         as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
             self.ecs,
-            |(mut commands, mut query)| {
-                let (entity_id_index, world_holder) = query.get_mut(self.player).unwrap();
+            |(mut commands, query)| {
+                let (entity_id_index, world_holder) = query.get(self.player).unwrap();
 
                 debug!("Got move entity pos packet {p:?}");
 
@@ -879,8 +883,8 @@ impl GamePacketHandler<'_> {
     pub fn move_entity_pos_rot(&mut self, p: &ClientboundMoveEntityPosRot) {
         as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
             self.ecs,
-            |(mut commands, mut query)| {
-                let (entity_id_index, world_holder) = query.get_mut(self.player).unwrap();
+            |(mut commands, query)| {
+                let (entity_id_index, world_holder) = query.get(self.player).unwrap();
 
                 debug!("Got move entity pos rot packet {p:?}");
 
@@ -929,8 +933,8 @@ impl GamePacketHandler<'_> {
     pub fn move_entity_rot(&mut self, p: &ClientboundMoveEntityRot) {
         as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
             self.ecs,
-            |(mut commands, mut query)| {
-                let (entity_id_index, world_holder) = query.get_mut(self.player).unwrap();
+            |(mut commands, query)| {
+                let (entity_id_index, world_holder) = query.get(self.player).unwrap();
 
                 let entity = entity_id_index.get_by_minecraft_entity(p.entity_id);
                 if let Some(entity) = entity {
@@ -1527,10 +1531,28 @@ impl GamePacketHandler<'_> {
     }
 
     pub fn entity_position_sync(&mut self, p: &ClientboundEntityPositionSync) {
-        as_system::<(Commands, Query<(&EntityIdIndex, &WorldHolder)>)>(
+        as_system::<(
+            Commands,
+            Query<(
+                &EntityIdIndex,
+                &WorldHolder,
+                Option<&LocalEntity>,
+                &mut Physics,
+                &mut Position,
+                &mut LookDirection,
+            )>,
+            EntityUpdateQuery,
+        )>(
             self.ecs,
-            |(mut commands, mut query)| {
-                let (entity_id_index, world_holder) = query.get_mut(self.player).unwrap();
+            |(mut commands, mut query, entity_update_query)| {
+                let (
+                    entity_id_index,
+                    world_holder,
+                    local_entity,
+                    mut physics,
+                    mut position,
+                    mut look_direction,
+                ) = query.get_mut(self.player).unwrap();
 
                 let Some(entity) = entity_id_index.get_by_minecraft_entity(p.id) else {
                     debug!("Got teleport entity packet for unknown entity id {}", p.id);
@@ -1541,28 +1563,32 @@ impl GamePacketHandler<'_> {
                 let new_on_ground = p.on_ground;
                 let new_look_direction = p.values.look_direction;
 
-                commands.entity(entity).queue(RelativeEntityUpdate::new(
-                    world_holder.partial.clone(),
-                    move |entity_mut| {
-                        let is_local_entity = entity_mut.get::<LocalEntity>().is_some();
-                        let mut physics = entity_mut.get_mut::<Physics>().unwrap();
+                if !should_apply_entity_update(
+                    &mut commands,
+                    &mut world_holder.partial.write(),
+                    entity,
+                    entity_update_query,
+                ) {
+                    return;
+                }
+                let is_local_entity = local_entity.is_some();
 
-                        physics.vec_delta_codec.set_base(new_position);
+                physics.vec_delta_codec.set_base(new_position);
 
-                        if is_local_entity {
-                            debug!("Ignoring entity position sync packet for local player");
-                            return;
-                        }
+                if is_local_entity {
+                    debug!("Ignoring entity position sync packet for local player");
+                    return;
+                }
 
-                        physics.set_on_ground(new_on_ground);
+                physics.set_on_ground(new_on_ground);
 
-                        let mut position = entity_mut.get_mut::<Position>().unwrap();
-                        **position = new_position;
+                if **position != new_position {
+                    **position = new_position;
+                }
 
-                        let mut look_direction = entity_mut.get_mut::<LookDirection>().unwrap();
-                        *look_direction = new_look_direction;
-                    },
-                ));
+                if *look_direction != new_look_direction {
+                    *look_direction = new_look_direction;
+                }
             },
         );
     }
diff --git a/azalea-client/src/plugins/packet/mod.rs b/azalea-client/src/plugins/packet/mod.rs
index 63a94ee0..4309850e 100644
--- a/azalea-client/src/plugins/packet/mod.rs
+++ b/azalea-client/src/plugins/packet/mod.rs
@@ -11,6 +11,7 @@ use crate::chat::ChatReceivedEvent;
 pub mod config;
 pub mod game;
 pub mod login;
+pub mod relative_updates;
 
 pub struct PacketPlugin;
 
@@ -30,23 +31,27 @@ pub fn death_event_on_0_health(
 
 impl Plugin for PacketPlugin {
     fn build(&self, app: &mut App) {
-        app.add_observer(game::handle_outgoing_packets_observer)
-            .add_observer(config::handle_outgoing_packets_observer)
-            .add_observer(login::handle_outgoing_packets_observer)
-            .add_systems(Update, death_event_on_0_health)
-            .add_message::<game::ReceiveGamePacketEvent>()
-            .add_message::<config::ReceiveConfigPacketEvent>()
-            .add_message::<login::ReceiveLoginPacketEvent>()
-            //
-            .add_message::<game::AddPlayerEvent>()
-            .add_message::<game::RemovePlayerEvent>()
-            .add_message::<game::UpdatePlayerEvent>()
-            .add_message::<ChatReceivedEvent>()
-            .add_message::<game::DeathEvent>()
-            .add_message::<game::KeepAliveEvent>()
-            .add_message::<game::ResourcePackEvent>()
-            .add_message::<game::WorldLoadedEvent>()
-            .add_message::<login::ReceiveCustomQueryEvent>();
+        app.add_systems(
+            Update,
+            relative_updates::debug_detect_updates_received_on_local_entities,
+        )
+        .add_observer(game::handle_outgoing_packets_observer)
+        .add_observer(config::handle_outgoing_packets_observer)
+        .add_observer(login::handle_outgoing_packets_observer)
+        .add_systems(Update, death_event_on_0_health)
+        .add_message::<game::ReceiveGamePacketEvent>()
+        .add_message::<config::ReceiveConfigPacketEvent>()
+        .add_message::<login::ReceiveLoginPacketEvent>()
+        //
+        .add_message::<game::AddPlayerEvent>()
+        .add_message::<game::RemovePlayerEvent>()
+        .add_message::<game::UpdatePlayerEvent>()
+        .add_message::<ChatReceivedEvent>()
+        .add_message::<game::DeathEvent>()
+        .add_message::<game::KeepAliveEvent>()
+        .add_message::<game::ResourcePackEvent>()
+        .add_message::<game::WorldLoadedEvent>()
+        .add_message::<login::ReceiveCustomQueryEvent>();
     }
 }
 
@@ -70,12 +75,21 @@ macro_rules! __declare_packet_handlers {
 
 pub(crate) use __declare_packet_handlers as declare_packet_handlers;
 
+#[derive(Resource)]
+struct CachedSystemState<T: SystemParam + 'static>(SystemState<T>);
+
 pub(crate) fn as_system<T>(ecs: &mut World, f: impl FnOnce(T::Item<'_, '_>))
 where
     T: SystemParam + 'static,
 {
-    let mut system_state = SystemState::<T>::new(ecs);
+    // creating a new SystemState is expensive, so we save them as a Resource in the
+    // ecs
+    let mut system_state = match ecs.remove_resource::<CachedSystemState<T>>() {
+        Some(s) => s.0,
+        None => SystemState::<T>::new(ecs),
+    };
     let values = system_state.get_mut(ecs);
     f(values);
     system_state.apply(ecs);
+    ecs.insert_resource(CachedSystemState(system_state));
 }
diff --git a/azalea-client/src/plugins/packet/relative_updates.rs b/azalea-client/src/plugins/packet/relative_updates.rs
new file mode 100644
index 00000000..2f7112b8
--- /dev/null
+++ b/azalea-client/src/plugins/packet/relative_updates.rs
@@ -0,0 +1,163 @@
+// How entity updates are processed (to avoid issues with shared worlds)
+// - each bot contains a map of { entity id: updates received }
+// - the shared world also contains a canonical "true" updates received for each
+//   entity
+// - when a client loads an entity, its "updates received" is set to the same as
+//   the global "updates received"
+// - when the shared world sees an entity for the first time, the "updates
+//   received" is set to 1.
+// - clients can force the shared "updates received" to 0 to make it so certain
+//   entities (i.e. other bots in our swarm) don't get confused and updated by
+//   other bots
+// - when a client gets an update to an entity, we check if our "updates
+//   received" is the same as the shared world's "updates received": if it is,
+//   then process the update and increment the client's and shared world's
+//   "updates received" if not, then we simply increment our local "updates
+//   received" and do nothing else
+
+use std::sync::Arc;
+
+use azalea_core::entity_id::MinecraftEntityId;
+use azalea_entity::LocalEntity;
+use azalea_world::PartialWorld;
+use bevy_ecs::prelude::*;
+use derive_more::{Deref, DerefMut};
+use parking_lot::RwLock;
+use tracing::warn;
+
+use crate::packet::as_system;
+
+/// An [`EntityCommand`] that applies a "relative update" to an entity, which
+/// means this update won't be run multiple times by different clients in the
+/// same world.
+///
+/// This is used to avoid a bug where when there's multiple clients in the same
+/// world and an entity sends a relative move packet to all clients, its
+/// position gets desynced since the relative move is applied multiple times.
+///
+/// Don't use this unless you actually got an entity update packet that all
+/// other clients within render distance will get too. You usually don't need
+/// this when the change isn't relative either.
+pub struct RelativeEntityUpdate {
+    pub partial_world: Arc<RwLock<PartialWorld>>,
+    // a function that takes the entity and updates it
+    pub update: Box<dyn FnOnce(&mut EntityWorldMut) + Send + Sync>,
+}
+impl RelativeEntityUpdate {
+    pub fn new(
+        partial_world: Arc<RwLock<PartialWorld>>,
+        update: impl FnOnce(&mut EntityWorldMut) + Send + Sync + 'static,
+    ) -> Self {
+        Self {
+            partial_world,
+            update: Box::new(update),
+        }
+    }
+}
+
+/// A component that counts the number of times this entity has been modified.
+///
+/// This is used for making sure two clients don't do the same relative update
+/// on an entity.
+///
+/// If an entity is local (i.e. it's a client/LocalEntity), this component
+/// should NOT be present in the entity.
+#[derive(Component, Debug, Deref, DerefMut)]
+pub struct UpdatesReceived(u32);
+
+pub type EntityUpdateQuery<'world, 'state, 'a> = Query<
+    'world,
+    'state,
+    (
+        &'a MinecraftEntityId,
+        Option<&'a UpdatesReceived>,
+        Option<&'a LocalEntity>,
+    ),
+>;
+
+/// See [`RelativeEntityUpdate`] for details.
+///
+/// Calling this function will have the same effect as using the Command, but
+/// it's more performant than the Command.
+pub fn should_apply_entity_update(
+    commands: &mut Commands,
+    partial_world: &mut PartialWorld,
+    entity: Entity,
+    entity_update_query: EntityUpdateQuery,
+) -> bool {
+    let partial_entity_infos = &mut partial_world.entity_infos;
+
+    if Some(entity) == partial_entity_infos.owner_entity {
+        // if the entity owns this partial world, it's always allowed to update itself
+        return true;
+    };
+
+    let Ok((minecraft_entity_id, updates_received, local_entity)) = entity_update_query.get(entity)
+    else {
+        warn!("called should_apply_entity_update on an entity with missing components");
+        return false;
+    };
+
+    if local_entity.is_some() {
+        // a client tried to update another client, which isn't allowed
+        return false;
+    }
+
+    let this_client_updates_received = partial_entity_infos
+        .updates_received
+        .get(&minecraft_entity_id)
+        .copied();
+
+    let can_update = if let Some(updates_received) = updates_received {
+        this_client_updates_received.unwrap_or(1) == **updates_received
+    } else {
+        // no UpdatesReceived means the entity was just spawned
+        true
+    };
+    if can_update {
+        let new_updates_received = this_client_updates_received.unwrap_or(0) + 1;
+        partial_entity_infos
+            .updates_received
+            .insert(*minecraft_entity_id, new_updates_received);
+
+        commands
+            .entity(entity)
+            .insert(UpdatesReceived(new_updates_received));
+
+        return true;
+    }
+    false
+}
+
+impl EntityCommand for RelativeEntityUpdate {
+    fn apply(self, mut entity_mut: EntityWorldMut) {
+        let partial_world = self.partial_world.clone();
+        let mut should_update = false;
+        let entity = entity_mut.id();
+
+        entity_mut.world_scope(|ecs| {
+            as_system::<(Commands, EntityUpdateQuery)>(ecs, |(mut commands, query)| {
+                should_update = should_apply_entity_update(
+                    &mut commands,
+                    &mut partial_world.write(),
+                    entity,
+                    query,
+                );
+            });
+        });
+
+        if should_update {
+            (self.update)(&mut entity_mut);
+        }
+    }
+}
+
+/// A system that logs a warning if an entity has both [`UpdatesReceived`]
+/// and [`LocalEntity`].
+pub fn debug_detect_updates_received_on_local_entities(
+    query: Query<Entity, (With<LocalEntity>, With<UpdatesReceived>)>,
+) {
+    for entity in &query {
+        warn!("Entity {entity:?} has both LocalEntity and UpdatesReceived");
+    }
+}