make ClientBuilder use SwarmBuilder internally

2 files changed, 32 insertions, 92 deletions

diff --git a/azalea/src/lib.rs b/azalea/src/lib.rs
index 18a3c314..0a243ada 100644
--- a/azalea/src/lib.rs
+++ b/azalea/src/lib.rs
@@ -14,7 +14,7 @@ pub mod pathfinder;
 pub mod prelude;
 pub mod swarm;
 
-use app::{App, Plugins};
+use app::Plugins;
 pub use azalea_auth as auth;
 pub use azalea_block as blocks;
 pub use azalea_brigadier as brigadier;
@@ -34,12 +34,9 @@ pub use azalea_world as world;
 pub use bot::*;
 use ecs::component::Component;
 use futures::{future::BoxFuture, Future};
-use protocol::{
-    resolver::{self, ResolverError},
-    ServerAddress,
-};
+use protocol::{resolver::ResolverError, ServerAddress};
+use swarm::SwarmBuilder;
 use thiserror::Error;
-use tokio::sync::mpsc;
 
 pub use bevy_app as app;
 pub use bevy_ecs as ecs;
@@ -54,8 +51,6 @@ pub enum StartError {
     InvalidAddress,
     #[error(transparent)]
     ResolveAddress(#[from] ResolverError),
-    #[error("Join error: {0}")]
-    Join(#[from] azalea_client::JoinError),
 }
 
 /// A builder for creating new [`Client`]s. This is the recommended way of
@@ -80,10 +75,10 @@ pub struct ClientBuilder<S>
 where
     S: Default + Send + Sync + Clone + Component + 'static,
 {
-    app: App,
-    /// The function that's called every time a bot receives an [`Event`].
-    handler: Option<BoxHandleFn<S>>,
-    state: S,
+    /// Internally, ClientBuilder is just a wrapper over SwarmBuilder since it's
+    /// technically just a subset of it so we can avoid duplicating code this
+    /// way.
+    swarm: SwarmBuilder<S, swarm::NoSwarmState>,
 }
 impl ClientBuilder<NoState> {
     /// Start building a client that can join the world.
@@ -120,11 +115,7 @@ impl ClientBuilder<NoState> {
     #[must_use]
     pub fn new_without_plugins() -> ClientBuilder<NoState> {
         Self {
-            // we create the app here so plugins can add onto it.
-            // the schedules won't run until [`Self::start`] is called.
-            app: App::new(),
-            handler: None,
-            state: NoState,
+            swarm: SwarmBuilder::new_without_plugins(),
         }
     }
 
@@ -151,11 +142,7 @@ impl ClientBuilder<NoState> {
         Fut: Future<Output = Result<(), anyhow::Error>> + Send + 'static,
     {
         ClientBuilder {
-            handler: Some(Box::new(move |bot, event, state| {
-                Box::pin(handler(bot, event, state))
-            })),
-            state: S::default(),
-            ..self
+            swarm: self.swarm.set_handler(handler),
         }
     }
 }
@@ -166,62 +153,36 @@ where
     /// Set the client state instead of initializing defaults.
     #[must_use]
     pub fn set_state(mut self, state: S) -> Self {
-        self.state = state;
+        self.swarm.states = vec![state];
         self
     }
     /// Add a group of plugins to the client.
     #[must_use]
     pub fn add_plugins<M>(mut self, plugins: impl Plugins<M>) -> Self {
-        self.app.add_plugins(plugins);
+        self.swarm = self.swarm.add_plugins(plugins);
         self
     }
 
     /// Build this `ClientBuilder` into an actual [`Client`] and join the given
-    /// server.
+    /// server. If the client can't join, it'll keep retrying forever until it
+    /// can.
     ///
     /// The `address` argument can be a `&str`, [`ServerAddress`], or anything
     /// that implements `TryInto<ServerAddress>`.
     ///
+    /// # Errors
+    ///
+    /// This will error if the given address is invalid or couldn't be resolved
+    /// to a Minecraft server.
+    ///
     /// [`ServerAddress`]: azalea_protocol::ServerAddress
     pub async fn start(
-        self,
+        mut self,
         account: Account,
         address: impl TryInto<ServerAddress>,
     ) -> Result<(), StartError> {
-        let address: ServerAddress = address.try_into().map_err(|_| JoinError::InvalidAddress)?;
-        let resolved_address = resolver::resolve_address(&address).await?;
-
-        // An event that causes the schedule to run. This is only used internally.
-        let (run_schedule_sender, run_schedule_receiver) = mpsc::unbounded_channel();
-
-        let main_schedule_label = self.app.main_schedule_label;
-
-        let ecs_lock =
-            start_ecs_runner(self.app, run_schedule_receiver, run_schedule_sender.clone());
-
-        // run the main schedule so the startup systems run
-        {
-            let mut ecs = ecs_lock.lock();
-            ecs.run_schedule(main_schedule_label);
-            ecs.clear_trackers();
-        }
-
-        let (bot, mut rx) = Client::start_client(
-            ecs_lock,
-            &account,
-            &address,
-            &resolved_address,
-            run_schedule_sender,
-        )
-        .await?;
-
-        while let Some(event) = rx.recv().await {
-            if let Some(handler) = &self.handler {
-                tokio::spawn((handler)(bot.clone(), event.clone(), self.state.clone()));
-            }
-        }
-
-        Ok(())
+        self.swarm.accounts = vec![account];
+        self.swarm.start(address).await
     }
 }
 impl Default for ClientBuilder<NoState> {
diff --git a/azalea/src/swarm/mod.rs b/azalea/src/swarm/mod.rs
index 41a54fdd..5c7a04be 100644
--- a/azalea/src/swarm/mod.rs
+++ b/azalea/src/swarm/mod.rs
@@ -7,22 +7,17 @@ pub mod prelude;
 use azalea_client::{
     chat::ChatPacket, start_ecs_runner, Account, Client, DefaultPlugins, Event, JoinError,
 };
-use azalea_protocol::{
-    connect::ConnectionError,
-    resolver::{self, ResolverError},
-    ServerAddress,
-};
+use azalea_protocol::{resolver, ServerAddress};
 use azalea_world::InstanceContainer;
 use bevy_app::{App, PluginGroup, PluginGroupBuilder, Plugins};
 use bevy_ecs::{component::Component, entity::Entity, system::Resource, world::World};
 use futures::future::{join_all, BoxFuture};
 use parking_lot::{Mutex, RwLock};
 use std::{collections::HashMap, future::Future, net::SocketAddr, sync::Arc, time::Duration};
-use thiserror::Error;
 use tokio::sync::mpsc;
 use tracing::error;
 
-use crate::{BoxHandleFn, DefaultBotPlugins, HandleFn, NoState};
+use crate::{BoxHandleFn, DefaultBotPlugins, HandleFn, NoState, StartError};
 
 /// A swarm is a way to conveniently control many bots at once, while also
 /// being able to control bots at an individual level when desired.
@@ -55,25 +50,25 @@ where
     S: Send + Sync + Clone + Component + 'static,
     SS: Default + Send + Sync + Clone + Resource + 'static,
 {
-    app: App,
+    pub(crate) app: App,
     /// The accounts that are going to join the server.
-    accounts: Vec<Account>,
+    pub(crate) accounts: Vec<Account>,
     /// The individual bot states. This must be the same length as `accounts`,
     /// since each bot gets one state.
-    states: Vec<S>,
+    pub(crate) states: Vec<S>,
     /// The state for the overall swarm.
-    swarm_state: SS,
+    pub(crate) swarm_state: SS,
     /// The function that's called every time a bot receives an [`Event`].
-    handler: Option<BoxHandleFn<S>>,
+    pub(crate) handler: Option<BoxHandleFn<S>>,
     /// The function that's called every time the swarm receives a
     /// [`SwarmEvent`].
-    swarm_handler: Option<BoxSwarmHandleFn<SS>>,
+    pub(crate) swarm_handler: Option<BoxSwarmHandleFn<SS>>,
 
     /// How long we should wait between each bot joining the server. Set to
     /// None to have every bot connect at the same time. None is different than
     /// a duration of 0, since if a duration is present the bots will wait for
     /// the previous one to be ready.
-    join_delay: Option<std::time::Duration>,
+    pub(crate) join_delay: Option<std::time::Duration>,
 }
 impl SwarmBuilder<NoState, NoSwarmState> {
     /// Start creating the swarm.
@@ -297,7 +292,7 @@ where
     /// that implements `TryInto<ServerAddress>`.
     ///
     /// [`ServerAddress`]: azalea_protocol::ServerAddress
-    pub async fn start(self, address: impl TryInto<ServerAddress>) -> Result<(), SwarmStartError> {
+    pub async fn start(self, address: impl TryInto<ServerAddress>) -> Result<(), StartError> {
         assert_eq!(
             self.accounts.len(),
             self.states.len(),
@@ -307,7 +302,7 @@ where
         // convert the TryInto<ServerAddress> into a ServerAddress
         let address: ServerAddress = match address.try_into() {
             Ok(address) => address,
-            Err(_) => return Err(SwarmStartError::InvalidAddress),
+            Err(_) => return Err(StartError::InvalidAddress),
         };
 
         // resolve the address
@@ -450,16 +445,6 @@ pub type SwarmHandleFn<SS, Fut> = fn(Swarm, SwarmEvent, SS) -> Fut;
 pub type BoxSwarmHandleFn<SS> =
     Box<dyn Fn(Swarm, SwarmEvent, SS) -> BoxFuture<'static, Result<(), anyhow::Error>> + Send>;
 
-#[derive(Error, Debug)]
-pub enum SwarmStartError {
-    #[error("Invalid address")]
-    InvalidAddress,
-    #[error(transparent)]
-    ResolveAddress(#[from] ResolverError),
-    #[error("Join error: {0}")]
-    Join(#[from] azalea_client::JoinError),
-}
-
 /// Make a bot [`Swarm`].
 ///
 /// [`Swarm`]: struct.Swarm.html
@@ -632,12 +617,6 @@ impl IntoIterator for Swarm {
     }
 }
 
-impl From<ConnectionError> for SwarmStartError {
-    fn from(e: ConnectionError) -> Self {
-        SwarmStartError::from(JoinError::from(e))
-    }
-}
-
 /// This plugin group will add all the default plugins necessary for swarms to
 /// work.
 pub struct DefaultSwarmPlugins;