From f5f15362f2cb48088eb8ccf80988f994fecd81c9 Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Fri, 21 Feb 2025 21:51:26 +0000
Subject: [PATCH] fix some components not being removed from clients and add
 debugecsleak testbot command

---
 azalea-client/src/disconnect.rs           | 13 +++++--
 azalea-protocol/src/read.rs               |  2 +-
 azalea/Cargo.toml                         |  7 +++-
 azalea/examples/testbot/commands/debug.rs | 46 +++++++++++++++++++++++
 4 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
index 06648691..a92423d7 100644
--- a/azalea-client/src/disconnect.rs
+++ b/azalea-client/src/disconnect.rs
@@ -1,7 +1,7 @@
 //! Disconnect a client from the server.
 
 use azalea_chat::FormattedText;
-use azalea_entity::{EntityBundle, LocalEntity};
+use azalea_entity::{metadata::PlayerMetadataBundle, EntityBundle, LocalEntity};
 use bevy_app::{App, Plugin, PostUpdate};
 use bevy_ecs::{
     component::Component,
@@ -15,7 +15,10 @@ use bevy_ecs::{
 use derive_more::Deref;
 use tracing::trace;
 
-use crate::{client::JoinedClientBundle, events::LocalPlayerEvents, raw_connection::RawConnection};
+use crate::{
+    client::JoinedClientBundle, events::LocalPlayerEvents, raw_connection::RawConnection,
+    InstanceHolder,
+};
 
 pub struct DisconnectPlugin;
 impl Plugin for DisconnectPlugin {
@@ -39,8 +42,8 @@ pub struct DisconnectEvent {
     pub reason: Option<FormattedText>,
 }
 
-/// System that removes the [`JoinedClientBundle`] from the entity when it
-/// receives a [`DisconnectEvent`].
+/// A system that removes the several components from our clients when they get
+/// a [`DisconnectEvent`].
 pub fn remove_components_from_disconnected_players(
     mut commands: Commands,
     mut events: EventReader<DisconnectEvent>,
@@ -51,6 +54,8 @@ pub fn remove_components_from_disconnected_players(
             .entity(*entity)
             .remove::<JoinedClientBundle>()
             .remove::<EntityBundle>()
+            .remove::<InstanceHolder>()
+            .remove::<PlayerMetadataBundle>()
             // this makes it close the tcp connection
             .remove::<RawConnection>()
             // swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect
diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs
index 50764c88..01744169 100755
--- a/azalea-protocol/src/read.rs
+++ b/azalea-protocol/src/read.rs
@@ -116,7 +116,7 @@ fn parse_frame(buffer: &mut Cursor<Vec<u8>>) -> Result<Box<[u8]>, FrameSplitterE
     if buffer.position() == buffer.get_ref().len() as u64 {
         // reset the inner vec once we've reached the end of the buffer so we don't keep
         // leaking memory
-        *buffer.get_mut() = Vec::new();
+        buffer.get_mut().clear();
         buffer.set_position(0);
     }
 
diff --git a/azalea/Cargo.toml b/azalea/Cargo.toml
index bd3e54b9..33730361 100644
--- a/azalea/Cargo.toml
+++ b/azalea/Cargo.toml
@@ -56,7 +56,12 @@ anyhow.workspace = true
 default = ["log", "serde"]
 # enables bevy_log::LogPlugin by default
 log = ["azalea-client/log"]
-serde = ["dep:serde", "azalea-core/serde", "azalea-registry/serde", "azalea-world/serde"]
+serde = [
+    "dep:serde",
+    "azalea-core/serde",
+    "azalea-registry/serde",
+    "azalea-world/serde",
+]
 
 [[bench]]
 name = "pathfinder"
diff --git a/azalea/examples/testbot/commands/debug.rs b/azalea/examples/testbot/commands/debug.rs
index 1e0846f4..a7f15d2b 100644
--- a/azalea/examples/testbot/commands/debug.rs
+++ b/azalea/examples/testbot/commands/debug.rs
@@ -1,5 +1,7 @@
 //! Commands for debugging and getting the current state of the bot.
 
+use std::{env, fs::File, io::Write, thread, time::Duration};
+
 use azalea::{
     brigadier::prelude::*,
     entity::{LookDirection, Position},
@@ -161,4 +163,48 @@ pub fn register(commands: &mut CommandDispatcher<Mutex<CommandSource>>) {
         ));
         1
     }));
+
+    commands.register(literal("debugecsleak").executes(|ctx: &Ctx| {
+        let source = ctx.source.lock();
+
+        source.reply("Ok!");
+        source.bot.disconnect();
+
+        let ecs = source.bot.ecs.clone();
+        thread::spawn(move || {
+            thread::sleep(Duration::from_secs(1));
+            // dump the ecs
+
+            let ecs = ecs.lock();
+
+            let report_path = env::temp_dir().join("azalea-ecs-leak-report.txt");
+            let mut report = File::create(&report_path).unwrap();
+
+            for entity in ecs.iter_entities() {
+                writeln!(report, "Entity: {}", entity.id()).unwrap();
+                let archetype = entity.archetype();
+                let component_count = archetype.component_count();
+
+                let component_names = archetype
+                    .components()
+                    .map(|c| ecs.components().get_info(c).unwrap().name())
+                    .collect::<Vec<_>>();
+                writeln!(
+                    report,
+                    "- {component_count} components: {}",
+                    component_names.join(", ")
+                )
+                .unwrap();
+            }
+
+            for (info, _) in ecs.iter_resources() {
+                writeln!(report, "Resource: {}", info.name()).unwrap();
+                writeln!(report, "- Size: {} bytes", info.layout().size()).unwrap();
+            }
+
+            println!("\x1b[1mWrote report to {}\x1b[m", report_path.display());
+        });
+
+        1
+    }));
 }