azalea/swarm/
mod.rs

1//! Swarms are a way to conveniently control many bots.
2//!
3//! See [`Swarm`] for more information.
4
5mod builder;
6mod chat;
7mod events;
8pub mod prelude;
9
10use std::sync::{
11    Arc,
12    atomic::{self, AtomicBool},
13};
14
15use azalea_client::{account::Account, chat::ChatPacket, join::ConnectOpts};
16use azalea_entity::LocalEntity;
17use azalea_protocol::address::ResolvedAddr;
18use azalea_world::Worlds;
19use bevy_app::{AppExit, PluginGroup, PluginGroupBuilder};
20use bevy_ecs::prelude::*;
21pub use builder::SwarmBuilder;
22use futures::future::BoxFuture;
23use parking_lot::RwLock;
24use tokio::{sync::mpsc, task};
25use tracing::{debug, error, warn};
26
27use crate::{Client, JoinOpts, client_impl::StartClientOpts};
28
29/// A swarm is a way to conveniently control many bots at once, while also
30/// being able to control bots at an individual level when desired.
31///
32/// It can safely be cloned, so there should be no need to wrap them in a Mutex.
33///
34/// Swarms are created from [`SwarmBuilder`].
35///
36/// Clients can be added to the swarm later via [`Swarm::add`], and can be
37/// removed with [`Client::disconnect`].
38#[derive(Clone, Resource)]
39pub struct Swarm {
40    /// A way to directly access the ECS.
41    ///
42    /// This will not work if called within a system, as the ECS is already
43    /// locked.
44    #[doc(alias = "ecs_lock")] // former type name
45    pub ecs: Arc<RwLock<World>>,
46
47    // the address is public and mutable so plugins can change it
48    pub address: Arc<RwLock<ResolvedAddr>>,
49
50    pub worlds: Arc<RwLock<Worlds>>,
51
52    /// This is used internally to make the client handler function work.
53    pub(crate) bots_tx: mpsc::UnboundedSender<(Option<crate::Event>, Client)>,
54    /// This is used internally to make the swarm handler function work.
55    pub(crate) swarm_tx: mpsc::UnboundedSender<SwarmEvent>,
56}
57
58/// An event about something that doesn't have to do with a single bot.
59#[derive(Clone, Debug)]
60#[non_exhaustive]
61pub enum SwarmEvent {
62    /// All the bots in the swarm have successfully joined the server.
63    Login,
64    /// The swarm was created.
65    ///
66    /// This is only fired once, and it's guaranteed to be the first event to
67    /// fire.
68    Init,
69    /// A bot got disconnected from the server.
70    ///
71    /// If you'd like to implement special auto-reconnect behavior beyond what's
72    /// built-in, you can disable that with [`SwarmBuilder::reconnect_delay`]
73    /// and then call [`Swarm::add_with_opts`] with the account and options
74    /// from this event.
75    ///
76    /// [`SwarmBuilder::reconnect_delay`]: crate::swarm::SwarmBuilder::reconnect_after
77    Disconnect(Box<Account>, Box<JoinOpts>),
78    /// At least one bot received a chat message.
79    Chat(ChatPacket),
80}
81
82pub type SwarmHandleFn<SS, Fut> = fn(Swarm, SwarmEvent, SS) -> Fut;
83pub type BoxSwarmHandleFn<SS, R> =
84    Box<dyn Fn(Swarm, SwarmEvent, SS) -> BoxFuture<'static, R> + Send + Sync>;
85
86/// Make a bot [`Swarm`].
87///
88/// [`Swarm`]: struct.Swarm.html
89///
90/// # Examples
91/// ```rust,no_run
92/// use azalea::{prelude::*, swarm::prelude::*};
93/// use std::time::Duration;
94///
95/// #[derive(Clone, Component, Default)]
96/// struct State {}
97///
98/// #[derive(Clone, Default, Resource)]
99/// struct SwarmState {}
100///
101/// #[tokio::main]
102/// async fn main() -> AppExit {
103///     let mut accounts = Vec::new();
104///     let mut states = Vec::new();
105///
106///     for i in 0..10 {
107///         accounts.push(Account::offline(&format!("bot{i}")));
108///         states.push(State::default());
109///     }
110///
111///     SwarmBuilder::new()
112///         .add_accounts(accounts.clone())
113///         .set_handler(handle)
114///         .set_swarm_handler(swarm_handle)
115///         .join_delay(Duration::from_millis(1000))
116///         .start("localhost")
117///         .await
118/// }
119///
120/// async fn handle(bot: Client, event: Event, _state: State) -> anyhow::Result<()> {
121///     match &event {
122///         _ => {}
123///     }
124///     Ok(())
125/// }
126///
127/// async fn swarm_handle(
128///     mut swarm: Swarm,
129///     event: SwarmEvent,
130///     _state: SwarmState,
131/// ) -> anyhow::Result<()> {
132///     match &event {
133///         SwarmEvent::Chat(m) => {
134///             println!("{}", m.message().to_ansi());
135///         }
136///         _ => {}
137///     }
138///     Ok(())
139/// }
140impl Swarm {
141    /// Add a new account to the swarm.
142    ///
143    /// You can remove it later by calling [`Client::disconnect`].
144    ///
145    /// # Errors
146    ///
147    /// Returns an error if the server's address could not be resolved.
148    pub async fn add<S: Component + Clone>(&self, account: &Account, state: S) -> Client {
149        self.add_with_opts(account, state, &JoinOpts::default())
150            .await
151    }
152    /// Add a new account to the swarm, using custom options.
153    ///
154    /// This is useful if you want bots in the same swarm to connect to
155    /// different addresses. Usually you'll just want [`Self::add`] though.
156    ///
157    /// # Errors
158    ///
159    /// Returns an error if the server's address could not be resolved.
160    pub async fn add_with_opts<S: Component + Clone>(
161        &self,
162        account: &Account,
163        state: S,
164        join_opts: &JoinOpts,
165    ) -> Client {
166        debug!(
167            "add_with_opts called for account {} with opts {join_opts:?}",
168            account.username()
169        );
170
171        let mut address = self.address.read().clone();
172        if let Some(custom_server_addr) = join_opts.custom_server_addr.clone() {
173            address.server = custom_server_addr;
174        }
175        if let Some(custom_socket_addr) = join_opts.custom_socket_addr {
176            address.socket = custom_socket_addr;
177        }
178        let server_proxy = join_opts.server_proxy.clone();
179        let sessionserver_proxy = join_opts.sessionserver_proxy.clone();
180
181        let (tx, rx) = mpsc::unbounded_channel();
182
183        let client = Client::start_client(StartClientOpts {
184            ecs_lock: self.ecs.clone(),
185            account: account.clone(),
186            connect_opts: ConnectOpts {
187                address,
188                server_proxy,
189                sessionserver_proxy,
190            },
191            event_sender: Some(tx),
192        })
193        .await;
194        // add the state to the client
195        {
196            let mut ecs = self.ecs.write();
197            ecs.entity_mut(client.entity).insert(state);
198        }
199
200        let cloned_bot = client.clone();
201        let swarm_tx = self.swarm_tx.clone();
202        let bots_tx = self.bots_tx.clone();
203
204        let join_opts = join_opts.clone();
205        task::spawn_local(Self::event_copying_task(
206            rx, swarm_tx, bots_tx, cloned_bot, join_opts,
207        ));
208
209        client
210    }
211
212    /// Copy the events from a client's receiver into bots_tx, until the bot is
213    /// removed from the ECS.
214    async fn event_copying_task(
215        mut rx: mpsc::UnboundedReceiver<crate::Event>,
216        swarm_tx: mpsc::UnboundedSender<SwarmEvent>,
217        bots_tx: mpsc::UnboundedSender<(Option<crate::Event>, Client)>,
218        bot: Client,
219        join_opts: JoinOpts,
220    ) {
221        while let Some(event) = rx.recv().await {
222            if rx.len() > 1_000 {
223                static WARNED_1_000: AtomicBool = AtomicBool::new(false);
224                if !WARNED_1_000.swap(true, atomic::Ordering::Relaxed) {
225                    warn!(
226                        "The client's Event channel has more than 1,000 items! If you don't need it, consider disabling the `packet-event` feature for `azalea`."
227                    )
228                }
229
230                if rx.len() > 10_000 {
231                    static WARNED_10_000: AtomicBool = AtomicBool::new(false);
232                    if !WARNED_10_000.swap(true, atomic::Ordering::Relaxed) {
233                        warn!("The client's Event channel has more than 10,000 items!!")
234                    }
235
236                    if rx.len() > 100_000 {
237                        static WARNED_100_000: AtomicBool = AtomicBool::new(false);
238                        if !WARNED_100_000.swap(true, atomic::Ordering::Relaxed) {
239                            warn!("The client's Event channel has more than 100,000 items!!!")
240                        }
241
242                        if rx.len() > 1_000_000 {
243                            static WARNED_1_000_000: AtomicBool = AtomicBool::new(false);
244                            if !WARNED_1_000_000.swap(true, atomic::Ordering::Relaxed) {
245                                warn!(
246                                    "The client's Event channel has more than 1,000,000 items!!!! your code is almost certainly leaking memory"
247                                )
248                            }
249                        }
250                    }
251                }
252            }
253
254            if let crate::Event::Disconnect(_) = event {
255                debug!(
256                    "Sending SwarmEvent::Disconnect due to receiving an Event::Disconnect from client {}",
257                    bot.entity
258                );
259                let account = bot.account();
260                swarm_tx
261                    .send(SwarmEvent::Disconnect(
262                        Box::new(account),
263                        Box::new(join_opts.clone()),
264                    ))
265                    .unwrap();
266            }
267
268            // we can't handle events here (since we can't copy the handler),
269            // they're handled above in SwarmBuilder::start
270            if let Err(e) = bots_tx.send((Some(event), bot.clone())) {
271                error!(
272                    "Error sending event to swarm, aborting event_copying_task for {}: {e}",
273                    bot.entity
274                );
275                break;
276            }
277        }
278        debug!(
279            "client sender ended for {}, this won't trigger SwarmEvent::Disconnect unless the client already sent its own disconnect event",
280            bot.entity
281        );
282    }
283
284    /// Get an array of ECS [`Entity`]s for all [`LocalEntity`]s in our world.
285    /// This will include clients that were disconnected without being removed
286    /// from the ECS.
287    ///
288    /// [`LocalEntity`]: azalea_entity::LocalEntity
289    pub fn client_entities(&self) -> Box<[Entity]> {
290        let mut ecs = self.ecs.write();
291        let mut query = ecs.query_filtered::<Entity, With<LocalEntity>>();
292        query.iter(&ecs).collect::<Box<[Entity]>>()
293    }
294
295    /// End the entire swarm and return from [`SwarmBuilder::start`].
296    ///
297    /// You should typically avoid calling this if you intend on creating the
298    /// swarm again, because creating an entirely new swarm can be a
299    /// relatively expensive process.
300    ///
301    /// If you only want to change the server that the bots are connecting to,
302    /// it may be better to call [`Swarm::add_with_opts`] with a different
303    /// server address.
304    ///
305    /// This is also implemented on [`Client`] as [`Client::exit`].
306    pub fn exit(&self) {
307        self.ecs.write().write_message(AppExit::Success);
308    }
309}
310
311impl IntoIterator for Swarm {
312    type Item = Client;
313    type IntoIter = std::vec::IntoIter<Self::Item>;
314
315    /// Iterate over the bots in this swarm.
316    ///
317    /// ```rust,no_run
318    /// # use azalea::{prelude::*, swarm::prelude::*};
319    /// #[derive(Clone, Component)]
320    /// # pub struct State;
321    /// # fn example(swarm: Swarm) {
322    /// for bot in swarm {
323    ///     let state = bot.component::<State>();
324    ///     // ...
325    /// }
326    /// # }
327    /// ```
328    fn into_iter(self) -> Self::IntoIter {
329        let client_entities = self.client_entities();
330
331        client_entities
332            .into_iter()
333            .map(|entity| Client::new(entity, self.ecs.clone()))
334            .collect::<Box<[Client]>>()
335            .into_iter()
336    }
337}
338
339/// This plugin group will add all the default plugins necessary for swarms to
340/// work.
341pub struct DefaultSwarmPlugins;
342
343impl PluginGroup for DefaultSwarmPlugins {
344    fn build(self) -> PluginGroupBuilder {
345        PluginGroupBuilder::start::<Self>()
346            .add(chat::SwarmChatPlugin)
347            .add(events::SwarmPlugin)
348    }
349}
350
351/// A marker that can be used in place of a SwarmState in [`SwarmBuilder`].
352///
353/// You probably don't need to use this manually since the compiler will infer
354/// it for you.
355#[derive(Clone, Default, Resource)]
356pub struct NoSwarmState;
azalea/swarm/mod.rs

azalea/swarm/
mod.rs
