handle AppExit event

5 files changed, 127 insertions, 63 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 712fbe2b..909c39fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -31,6 +31,7 @@ is breaking anyways, semantic versioning is not followed.
 - Replaced `start_goto_without_mining` with `start_goto_with_opts`.
 - Rename `send_chat_packet` / `send_command_packet` to `write_chat_packet` / `write_command_packet` (for consistency with `write_packet`).
 - Split `ClientInformation` handling out of `BrandPlugin` to `ClientInformationPlugin`.
+- `ClientBuilder::start` and `SwarmBuilder::start` now return a `Result<AppExit>` instead of `Result<!>`.
 
 ### Fixed
 
@@ -46,6 +47,7 @@ is breaking anyways, semantic versioning is not followed.
 - Fix parsing some metadata fields of Display entities.
 - Mining blocks in creative mode now works. (@qwqawawow)
 - Improved matchers on the `ChatPacket` functions to work on more servers. (@ShayBox)
+- Bevy's `AppExit` Event is now handled by Azalea's ECS runner.
 
 ## [0.13.0+mc1.21.5] - 2025-06-15
 
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 884d8c20..98d8a3d3 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -27,8 +27,9 @@ use azalea_protocol::{
     resolver,
 };
 use azalea_world::{Instance, InstanceContainer, InstanceName, MinecraftEntityId, PartialInstance};
-use bevy_app::{App, Plugin, PluginsState, SubApp, Update};
+use bevy_app::{App, AppExit, Plugin, PluginsState, SubApp, Update};
 use bevy_ecs::{
+    event::EventCursor,
     prelude::*,
     schedule::{InternedScheduleLabel, LogLevel, ScheduleBuildSettings},
 };
@@ -36,7 +37,10 @@ use parking_lot::{Mutex, RwLock};
 use simdnbt::owned::NbtCompound;
 use thiserror::Error;
 use tokio::{
-    sync::mpsc::{self},
+    sync::{
+        mpsc::{self},
+        oneshot,
+    },
     time,
 };
 use tracing::{error, info, warn};
@@ -106,7 +110,9 @@ impl StartClientOpts {
         let mut app = App::new();
         app.add_plugins(DefaultPlugins);
 
-        let (ecs_lock, start_running_systems) = start_ecs_runner(app.main_mut());
+        // appexit_rx is unused here since the user should be able to handle it
+        // themselves if they're using StartClientOpts::new
+        let (ecs_lock, start_running_systems, _appexit_rx) = start_ecs_runner(app.main_mut());
         start_running_systems();
 
         Self {
@@ -610,7 +616,9 @@ impl Plugin for AzaleaPlugin {
 /// You can create your app with `App::new()`, but don't forget to add
 /// [`DefaultPlugins`].
 #[doc(hidden)]
-pub fn start_ecs_runner(app: &mut SubApp) -> (Arc<Mutex<World>>, impl FnOnce()) {
+pub fn start_ecs_runner(
+    app: &mut SubApp,
+) -> (Arc<Mutex<World>>, impl FnOnce(), oneshot::Receiver<AppExit>) {
     // this block is based on Bevy's default runner:
     // https://github.com/bevyengine/bevy/blob/390877cdae7a17095a75c8f9f1b4241fe5047e83/crates/bevy_app/src/schedule_runner.rs#L77-L85
     if app.plugins_state() != PluginsState::Cleaned {
@@ -632,14 +640,26 @@ pub fn start_ecs_runner(app: &mut SubApp) -> (Arc<Mutex<World>>, impl FnOnce())
 
     let ecs_clone = ecs.clone();
     let outer_schedule_label = *app.update_schedule.as_ref().unwrap();
+
+    let (appexit_tx, appexit_rx) = oneshot::channel();
     let start_running_systems = move || {
-        tokio::spawn(run_schedule_loop(ecs_clone, outer_schedule_label));
+        tokio::spawn(async move {
+            let appexit = run_schedule_loop(ecs_clone, outer_schedule_label).await;
+            appexit_tx.send(appexit)
+        });
     };
 
-    (ecs, start_running_systems)
+    (ecs, start_running_systems, appexit_rx)
 }
 
-async fn run_schedule_loop(ecs: Arc<Mutex<World>>, outer_schedule_label: InternedScheduleLabel) {
+/// Runs the `Update` schedule 60 times per second and the `GameTick` schedule
+/// 20 times per second.
+///
+/// Exits when we receive an `AppExit` event.
+async fn run_schedule_loop(
+    ecs: Arc<Mutex<World>>,
+    outer_schedule_label: InternedScheduleLabel,
+) -> AppExit {
     let mut last_update: Option<Instant> = None;
     let mut last_tick: Option<Instant> = None;
 
@@ -687,9 +707,38 @@ async fn run_schedule_loop(ecs: Arc<Mutex<World>>, outer_schedule_label: Interne
         }
 
         ecs.clear_trackers();
+        if let Some(exit) = should_exit(&mut ecs) {
+            // it's possible for references to the World to stay around, so we clear the ecs
+            ecs.clear_all();
+            // ^ note that this also forcefully disconnects all of our bots without sending
+            // a disconnect packet (which is fine because we want to disconnect immediately)
+
+            return exit;
+        }
     }
 }
 
+/// Checks whether the [`AppExit`] event was sent, and if so returns it.
+///
+/// This is based on Bevy's `should_exit` function: https://github.com/bevyengine/bevy/blob/b9fd7680e78c4073dfc90fcfdc0867534d92abe0/crates/bevy_app/src/app.rs#L1292
+fn should_exit(ecs: &mut World) -> Option<AppExit> {
+    let mut reader = EventCursor::default();
+
+    let events = ecs.get_resource::<Events<AppExit>>()?;
+    let mut events = reader.read(events);
+
+    if events.len() != 0 {
+        return Some(
+            events
+                .find(|exit| exit.is_error())
+                .cloned()
+                .unwrap_or(AppExit::Success),
+        );
+    }
+
+    None
+}
+
 pub struct AmbiguityLoggerPlugin;
 impl Plugin for AmbiguityLoggerPlugin {
     fn build(&self, app: &mut App) {
diff --git a/azalea/examples/testbot/commands/debug.rs b/azalea/examples/testbot/commands/debug.rs
index b3a8b419..06a49ce2 100644
--- a/azalea/examples/testbot/commands/debug.rs
+++ b/azalea/examples/testbot/commands/debug.rs
@@ -1,6 +1,6 @@
 //! Commands for debugging and getting the current state of the bot.
 
-use std::{env, fs::File, io::Write, process, thread, time::Duration};
+use std::{env, fs::File, io::Write, thread, time::Duration};
 
 use azalea::{
     BlockPos,
@@ -17,6 +17,7 @@ use azalea_core::hit_result::HitResult;
 use azalea_entity::EntityKindComponent;
 use azalea_inventory::components::MaxStackSize;
 use azalea_world::InstanceContainer;
+use bevy_app::AppExit;
 use bevy_ecs::event::Events;
 use parking_lot::Mutex;
 
@@ -316,10 +317,11 @@ pub fn register(commands: &mut CommandDispatcher<Mutex<CommandSource>>) {
 
         source.bot.disconnect();
 
+        let source = ctx.source.clone();
         thread::spawn(move || {
             thread::sleep(Duration::from_secs(1));
 
-            process::exit(0);
+            source.lock().bot.ecs.lock().send_event(AppExit::Success);
         });
 
         1
diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs
index 049fe2f0..7b4ec5ae 100644
--- a/azalea/src/lib.rs
+++ b/azalea/src/lib.rs
@@ -1,6 +1,5 @@
 #![doc = include_str!("../README.md")]
 #![feature(type_changing_struct_update)]
-#![feature(never_type)]
 
 pub mod accept_resource_packs;
 pub mod auto_respawn;
@@ -33,6 +32,7 @@ pub use azalea_protocol as protocol;
 pub use azalea_registry as registry;
 pub use azalea_world as world;
 pub use bevy_app as app;
+use bevy_app::AppExit;
 pub use bevy_ecs as ecs;
 pub use bot::*;
 use ecs::component::Component;
@@ -215,7 +215,7 @@ where
         mut self,
         account: Account,
         address: impl TryInto<ServerAddress>,
-    ) -> Result<!, StartError> {
+    ) -> Result<AppExit, StartError> {
         self.swarm.accounts = vec![(account, JoinOpts::default())];
         if self.swarm.states.is_empty() {
             self.swarm.states = vec![S::default()];
@@ -230,7 +230,7 @@ where
         account: Account,
         address: impl TryInto<ServerAddress>,
         opts: JoinOpts,
-    ) -> Result<!, StartError> {
+    ) -> Result<AppExit, StartError> {
         self.swarm.accounts = vec![(account, opts.clone())];
         if self.swarm.states.is_empty() {
             self.swarm.states = vec![S::default()];
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index b39d6911..25f11f1e 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -28,7 +28,7 @@ use azalea_client::{
 use azalea_entity::LocalEntity;
 use azalea_protocol::{ServerAddress, resolver};
 use azalea_world::InstanceContainer;
-use bevy_app::{App, PluginGroup, PluginGroupBuilder, Plugins, SubApp};
+use bevy_app::{App, AppExit, PluginGroup, PluginGroupBuilder, Plugins, SubApp};
 use bevy_ecs::prelude::*;
 use futures::future::{BoxFuture, join_all};
 use parking_lot::{Mutex, RwLock};
@@ -399,7 +399,7 @@ where
     /// that implements `TryInto<ServerAddress>`.
     ///
     /// [`ServerAddress`]: azalea_protocol::ServerAddress
-    pub async fn start(self, address: impl TryInto<ServerAddress>) -> Result<!, StartError> {
+    pub async fn start(self, address: impl TryInto<ServerAddress>) -> Result<AppExit, StartError> {
         // convert the TryInto<ServerAddress> into a ServerAddress
         let address: ServerAddress = match address.try_into() {
             Ok(address) => address,
@@ -416,7 +416,7 @@ where
         mut self,
         address: impl TryInto<ServerAddress>,
         default_join_opts: JoinOpts,
-    ) -> Result<!, StartError> {
+    ) -> Result<AppExit, StartError> {
         assert_eq!(
             self.accounts.len(),
             self.states.len(),
@@ -448,7 +448,7 @@ where
 
         let main_schedule_label = self.app.update_schedule.unwrap();
 
-        let (ecs_lock, start_running_systems) = start_ecs_runner(&mut self.app);
+        let (ecs_lock, start_running_systems, appexit_rx) = start_ecs_runner(&mut self.app);
 
         let swarm = Swarm {
             ecs_lock: ecs_lock.clone(),
@@ -521,7 +521,7 @@ where
 
         // Watch swarm_rx and send those events to the swarm_handle.
         let swarm_clone = swarm.clone();
-        tokio::spawn(async move {
+        let swarm_handler_task = tokio::spawn(async move {
             while let Some(event) = swarm_rx.recv().await {
                 if let Some(swarm_handler) = &self.swarm_handler {
                     tokio::spawn((swarm_handler)(
@@ -531,61 +531,72 @@ where
                     ));
                 }
             }
+
+            unreachable!(
+                "The `Swarm` here contains a sender for the `SwarmEvent`s, so swarm_rx.recv() will never fail"
+            );
         });
 
         // bot events
-        while let Some((Some(first_event), first_bot)) = bots_rx.recv().await {
-            if bots_rx.len() > 1_000 {
-                static WARNED: AtomicBool = AtomicBool::new(false);
-                if !WARNED.swap(true, atomic::Ordering::Relaxed) {
-                    warn!("the Client Event channel has more than 1000 items!")
+        let client_handler_task = tokio::spawn(async move {
+            while let Some((Some(first_event), first_bot)) = bots_rx.recv().await {
+                if bots_rx.len() > 1_000 {
+                    static WARNED: AtomicBool = AtomicBool::new(false);
+                    if !WARNED.swap(true, atomic::Ordering::Relaxed) {
+                        warn!("the Client Event channel has more than 1000 items!")
+                    }
                 }
-            }
 
-            if let Some(handler) = &self.handler {
-                let ecs_mutex = first_bot.ecs.clone();
-                let mut ecs = ecs_mutex.lock();
-                let mut query = ecs.query::<Option<&S>>();
-                let Ok(Some(first_bot_state)) = query.get(&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(),
-                        first_bot.entity
-                    );
-                    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()));
-
-                // this makes it not have to keep locking the ecs
-                let mut states = HashMap::new();
-                states.insert(first_bot_entity, first_bot_state);
-                while let Ok((Some(event), bot)) = bots_rx.try_recv() {
-                    let state = match states.entry(bot.entity) {
-                        hash_map::Entry::Occupied(e) => e.into_mut(),
-                        hash_map::Entry::Vacant(e) => {
-                            let Ok(Some(state)) = query.get(&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(),
-                                    bot.entity
-                                );
-                                continue;
-                            };
-                            let state = state.clone();
-                            e.insert(state)
-                        }
+                if let Some(handler) = &self.handler {
+                    let ecs_mutex = first_bot.ecs.clone();
+                    let mut ecs = ecs_mutex.lock();
+                    let mut query = ecs.query::<Option<&S>>();
+                    let Ok(Some(first_bot_state)) = query.get(&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(),
+                            first_bot.entity
+                        );
+                        continue;
                     };
-                    tokio::spawn((handler)(bot, event, state.clone()));
+                    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()));
+
+                    // this makes it not have to keep locking the ecs
+                    let mut states = HashMap::new();
+                    states.insert(first_bot_entity, first_bot_state);
+                    while let Ok((Some(event), bot)) = bots_rx.try_recv() {
+                        let state = match states.entry(bot.entity) {
+                            hash_map::Entry::Occupied(e) => e.into_mut(),
+                            hash_map::Entry::Vacant(e) => {
+                                let Ok(Some(state)) = query.get(&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(),
+                                        bot.entity
+                                    );
+                                    continue;
+                                };
+                                let state = state.clone();
+                                e.insert(state)
+                            }
+                        };
+                        tokio::spawn((handler)(bot, event, state.clone()));
+                    }
                 }
             }
-        }
+        });
 
-        unreachable!(
-            "bots_rx.recv() should never be None because the bots_tx channel is never closed"
-        );
+        let appexit = appexit_rx
+            .await
+            .expect("appexit_tx shouldn't be dropped by the ECS runner before sending");
+
+        swarm_handler_task.abort();
+        client_handler_task.abort();
+
+        Ok(appexit)
     }
 }