properly disconnect on ungraceful disconnect

1 files changed, 17 insertions, 4 deletions

diff --git a/azalea-client/src/disconnect.rs b/azalea-client/src/disconnect.rs
index 100618c3..e01ffe1b 100644
--- a/azalea-client/src/disconnect.rs
+++ b/azalea-client/src/disconnect.rs
@@ -3,17 +3,19 @@
 use azalea_ecs::{
     app::{App, CoreStage, Plugin},
     entity::Entity,
-    event::EventReader,
-    system::Commands,
+    event::{EventReader, EventWriter},
+    system::{Commands, Query},
+    AppTickExt,
 };
 
-use crate::client::JoinedClientBundle;
+use crate::{client::JoinedClientBundle, LocalPlayer};
 
 pub struct DisconnectPlugin;
 impl Plugin for DisconnectPlugin {
     fn build(&self, app: &mut App) {
         app.add_event::<DisconnectEvent>()
-            .add_system_to_stage(CoreStage::PostUpdate, handle_disconnect);
+            .add_system_to_stage(CoreStage::PostUpdate, handle_disconnect)
+            .add_tick_system(disconnect_on_read_packets_ended);
     }
 }
 
@@ -29,3 +31,14 @@ pub fn handle_disconnect(mut commands: Commands, mut events: EventReader<Disconn
         commands.entity(*entity).remove::<JoinedClientBundle>();
     }
 }
+
+fn disconnect_on_read_packets_ended(
+    local_player: Query<(Entity, &LocalPlayer)>,
+    mut disconnect_events: EventWriter<DisconnectEvent>,
+) {
+    for (entity, local_player) in &local_player {
+        if local_player.read_packets_task.is_finished() {
+            disconnect_events.send(DisconnectEvent { entity });
+        }
+    }
+}