From 43d7c428e317f503dda899cade9e3ec8024de4de Mon Sep 17 00:00:00 2001
From: mat <git@matdoes.dev>
Date: Fri, 18 Apr 2025 00:32:31 +0200
Subject: [PATCH] fix another panic on disconnect and slightly optimize client
 events loop

---
 azalea-client/src/plugins/connection.rs | 13 ++++++++++---
 azalea/src/swarm/mod.rs                 |  7 +++++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/azalea-client/src/plugins/connection.rs b/azalea-client/src/plugins/connection.rs
index b462535e..30ac5d20 100644
--- a/azalea-client/src/plugins/connection.rs
+++ b/azalea-client/src/plugins/connection.rs
@@ -126,7 +126,11 @@ fn poll_all_writer_tasks(mut conn_query: Query<&mut RawConnection>) {
             // this needs to be done at some point every update to make sure packets are
             // actually sent to the network
 
-            net_conn.poll_writer();
+            if net_conn.poll_writer().is_some() {
+                // means the writer task ended
+                conn.network = None;
+                conn.is_alive = false;
+            }
         }
     }
 }
@@ -319,9 +323,11 @@ impl NetworkConnection {
         Ok(())
     }
 
-    pub fn poll_writer(&mut self) {
+    /// Makes sure packets get sent and returns Some(()) if the connection has
+    /// closed.
+    pub fn poll_writer(&mut self) -> Option<()> {
         let poll_once_res = future::poll_once(&mut self.writer_task);
-        future::block_on(poll_once_res);
+        future::block_on(poll_once_res)
     }
 
     pub fn set_compression_threshold(&mut self, threshold: Option<u32>) {
@@ -348,6 +354,7 @@ async fn write_task(
             break;
         };
     }
+
     trace!("write task is done");
 }
 
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index ba8fd7bb..ad5ac815 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -485,7 +485,8 @@ where
             if let Some(handler) = &self.handler {
                 let ecs_mutex = first_bot.ecs.clone();
                 let mut ecs = ecs_mutex.lock();
-                let Some(first_bot_state) = first_bot.query::<Option<&S>>(&mut ecs).cloned() else {
+                let mut query = ecs.query::<Option<&S>>();
+                let Ok(Some(first_bot_state)) = query.get(&mut ecs, first_bot.entity) else {
                     error!(
                         "the first bot ({} / {}) is missing the required state component! none of the client handler functions will be called.",
                         first_bot.username(),
@@ -494,6 +495,7 @@ where
                     continue;
                 };
                 let first_bot_entity = first_bot.entity;
+                let first_bot_state = first_bot_state.clone();
 
                 tokio::spawn((handler)(first_bot, first_event, first_bot_state.clone()));
 
@@ -504,7 +506,7 @@ where
                     let state = match states.entry(bot.entity) {
                         hash_map::Entry::Occupied(e) => e.into_mut(),
                         hash_map::Entry::Vacant(e) => {
-                            let Some(state) = bot.query::<Option<&S>>(&mut ecs).cloned() else {
+                            let Ok(Some(state)) = query.get(&mut ecs, bot.entity) else {
                                 error!(
                                     "one of our bots ({} / {}) is missing the required state component! its client handler function will not be called.",
                                     bot.username(),
@@ -512,6 +514,7 @@ where
                                 );
                                 continue;
                             };
+                            let state = state.clone();
                             e.insert(state)
                         }
                     };