add Client::exit and Swarm::exit, and write related docs

4 files changed, 72 insertions, 3 deletions

diff --git a/azalea/src/builder.rs b/azalea/src/builder.rs
index db51de63..0cb4abbf 100644
--- a/azalea/src/builder.rs
+++ b/azalea/src/builder.rs
@@ -150,7 +150,13 @@ where
     /// If this function isn't called, then our client will reconnect after
     /// [`DEFAULT_RECONNECT_DELAY`].
     ///
+    /// Note that disabling auto-reconnecting will not make
+    /// [`ClientBuilder::start`] return on disconnect, because Azalea will keep
+    /// the internal swarm around forever until it's forcibly exited. To learn
+    /// how to do that, see [`Client::exit`].
+    ///
     /// [`DEFAULT_RECONNECT_DELAY`]: crate::auto_reconnect::DEFAULT_RECONNECT_DELAY
+    /// [`Client::exit`]: crate::Client::exit
     #[must_use]
     pub fn reconnect_after(mut self, delay: impl Into<Option<Duration>>) -> Self {
         self.swarm.reconnect_after = delay.into();
diff --git a/azalea/src/client_impl/mod.rs b/azalea/src/client_impl/mod.rs
index 6725174c..66fa1a6f 100644
--- a/azalea/src/client_impl/mod.rs
+++ b/azalea/src/client_impl/mod.rs
@@ -26,7 +26,7 @@ use azalea_protocol::{
 };
 use azalea_registry::{DataRegistryKeyRef, identifier::Identifier};
 use azalea_world::{PartialWorld, World, WorldName};
-use bevy_app::App;
+use bevy_app::{App, AppExit};
 use bevy_ecs::{entity::Entity, resource::Resource, world::Mut};
 use parking_lot::RwLock;
 use tokio::sync::mpsc;
@@ -243,6 +243,9 @@ impl Client {
     ///
     /// The OwnedReadHalf for the TCP connection is in one of the tasks, so it
     /// automatically closes the connection when that's dropped.
+    ///
+    /// Note that this will not return from your client builder. If you need
+    /// that, consider using [`Self::exit`] instead.
     pub fn disconnect(&self) {
         self.ecs.write().write_message(DisconnectEvent {
             entity: self.entity,
@@ -250,6 +253,51 @@ impl Client {
         });
     }
 
+    /// End the entire client or swarm, and return from
+    /// [`ClientBuilder::start`] or [`SwarmBuilder::start`].
+    ///
+    /// You should typically avoid calling this if you intend on creating the
+    /// client again, because creating an entirely new swarm can be a
+    /// relatively expensive process.
+    ///
+    /// If you only want to change the server that the bots are connecting to,
+    /// it may be better to access the client's internal [`Swarm`] and call
+    /// [`Swarm::add_with_opts`] with a different server address.
+    ///
+    /// For convenience, this is also duplicated on `Swarm` as [`Swarm::exit`].
+    ///
+    /// [`ClientBuilder::start`]: crate::ClientBuilder::start
+    /// [`SwarmBuilder::start`]: crate::swarm::SwarmBuilder::start
+    /// [`Swarm`]: crate::swarm::Swarm
+    /// [`Swarm::add_with_opts`]: crate::swarm::Swarm::add_with_opts
+    /// [`Swarm::exit`]: crate::swarm::Swarm::exit
+    ///
+    /// ```
+    /// // a bot that joins a server and prints "done!" when it's disconnected or if it fails to connect.
+    /// use azalea::{NoState, prelude::*};
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let account = Account::offline("bot");
+    ///     ClientBuilder::new()
+    ///         .set_handler(handle)
+    ///         .start(account, "localhost")
+    ///         .await;
+    ///     println!("done!");
+    /// }
+    /// async fn handle(bot: Client, event: Event, _state: NoState) -> anyhow::Result<()> {
+    ///     match event {
+    ///         Event::Disconnect(_) | Event::ConnectionFailed(_) => {
+    ///             bot.exit();
+    ///         }
+    ///         _ => {}
+    ///     }
+    ///     Ok(())
+    /// }
+    /// ```
+    pub fn exit(&self) {
+        self.ecs.write().write_message(AppExit::Success);
+    }
+
     pub fn with_raw_connection<R>(&self, f: impl FnOnce(&RawConnection) -> R) -> R {
         self.query_self::<&RawConnection, _>(f)
     }
diff --git a/azalea/src/events.rs b/azalea/src/events.rs
index fcd47e2f..7d625456 100644
--- a/azalea/src/events.rs
+++ b/azalea/src/events.rs
@@ -161,7 +161,7 @@ impl Plugin for EventsPlugin {
                 keepalive_listener,
                 death_listener.after(azalea_client::packet::death_event_on_0_health),
                 disconnect_listener,
-                connection_failed_listener,
+                connection_failed_listener.after(azalea_client::join::poll_create_connection_task),
                 receive_chunk_listener,
             ),
         )
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index 4ecb0726..de4d2ebe 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -16,7 +16,7 @@ use azalea_client::{account::Account, chat::ChatPacket, join::ConnectOpts};
 use azalea_entity::LocalEntity;
 use azalea_protocol::address::ResolvedAddr;
 use azalea_world::Worlds;
-use bevy_app::{PluginGroup, PluginGroupBuilder};
+use bevy_app::{AppExit, PluginGroup, PluginGroupBuilder};
 use bevy_ecs::prelude::*;
 pub use builder::SwarmBuilder;
 use futures::future::BoxFuture;
@@ -291,6 +291,21 @@ impl Swarm {
         let mut query = ecs.query_filtered::<Entity, With<LocalEntity>>();
         query.iter(&ecs).collect::<Box<[Entity]>>()
     }
+
+    /// End the entire swarm and return from [`SwarmBuilder::start`].
+    ///
+    /// You should typically avoid calling this if you intend on creating the
+    /// swarm again, because creating an entirely new swarm can be a
+    /// relatively expensive process.
+    ///
+    /// If you only want to change the server that the bots are connecting to,
+    /// it may be better to call [`Swarm::add_with_opts`] with a different
+    /// server address.
+    ///
+    /// This is also implemented on [`Client`] as [`Client::exit`].
+    pub fn exit(&self) {
+        self.ecs.write().write_message(AppExit::Success);
+    }
 }
 
 impl IntoIterator for Swarm {