make loading chunks its own bevy system

azalea-client/src/chunk_batching.rs → azalea-client/src/chunks.rs

@@ -2,26 +2,37 @@
 //! for making the server spread out how often it sends us chunk packets
 //! depending on our receiving speed.
 
-use std::time::{Duration, Instant};
+use std::{
+    io::Cursor,
+    time::{Duration, Instant},
+};
 
-use azalea_protocol::packets::game::serverbound_chunk_batch_received_packet::ServerboundChunkBatchReceivedPacket;
+use azalea_core::position::ChunkPos;
+use azalea_nbt::NbtCompound;
+use azalea_protocol::packets::game::{
+    clientbound_level_chunk_with_light_packet::ClientboundLevelChunkWithLightPacket,
+    serverbound_chunk_batch_received_packet::ServerboundChunkBatchReceivedPacket,
+};
 use bevy_app::{App, Plugin, Update};
 use bevy_ecs::prelude::*;
+use tracing::{error, trace};
 
 use crate::{
     interact::handle_block_interact_event,
     inventory::InventorySet,
     local_player::{handle_send_packet_event, SendPacketEvent},
     respawn::perform_respawn,
+    InstanceHolder,
 };
 
-pub struct ChunkBatchingPlugin;
-impl Plugin for ChunkBatchingPlugin {
+pub struct ChunkPlugin;
+impl Plugin for ChunkPlugin {
     fn build(&self, app: &mut App) {
         app.add_systems(
             Update,
             (
                 handle_chunk_batch_start_event,
+                handle_receive_chunk_events,
                 handle_chunk_batch_finished_event,
             )
                 .chain()
@@ -30,18 +41,88 @@ impl Plugin for ChunkBatchingPlugin {
                 .before(handle_block_interact_event)
                 .before(perform_respawn),
         )
+        .add_event::<ReceiveChunkEvent>()
         .add_event::<ChunkBatchStartEvent>()
         .add_event::<ChunkBatchFinishedEvent>();
     }
 }
 
+#[derive(Event)]
+pub struct ReceiveChunkEvent {
+    pub entity: Entity,
+    pub packet: ClientboundLevelChunkWithLightPacket,
+}
+
 #[derive(Component, Clone, Debug)]
 pub struct ChunkBatchInfo {
     pub start_time: Instant,
     pub aggregated_duration_per_chunk: Duration,
     pub old_samples_weight: u32,
 }
 
+#[derive(Event)]
+pub struct ChunkBatchStartEvent {
+    pub entity: Entity,
+}
+#[derive(Event)]
+pub struct ChunkBatchFinishedEvent {
+    pub entity: Entity,
+    pub batch_size: u32,
+}
+
+fn handle_receive_chunk_events(
+    mut events: EventReader<ReceiveChunkEvent>,
+    mut query: Query<&mut InstanceHolder>,
+) {
+    for event in events.read() {
+        let pos = ChunkPos::new(event.packet.x, event.packet.z);
+
+        let local_player = query.get_mut(event.entity).unwrap();
+
+        // OPTIMIZATION: if we already know about the chunk from the
+        // shared world (and not ourselves), then we don't need to
+        // parse it again. This is only used when we have a shared
+        // world, since we check that the chunk isn't currently owned
+        // by this client.
+        let shared_chunk = local_player.instance.read().chunks.get(&pos);
+        let this_client_has_chunk = local_player
+            .partial_instance
+            .read()
+            .chunks
+            .limited_get(&pos)
+            .is_some();
+
+        let mut world = local_player.instance.write();
+        let mut partial_world = local_player.partial_instance.write();
+
+        if !this_client_has_chunk {
+            if let Some(shared_chunk) = shared_chunk {
+                trace!("Skipping parsing chunk {pos:?} because we already know about it");
+                partial_world.chunks.set_with_shared_reference(
+                    &pos,
+                    Some(shared_chunk.clone()),
+                    &mut world.chunks,
+                );
+                continue;
+            }
+        }
+
+        let heightmaps = event.packet.chunk_data.heightmaps.as_compound();
+        // necessary to make the unwrap_or work
+        let empty_nbt_compound = NbtCompound::default();
+        let heightmaps = heightmaps.unwrap_or(&empty_nbt_compound);
+
+        if let Err(e) = partial_world.chunks.replace_with_packet_data(
+            &pos,
+            &mut Cursor::new(&event.packet.chunk_data.data),
+            heightmaps,
+            &mut world.chunks,
+        ) {
+            error!("Couldn't set chunk data: {e}");
+        }
+    }
+}
+
 impl ChunkBatchInfo {
     pub fn batch_finished(&mut self, batch_size: u32) {
         if batch_size == 0 {
@@ -65,16 +146,6 @@ impl ChunkBatchInfo {
     }
 }
 
-#[derive(Event)]
-pub struct ChunkBatchStartEvent {
-    pub entity: Entity,
-}
-#[derive(Event)]
-pub struct ChunkBatchFinishedEvent {
-    pub entity: Entity,
-    pub batch_size: u32,
-}
-
 pub fn handle_chunk_batch_start_event(
     mut query: Query<&mut ChunkBatchInfo>,
     mut events: EventReader<ChunkBatchStartEvent>,

azalea-client/src/client.rs

@@ -1,7 +1,7 @@
 use crate::{
     attack::{self, AttackPlugin},
     chat::ChatPlugin,
-    chunk_batching::{ChunkBatchInfo, ChunkBatchingPlugin},
+    chunks::{ChunkBatchInfo, ChunkPlugin},
     disconnect::{DisconnectEvent, DisconnectPlugin},
     events::{Event, EventPlugin, LocalPlayerEvents},
     interact::{CurrentSequenceNumber, InteractPlugin},
@@ -782,7 +782,7 @@ impl PluginGroup for DefaultPlugins {
             .add(RespawnPlugin)
             .add(MinePlugin)
             .add(AttackPlugin)
-            .add(ChunkBatchingPlugin)
+            .add(ChunkPlugin)
             .add(TickBroadcastPlugin);
         #[cfg(feature = "log")]
         {

azalea-client/src/lib.rs

@@ -12,7 +12,7 @@
 mod account;
 pub mod attack;
 pub mod chat;
-pub mod chunk_batching;
+pub mod chunks;
 mod client;
 pub mod disconnect;
 mod entity_query;

azalea-client/src/packet_handling/configuration.rs

@@ -140,7 +140,7 @@ pub fn process_packet_events(ecs: &mut World) {
                         abilities: crate::local_player::PlayerAbilities::default(),
                         permission_level: crate::local_player::PermissionLevel::default(),
                         hunger: Hunger::default(),
-                        chunk_batch_info: crate::chunk_batching::ChunkBatchInfo::default(),
+                        chunk_batch_info: crate::chunks::ChunkBatchInfo::default(),
 
                         entity_id_index: EntityIdIndex::default(),
 

azalea-client/src/packet_handling/game.rs

@@ -16,7 +16,6 @@ use azalea_entity::{
     Dead, EntityBundle, EntityKind, LastSentPosition, LoadedBy, LocalEntity, LookDirection,
     Physics, PlayerBundle, Position, RelativeEntityUpdate,
 };
-use azalea_nbt::NbtCompound;
 use azalea_protocol::{
     packets::game::{
         clientbound_player_combat_kill_packet::ClientboundPlayerCombatKillPacket,
@@ -34,7 +33,7 @@ use tracing::{debug, error, trace, warn};
 
 use crate::{
     chat::{ChatPacket, ChatReceivedEvent},
-    chunk_batching,
+    chunks,
     disconnect::DisconnectEvent,
     inventory::{
         ClientSideCloseContainerEvent, InventoryComponent, MenuOpenedEvent,
@@ -339,24 +338,22 @@ pub fn process_packet_events(ecs: &mut World) {
             ClientboundGamePacket::ChunkBatchStart(_p) => {
                 // the packet is empty, just a marker to tell us when the batch starts and ends
                 debug!("Got chunk batch start");
-                let mut system_state: SystemState<
-                    EventWriter<chunk_batching::ChunkBatchStartEvent>,
-                > = SystemState::new(ecs);
+                let mut system_state: SystemState<EventWriter<chunks::ChunkBatchStartEvent>> =
+                    SystemState::new(ecs);
                 let mut chunk_batch_start_events = system_state.get_mut(ecs);
 
-                chunk_batch_start_events.send(chunk_batching::ChunkBatchStartEvent {
+                chunk_batch_start_events.send(chunks::ChunkBatchStartEvent {
                     entity: player_entity,
                 });
             }
             ClientboundGamePacket::ChunkBatchFinished(p) => {
                 debug!("Got chunk batch finished {p:?}");
 
-                let mut system_state: SystemState<
-                    EventWriter<chunk_batching::ChunkBatchFinishedEvent>,
-                > = SystemState::new(ecs);
+                let mut system_state: SystemState<EventWriter<chunks::ChunkBatchFinishedEvent>> =
+                    SystemState::new(ecs);
                 let mut chunk_batch_start_events = system_state.get_mut(ecs);
 
-                chunk_batch_start_events.send(chunk_batching::ChunkBatchFinishedEvent {
+                chunk_batch_start_events.send(chunks::ChunkBatchFinishedEvent {
                     entity: player_entity,
                     batch_size: p.batch_size,
                 });
@@ -597,54 +594,14 @@ pub fn process_packet_events(ecs: &mut World) {
             }
             ClientboundGamePacket::LevelChunkWithLight(p) => {
                 debug!("Got chunk with light packet {} {}", p.x, p.z);
-                let pos = ChunkPos::new(p.x, p.z);
 
-                let mut system_state: SystemState<Query<&mut InstanceHolder>> =
+                let mut system_state: SystemState<EventWriter<chunks::ReceiveChunkEvent>> =
                     SystemState::new(ecs);
-                let mut query = system_state.get_mut(ecs);
-                let local_player = query.get_mut(player_entity).unwrap();
-
-                // OPTIMIZATION: if we already know about the chunk from the
-                // shared world (and not ourselves), then we don't need to
-                // parse it again. This is only used when we have a shared
-                // world, since we check that the chunk isn't currently owned
-                // by this client.
-                let shared_chunk = local_player.instance.read().chunks.get(&pos);
-                let this_client_has_chunk = local_player
-                    .partial_instance
-                    .read()
-                    .chunks
-                    .limited_get(&pos)
-                    .is_some();
-
-                let mut world = local_player.instance.write();
-                let mut partial_world = local_player.partial_instance.write();
-
-                if !this_client_has_chunk {
-                    if let Some(shared_chunk) = shared_chunk {
-                        trace!("Skipping parsing chunk {pos:?} because we already know about it");
-                        partial_world.chunks.set_with_shared_reference(
-                            &pos,
-                            Some(shared_chunk.clone()),
-                            &mut world.chunks,
-                        );
-                        continue;
-                    }
-                }
-
-                let heightmaps = p.chunk_data.heightmaps.as_compound();
-                // necessary to make the unwrap_or work
-                let empty_nbt_compound = NbtCompound::default();
-                let heightmaps = heightmaps.unwrap_or(&empty_nbt_compound);
-
-                if let Err(e) = partial_world.chunks.replace_with_packet_data(
-                    &pos,
-                    &mut Cursor::new(&p.chunk_data.data),
-                    heightmaps,
-                    &mut world.chunks,
-                ) {
-                    error!("Couldn't set chunk data: {e}");
-                }
+                let mut receive_chunk_events = system_state.get_mut(ecs);
+                receive_chunk_events.send(chunks::ReceiveChunkEvent {
+                    entity: player_entity,
+                    packet: p.clone(),
+                });
             }
             ClientboundGamePacket::AddEntity(p) => {
                 debug!("Got add entity packet {p:?}");

azalea/src/accept_resource_packs.rs

@@ -1,5 +1,5 @@
 use crate::app::{App, Plugin};
-use azalea_client::chunk_batching::handle_chunk_batch_finished_event;
+use azalea_client::chunks::handle_chunk_batch_finished_event;
 use azalea_client::inventory::InventorySet;
 use azalea_client::packet_handling::{death_event_on_0_health, game::ResourcePackEvent};
 use azalea_client::respawn::perform_respawn;