From 48640ca303784a3761dabc7436cd96ac93db279b Mon Sep 17 00:00:00 2001
From: mat <github@matdoes.dev>
Date: Tue, 21 Feb 2023 21:53:22 -0600
Subject: [PATCH] properly disconnect on ungraceful disconnect

---
 azalea-client/src/client.rs          | 20 ++++++++++----------
 azalea-client/src/disconnect.rs      | 21 +++++++++++++++++----
 azalea-client/src/local_player.rs    | 18 +++++++++++-------
 azalea-client/src/packet_handling.rs |  4 ----
 4 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index d9214329..69a50c33 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -207,14 +207,6 @@ impl Client {
 
         let (packet_writer_sender, packet_writer_receiver) = mpsc::unbounded_channel();
 
-        let mut local_player = crate::local_player::LocalPlayer::new(
-            entity,
-            packet_writer_sender,
-            // default to an empty world, it'll be set correctly later when we
-            // get the login packet
-            Arc::new(RwLock::new(World::default())),
-        );
-
         // start receiving packets
         let packet_receiver = packet_handling::PacketReceiver {
             packets: Arc::new(Mutex::new(Vec::new())),
@@ -227,8 +219,16 @@ impl Client {
                 .clone()
                 .write_task(write_conn, packet_writer_receiver),
         );
-        local_player.tasks.push(read_packets_task);
-        local_player.tasks.push(write_packets_task);
+
+        let local_player = crate::local_player::LocalPlayer::new(
+            entity,
+            packet_writer_sender,
+            // default to an empty world, it'll be set correctly later when we
+            // get the login packet
+            Arc::new(RwLock::new(World::default())),
+            read_packets_task,
+            write_packets_task,
+        );
 
         ecs.entity_mut(entity).insert(JoinedClientBundle {
             local_player,
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 });
+        }
+    }
+}
diff --git a/azalea-client/src/local_player.rs b/azalea-client/src/local_player.rs
index 8298cf0a..d31f840f 100644
--- a/azalea-client/src/local_player.rs
+++ b/azalea-client/src/local_player.rs
@@ -46,9 +46,11 @@ pub struct LocalPlayer {
     /// world. (Only relevant if you're using a shared world, i.e. a swarm)
     pub world: Arc<RwLock<World>>,
 
-    /// A list of async tasks that are running and will stop running when this
-    /// LocalPlayer is dropped or disconnected with [`Self::disconnect`]
-    pub(crate) tasks: Vec<JoinHandle<()>>,
+    /// A task that reads packets from the server. The client is disconnected
+    /// when this task ends.
+    pub(crate) read_packets_task: JoinHandle<()>,
+    /// A task that writes packets from the server.
+    pub(crate) write_packets_task: JoinHandle<()>,
 }
 
 /// Component for entities that can move and sprint. Usually only in
@@ -87,6 +89,8 @@ impl LocalPlayer {
         entity: Entity,
         packet_writer: mpsc::UnboundedSender<ServerboundGamePacket>,
         world: Arc<RwLock<World>>,
+        read_packets_task: JoinHandle<()>,
+        write_packets_task: JoinHandle<()>,
     ) -> Self {
         let client_information = ClientInformation::default();
 
@@ -102,7 +106,8 @@ impl LocalPlayer {
                 Some(entity),
             ))),
 
-            tasks: Vec::new(),
+            read_packets_task,
+            write_packets_task,
         }
     }
 
@@ -117,9 +122,8 @@ impl LocalPlayer {
 impl Drop for LocalPlayer {
     /// Stop every active task when the `LocalPlayer` is dropped.
     fn drop(&mut self) {
-        for task in &self.tasks {
-            task.abort();
-        }
+        self.read_packets_task.abort();
+        self.write_packets_task.abort();
     }
 }
 
diff --git a/azalea-client/src/packet_handling.rs b/azalea-client/src/packet_handling.rs
index 873ac3f5..d5cce3a8 100644
--- a/azalea-client/src/packet_handling.rs
+++ b/azalea-client/src/packet_handling.rs
@@ -961,10 +961,6 @@ impl PacketReceiver {
                 }
             }
         }
-        // TODO: it should send a DisconnectEvent here somehow
-        // maybe use a tokio::sync::oneshot that tells it to close and have the
-        // receiver in localplayer and have a system that watches that or
-        // something?
     }
 
     /// Consume the [`ServerboundGamePacket`] queue and actually write the