Add AutoReconnectPlugin (#221)

* add AutoReconnectPlugin

* merge main

* start simplifying swarm internals

* fix Swarm::into_iter, handler functions, DisconnectEvent, and add some more docs

* add ClientBuilder/SwarmBuilder::reconnect_after

* fix a doctest

* reword SwarmEvent::Disconnect doc

* better behavior when we try to join twice

* reconnect on ConnectionFailedEvent too

* autoreconnect is less breaking now

5 files changed, 252 insertions, 34 deletions

diff --git a/azalea-client/src/plugins/auto_reconnect.rs b/azalea-client/src/plugins/auto_reconnect.rs
new file mode 100644
index 00000000..280aaa65
--- /dev/null
+++ b/azalea-client/src/plugins/auto_reconnect.rs
@@ -0,0 +1,138 @@
+//! Auto-reconnect to the server when the client is kicked.
+//!
+//! See [`AutoReconnectPlugin`] for more information.
+
+use std::time::{Duration, Instant};
+
+use bevy_app::prelude::*;
+use bevy_ecs::prelude::*;
+
+use super::{
+    disconnect::DisconnectEvent,
+    events::LocalPlayerEvents,
+    join::{ConnectOpts, ConnectionFailedEvent, StartJoinServerEvent},
+};
+use crate::Account;
+
+/// The default delay that Azalea will use for reconnecting our clients. See
+/// [`AutoReconnectPlugin`] for more information.
+pub const DEFAULT_RECONNECT_DELAY: Duration = Duration::from_secs(5);
+
+/// A default plugin that makes clients automatically rejoin the server when
+/// they're disconnected. The reconnect delay is configurable globally or
+/// per-client with the [`AutoReconnectDelay`] resource/component. Auto
+/// reconnecting can be disabled by removing the resource from the ECS.
+///
+/// The delay defaults to [`DEFAULT_RECONNECT_DELAY`].
+pub struct AutoReconnectPlugin;
+impl Plugin for AutoReconnectPlugin {
+    fn build(&self, app: &mut App) {
+        app.insert_resource(AutoReconnectDelay::new(DEFAULT_RECONNECT_DELAY))
+            .add_systems(
+                Update,
+                (start_rejoin_on_disconnect, rejoin_after_delay)
+                    .chain()
+                    .before(super::join::handle_start_join_server_event),
+            );
+    }
+}
+
+pub fn start_rejoin_on_disconnect(
+    mut commands: Commands,
+    mut disconnect_events: EventReader<DisconnectEvent>,
+    mut connection_failed_events: EventReader<ConnectionFailedEvent>,
+    auto_reconnect_delay_res: Option<Res<AutoReconnectDelay>>,
+    auto_reconnect_delay_query: Query<&AutoReconnectDelay>,
+) {
+    for entity in disconnect_events
+        .read()
+        .map(|e| e.entity)
+        .chain(connection_failed_events.read().map(|e| e.entity))
+    {
+        let Some(delay) = get_delay(
+            &auto_reconnect_delay_res,
+            auto_reconnect_delay_query,
+            entity,
+        ) else {
+            // no auto reconnect
+            continue;
+        };
+
+        let reconnect_after = Instant::now() + delay;
+        commands.entity(entity).insert(InternalReconnectAfter {
+            instant: reconnect_after,
+        });
+    }
+}
+
+fn get_delay(
+    auto_reconnect_delay_res: &Option<Res<AutoReconnectDelay>>,
+    auto_reconnect_delay_query: Query<&AutoReconnectDelay>,
+    entity: Entity,
+) -> Option<Duration> {
+    if let Ok(c) = auto_reconnect_delay_query.get(entity) {
+        Some(c.delay)
+    } else if let Some(r) = &auto_reconnect_delay_res {
+        Some(r.delay)
+    } else {
+        None
+    }
+}
+
+pub fn rejoin_after_delay(
+    mut commands: Commands,
+    mut join_events: EventWriter<StartJoinServerEvent>,
+    query: Query<(
+        Entity,
+        &InternalReconnectAfter,
+        &Account,
+        &ConnectOpts,
+        Option<&LocalPlayerEvents>,
+    )>,
+) {
+    for (entity, reconnect_after, account, connect_opts, local_player_events) in query.iter() {
+        if Instant::now() >= reconnect_after.instant {
+            // don't keep trying to reconnect
+            commands.entity(entity).remove::<InternalReconnectAfter>();
+
+            // our Entity will be reused since the account has the same uuid
+            join_events.write(StartJoinServerEvent {
+                account: account.clone(),
+                connect_opts: connect_opts.clone(),
+                // not actually necessary since we're reusing the same entity and LocalPlayerEvents
+                // isn't removed, but this is more readable and just in case it's changed in the
+                // future
+                event_sender: local_player_events.map(|e| e.0.clone()),
+                start_join_callback_tx: None,
+            });
+        }
+    }
+}
+
+/// A resource *and* component that indicates how long to wait before
+/// reconnecting when we're kicked.
+///
+/// Initially, it's a resource in the ECS set to 5 seconds. You can modify
+/// the resource to update the global reconnect delay, or insert it as a
+/// component to set the individual delay for a single client.
+///
+/// You can also remove this resource from the ECS to disable the default
+/// auto-reconnecting behavior. Inserting the resource/component again will not
+/// make clients that were already disconnected automatically reconnect.
+#[derive(Resource, Component, Debug, Clone)]
+pub struct AutoReconnectDelay {
+    pub delay: Duration,
+}
+impl AutoReconnectDelay {
+    pub fn new(delay: Duration) -> Self {
+        Self { delay }
+    }
+}
+
+/// This is inserted when we're disconnected and indicates when we'll reconnect.
+///
+/// This is set based on [`AutoReconnectDelay`].
+#[derive(Component, Debug, Clone)]
+pub struct InternalReconnectAfter {
+    pub instant: Instant,
+}
diff --git a/azalea-client/src/plugins/disconnect.rs b/azalea-client/src/plugins/disconnect.rs
index 343c25d8..987007c2 100644
--- a/azalea-client/src/plugins/disconnect.rs
+++ b/azalea-client/src/plugins/disconnect.rs
@@ -7,10 +7,7 @@ use bevy_ecs::prelude::*;
 use derive_more::Deref;
 use tracing::info;
 
-use crate::{
-    InstanceHolder, client::JoinedClientBundle, connection::RawConnection,
-    events::LocalPlayerEvents,
-};
+use crate::{InstanceHolder, client::JoinedClientBundle, connection::RawConnection};
 
 pub struct DisconnectPlugin;
 impl Plugin for DisconnectPlugin {
@@ -19,15 +16,28 @@ impl Plugin for DisconnectPlugin {
             PostUpdate,
             (
                 update_read_packets_task_running_component,
-                disconnect_on_connection_dead,
                 remove_components_from_disconnected_players,
+                // this happens after `remove_components_from_disconnected_players` since that
+                // system removes `IsConnectionAlive`, which ensures that
+                // `DisconnectEvent` won't get called again from
+                // `disconnect_on_connection_dead`
+                disconnect_on_connection_dead,
             )
                 .chain(),
         );
     }
 }
 
-/// An event sent when a client is getting disconnected.
+/// An event sent when a client got disconnected from the server.
+///
+/// If the client was kicked with a reason, that reason will be present in the
+/// [`reason`](DisconnectEvent::reason) field.
+///
+/// This event won't be sent if creating the initial connection to the server
+/// failed, for that see [`ConnectionFailedEvent`].
+///
+/// [`ConnectionFailedEvent`]: crate::join::ConnectionFailedEvent
+
 #[derive(Event)]
 pub struct DisconnectEvent {
     pub entity: Entity,
@@ -59,8 +69,8 @@ pub fn remove_components_from_disconnected_players(
             .remove::<InLoadedChunk>()
             // this makes it close the tcp connection
             .remove::<RawConnection>()
-            // swarm detects when this tx gets dropped to fire SwarmEvent::Disconnect
-            .remove::<LocalPlayerEvents>();
+            // this makes it not send DisconnectEvent again
+            .remove::<IsConnectionAlive>();
         // note that we don't remove the client from the ECS, so if they decide
         // to reconnect they'll keep their state
 
diff --git a/azalea-client/src/plugins/join.rs b/azalea-client/src/plugins/join.rs
index 3f47d90c..e31c64c4 100644
--- a/azalea-client/src/plugins/join.rs
+++ b/azalea-client/src/plugins/join.rs
@@ -1,4 +1,4 @@
-use std::{net::SocketAddr, sync::Arc};
+use std::{io, net::SocketAddr, sync::Arc};
 
 use azalea_entity::{LocalEntity, indexing::EntityUuidIndex};
 use azalea_protocol::{
@@ -29,22 +29,54 @@ use crate::{
 pub struct JoinPlugin;
 impl Plugin for JoinPlugin {
     fn build(&self, app: &mut App) {
-        app.add_event::<StartJoinServerEvent>().add_systems(
-            Update,
-            (handle_start_join_server_event, poll_create_connection_task),
-        );
+        app.add_event::<StartJoinServerEvent>()
+            .add_event::<ConnectionFailedEvent>()
+            .add_systems(
+                Update,
+                (
+                    handle_start_join_server_event.before(super::login::poll_auth_task),
+                    poll_create_connection_task,
+                    handle_connection_failed_events,
+                )
+                    .chain(),
+            );
     }
 }
 
+/// An event to make a client join the server and be added to our swarm.
+///
+/// This won't do anything if a client with the Account UUID is already
+/// connected to the server.
 #[derive(Event, Debug)]
 pub struct StartJoinServerEvent {
     pub account: Account,
+    pub connect_opts: ConnectOpts,
+    pub event_sender: Option<mpsc::UnboundedSender<crate::Event>>,
+
+    pub start_join_callback_tx: Option<StartJoinCallback>,
+}
+
+/// Options for how the connection to the server will be made. These are
+/// persisted on reconnects.
+///
+/// This is inserted as a component on clients to make auto-reconnecting work.
+#[derive(Debug, Clone, Component)]
+pub struct ConnectOpts {
     pub address: ServerAddress,
     pub resolved_address: SocketAddr,
     pub proxy: Option<Proxy>,
-    pub event_sender: Option<mpsc::UnboundedSender<crate::Event>>,
+}
 
-    pub start_join_callback_tx: Option<StartJoinCallback>,
+/// An event that's sent when creating the TCP connection and sending the first
+/// packet fails.
+///
+/// This isn't sent if we're kicked later, see [`DisconnectEvent`].
+///
+/// [`DisconnectEvent`]: crate::disconnect::DisconnectEvent
+#[derive(Event)]
+pub struct ConnectionFailedEvent {
+    pub entity: Entity,
+    pub error: ConnectionError,
 }
 
 // this is mpsc instead of oneshot so it can be cloned (since it's sent in an
@@ -56,11 +88,30 @@ pub fn handle_start_join_server_event(
     mut commands: Commands,
     mut events: EventReader<StartJoinServerEvent>,
     mut entity_uuid_index: ResMut<EntityUuidIndex>,
+    connection_query: Query<&RawConnection>,
 ) {
     for event in events.read() {
         let uuid = event.account.uuid_or_offline();
         let entity = if let Some(entity) = entity_uuid_index.get(&uuid) {
             debug!("Reusing entity {entity:?} for client");
+
+            // check if it's already connected
+            if let Ok(conn) = connection_query.get(entity) {
+                if conn.is_alive() {
+                    if let Some(start_join_callback_tx) = &event.start_join_callback_tx {
+                        warn!(
+                            "Received StartJoinServerEvent for {entity:?} but it's already connected. Ignoring the event but replying with Ok."
+                        );
+                        let _ = start_join_callback_tx.0.send(Ok(entity));
+                    } else {
+                        warn!(
+                            "Received StartJoinServerEvent for {entity:?} but it's already connected. Ignoring the event."
+                        );
+                    }
+                    return;
+                }
+            }
+
             entity
         } else {
             let entity = commands.spawn_empty().id();
@@ -71,12 +122,15 @@ pub fn handle_start_join_server_event(
         };
 
         let mut entity_mut = commands.entity(entity);
+
         entity_mut.insert((
             // add the Account to the entity now so plugins can access it earlier
             event.account.to_owned(),
             // localentity is always present for our clients, even if we're not actually logged
             // in
             LocalEntity,
+            // ConnectOpts is inserted as a component here
+            event.connect_opts.clone(),
             // we don't insert InLoginState until we actually create the connection. note that
             // there's no InHandshakeState component since we switch off of the handshake state
             // immediately when the connection is created
@@ -92,11 +146,9 @@ pub fn handle_start_join_server_event(
         }
 
         let task_pool = IoTaskPool::get();
-        let resolved_addr = event.resolved_address;
-        let address = event.address.clone();
-        let proxy = event.proxy.clone();
+        let connect_opts = event.connect_opts.clone();
         let task = task_pool.spawn(async_compat::Compat::new(
-            create_conn_and_send_intention_packet(resolved_addr, address, proxy),
+            create_conn_and_send_intention_packet(connect_opts),
         ));
 
         entity_mut.insert(CreateConnectionTask(task));
@@ -104,20 +156,18 @@ pub fn handle_start_join_server_event(
 }
 
 async fn create_conn_and_send_intention_packet(
-    resolved_addr: SocketAddr,
-    address: ServerAddress,
-    proxy: Option<Proxy>,
+    opts: ConnectOpts,
 ) -> Result<LoginConn, ConnectionError> {
-    let mut conn = if let Some(proxy) = proxy {
-        Connection::new_with_proxy(&resolved_addr, proxy).await?
+    let mut conn = if let Some(proxy) = opts.proxy {
+        Connection::new_with_proxy(&opts.resolved_address, proxy).await?
     } else {
-        Connection::new(&resolved_addr).await?
+        Connection::new(&opts.resolved_address).await?
     };
 
     conn.write(ServerboundIntention {
         protocol_version: PROTOCOL_VERSION,
-        hostname: address.host.clone(),
-        port: address.port,
+        hostname: opts.address.host.clone(),
+        port: opts.address.port,
         intention: ClientIntention::Login,
     })
     .await?;
@@ -140,6 +190,7 @@ pub fn poll_create_connection_task(
         &Account,
         Option<&StartJoinCallback>,
     )>,
+    mut connection_failed_events: EventWriter<ConnectionFailedEvent>,
 ) {
     for (entity, mut task, account, mut start_join_callback) in query.iter_mut() {
         if let Some(poll_res) = future::block_on(future::poll_once(&mut task.0)) {
@@ -147,11 +198,9 @@ pub fn poll_create_connection_task(
             entity_mut.remove::<CreateConnectionTask>();
             let conn = match poll_res {
                 Ok(conn) => conn,
-                Err(err) => {
-                    warn!("failed to create connection: {err}");
-                    if let Some(cb) = start_join_callback.take() {
-                        let _ = cb.0.send(Err(err.into()));
-                    }
+                Err(error) => {
+                    warn!("failed to create connection: {error}");
+                    connection_failed_events.write(ConnectionFailedEvent { entity, error });
                     return;
                 }
             };
@@ -196,3 +245,22 @@ pub fn poll_create_connection_task(
         }
     }
 }
+
+pub fn handle_connection_failed_events(
+    mut events: EventReader<ConnectionFailedEvent>,
+    query: Query<&StartJoinCallback>,
+) {
+    for event in events.read() {
+        let Ok(start_join_callback) = query.get(event.entity) else {
+            // the StartJoinCallback isn't required to be present, so this is fine
+            continue;
+        };
+
+        // io::Error isn't clonable, so we create a new one based on the `kind` and
+        // `to_string`,
+        let ConnectionError::Io(err) = &event.error;
+        let cloned_err = ConnectionError::Io(io::Error::new(err.kind(), err.to_string()));
+
+        let _ = start_join_callback.0.send(Err(cloned_err.into()));
+    }
+}
diff --git a/azalea-client/src/plugins/login.rs b/azalea-client/src/plugins/login.rs
index 9a871ac3..357769e9 100644
--- a/azalea-client/src/plugins/login.rs
+++ b/azalea-client/src/plugins/login.rs
@@ -33,7 +33,7 @@ fn handle_receive_hello_event(trigger: Trigger<ReceiveHelloEvent>, mut commands:
     commands.entity(player).insert(AuthTask(task));
 }
 
-fn poll_auth_task(
+pub fn poll_auth_task(
     mut commands: Commands,
     mut query: Query<(Entity, &mut AuthTask, &mut RawConnection)>,
 ) {
diff --git a/azalea-client/src/plugins/mod.rs b/azalea-client/src/plugins/mod.rs
index 431d59b2..f657b9e9 100644
--- a/azalea-client/src/plugins/mod.rs
+++ b/azalea-client/src/plugins/mod.rs
@@ -1,6 +1,7 @@
 use bevy_app::{PluginGroup, PluginGroupBuilder};
 
 pub mod attack;
+pub mod auto_reconnect;
 pub mod brand;
 pub mod chat;
 pub mod chunks;
@@ -51,7 +52,8 @@ impl PluginGroup for DefaultPlugins {
             .add(pong::PongPlugin)
             .add(connection::ConnectionPlugin)
             .add(login::LoginPlugin)
-            .add(join::JoinPlugin);
+            .add(join::JoinPlugin)
+            .add(auto_reconnect::AutoReconnectPlugin);
         #[cfg(feature = "log")]
         {
             group = group.add(bevy_log::LogPlugin::default());