azalea/swarm/builder.rs
1use std::{
2 collections::{HashMap, hash_map},
3 mem,
4 sync::{
5 Arc,
6 atomic::{self, AtomicBool},
7 },
8 time::Duration,
9};
10
11use azalea_client::{DefaultPlugins, account::Account, start_ecs_runner};
12use azalea_protocol::address::{ResolvableAddr, ResolvedAddr};
13use azalea_world::Worlds;
14use bevy_app::{App, AppExit, Plugins, SubApp};
15use bevy_ecs::{component::Component, resource::Resource};
16use futures::future::join_all;
17use parking_lot::RwLock;
18use tokio::{sync::mpsc, task};
19use tracing::{debug, error, warn};
20
21use crate::{
22 BoxHandleFn, HandleFn, JoinOpts, NoState, TokioRuntimeHandle,
23 auto_reconnect::{AutoReconnectDelay, DEFAULT_RECONNECT_DELAY},
24 bot::DefaultBotPlugins,
25 swarm::{
26 BoxSwarmHandleFn, DefaultSwarmPlugins, NoSwarmState, Swarm, SwarmEvent, SwarmHandleFn,
27 },
28};
29
30/// Create a new [`Swarm`].
31///
32/// The generics of this struct stand for the following:
33/// - S: State
34/// - SS: Swarm State
35/// - R: Return type of the handler
36/// - SR: Return type of the swarm handler
37///
38/// You shouldn't have to manually set them though, they'll be inferred for you.
39pub struct SwarmBuilder<S, SS, R, SR>
40where
41 S: Send + Sync + Clone + Component + 'static,
42 SS: Default + Send + Sync + Clone + Resource + 'static,
43 Self: Send,
44{
45 // SubApp is used instead of App to make it Send
46 pub(crate) app: SubApp,
47 /// The accounts and proxies that are going to join the server.
48 pub(crate) accounts: Vec<(Account, JoinOpts)>,
49 /// The individual bot states.
50 ///
51 /// This must be the same length as `accounts`, since each bot gets one
52 /// state.
53 pub(crate) states: Vec<S>,
54 /// The state for the overall swarm.
55 pub(crate) swarm_state: SS,
56 /// The function that's called every time a bot receives an [`Event`].
57 pub(crate) handler: Option<BoxHandleFn<S, R>>,
58 /// The function that's called every time the swarm receives a
59 /// [`SwarmEvent`].
60 pub(crate) swarm_handler: Option<BoxSwarmHandleFn<SS, SR>>,
61
62 /// How long we should wait between each bot joining the server.
63 ///
64 /// If this is None, every bot will connect at the same time. None is
65 /// different than a duration of 0, since if a duration is present the
66 /// bots will wait for the previous one to be ready.
67 pub(crate) join_delay: Option<Duration>,
68
69 /// The default reconnection delay for our bots.
70 ///
71 /// This will change the value of the [`AutoReconnectDelay`] resource.
72 pub(crate) reconnect_after: Option<Duration>,
73}
74impl SwarmBuilder<NoState, NoSwarmState, (), ()> {
75 /// Start creating the swarm.
76 #[must_use]
77 pub fn new() -> Self {
78 Self::new_without_plugins().add_plugins((
79 DefaultPlugins,
80 DefaultBotPlugins,
81 DefaultSwarmPlugins,
82 ))
83 }
84
85 /// [`Self::new`] but without adding the plugins by default.
86 ///
87 /// This is useful if you want to disable a default plugin. This also exists
88 /// for `ClientBuilder`, see [`ClientBuilder::new_without_plugins`].
89 ///
90 /// You **must** add [`DefaultPlugins`], [`DefaultBotPlugins`], and
91 /// [`DefaultSwarmPlugins`] to this.
92 ///
93 /// ```
94 /// # use azalea::{prelude::*, swarm::{prelude::*, DefaultSwarmPlugins}, bot::DefaultBotPlugins};
95 /// use azalea::app::PluginGroup;
96 ///
97 /// let swarm_builder = SwarmBuilder::new_without_plugins()
98 /// .add_plugins((
99 /// DefaultBotPlugins,
100 /// DefaultSwarmPlugins,
101 /// azalea::DefaultPlugins
102 /// .build()
103 /// .disable::<azalea::chat_signing::ChatSigningPlugin>(),
104 /// ));
105 /// # swarm_builder.set_handler(handle).set_swarm_handler(swarm_handle);
106 /// # #[derive(Clone, Component, Default)]
107 /// # pub struct State;
108 /// # #[derive(Clone, Resource, Default)]
109 /// # pub struct SwarmState;
110 /// # async fn handle(mut bot: Client, event: Event, state: State) -> eyre::Result<()> {
111 /// # Ok(())
112 /// # }
113 /// # async fn swarm_handle(swarm: Swarm, event: SwarmEvent, state: SwarmState) -> eyre::Result<()> {
114 /// # Ok(())
115 /// # }
116 /// ```
117 ///
118 /// [`ClientBuilder::new_without_plugins`]: crate::ClientBuilder::new_without_plugins
119 #[must_use]
120 pub fn new_without_plugins() -> Self {
121 SwarmBuilder {
122 // we create the app here so plugins can add onto it.
123 // the schedules won't run until [`Self::start`] is called.
124
125 // `App::new()` is used instead of `SubApp::new()` so the necessary resources are
126 // initialized
127 app: mem::take(App::new().main_mut()),
128 accounts: Vec::new(),
129 states: Vec::new(),
130 swarm_state: NoSwarmState,
131 handler: None,
132 swarm_handler: None,
133 join_delay: None,
134 reconnect_after: Some(DEFAULT_RECONNECT_DELAY),
135 }
136 }
137}
138
139impl<SS, SR> SwarmBuilder<NoState, SS, (), SR>
140where
141 SS: Default + Send + Sync + Clone + Resource + 'static,
142{
143 /// Set the function that's called every time a bot receives an
144 /// [`Event`](crate::Event). This is the intended way to handle
145 /// normal per-bot events.
146 ///
147 /// Currently you can have up to one handler.
148 ///
149 /// Note that if you're creating clients directly from the ECS using
150 /// [`StartJoinServerEvent`] and the client wasn't already in the ECS, then
151 /// the handler function won't be called for that client. This also applies
152 /// to [`SwarmBuilder::set_swarm_handler`]. This shouldn't be a concern for
153 /// most bots, though.
154 ///
155 /// ```
156 /// # use azalea::{prelude::*, swarm::prelude::*};
157 /// # let swarm_builder = SwarmBuilder::new().set_swarm_handler(swarm_handle);
158 /// swarm_builder.set_handler(handle);
159 ///
160 /// #[derive(Clone, Component, Default)]
161 /// struct State {}
162 /// async fn handle(mut bot: Client, event: Event, state: State) -> eyre::Result<()> {
163 /// Ok(())
164 /// }
165 ///
166 /// # #[derive(Clone, Default, Resource)]
167 /// # struct SwarmState {}
168 /// # async fn swarm_handle(
169 /// # mut swarm: Swarm,
170 /// # event: SwarmEvent,
171 /// # state: SwarmState,
172 /// # ) -> eyre::Result<()> {
173 /// # Ok(())
174 /// # }
175 /// ```
176 ///
177 /// [`StartJoinServerEvent`]: azalea_client::join::StartJoinServerEvent
178 #[must_use]
179 pub fn set_handler<S, Fut, R>(self, handler: HandleFn<S, Fut>) -> SwarmBuilder<S, SS, R, SR>
180 where
181 Fut: Future<Output = R> + Send + 'static,
182 S: Send + Sync + Clone + Component + Default + 'static,
183 {
184 SwarmBuilder {
185 handler: Some(Box::new(move |bot, event, state: S| {
186 Box::pin(handler(bot, event, state))
187 })),
188 // if we added accounts before the State was set, we've gotta set it to the default now
189 states: vec![S::default(); self.accounts.len()],
190 app: self.app,
191 ..self
192 }
193 }
194}
195
196impl<S, R> SwarmBuilder<S, NoSwarmState, R, ()>
197where
198 S: Send + Sync + Clone + Component + 'static,
199{
200 /// Set the function that's called every time the swarm receives a
201 /// [`SwarmEvent`]. This is the intended way to handle global swarm events.
202 ///
203 /// Currently you can have up to one swarm handler.
204 ///
205 /// Note that if you're creating clients directly from the ECS using
206 /// [`StartJoinServerEvent`] and the client wasn't already in the ECS, then
207 /// this handler function won't be called for that client. This also applies
208 /// to [`SwarmBuilder::set_handler`]. This shouldn't be a concern for
209 /// most bots, though.
210 ///
211 /// ```
212 /// # use azalea::{prelude::*, swarm::prelude::*};
213 /// # let swarm_builder = SwarmBuilder::new().set_handler(handle);
214 /// swarm_builder.set_swarm_handler(swarm_handle);
215 ///
216 /// # #[derive(Clone, Component, Default)]
217 /// # struct State {}
218 ///
219 /// # async fn handle(mut bot: Client, event: Event, state: State) -> eyre::Result<()> {
220 /// # Ok(())
221 /// # }
222 ///
223 /// #[derive(Clone, Default, Resource)]
224 /// struct SwarmState {}
225 /// async fn swarm_handle(
226 /// mut swarm: Swarm,
227 /// event: SwarmEvent,
228 /// state: SwarmState,
229 /// ) -> eyre::Result<()> {
230 /// Ok(())
231 /// }
232 /// ```
233 ///
234 /// [`StartJoinServerEvent`]: azalea_client::join::StartJoinServerEvent
235 #[must_use]
236 pub fn set_swarm_handler<SS, Fut, SR>(
237 self,
238 handler: SwarmHandleFn<SS, Fut>,
239 ) -> SwarmBuilder<S, SS, R, SR>
240 where
241 SS: Default + Send + Sync + Clone + Resource + 'static,
242 Fut: Future<Output = SR> + Send + 'static,
243 {
244 SwarmBuilder {
245 handler: self.handler,
246 app: self.app,
247 accounts: self.accounts,
248 states: self.states,
249 swarm_state: SS::default(),
250 swarm_handler: Some(Box::new(move |swarm, event, state| {
251 Box::pin(handler(swarm, event, state))
252 })),
253 join_delay: self.join_delay,
254 reconnect_after: self.reconnect_after,
255 }
256 }
257}
258
259impl<S, SS, R, SR> SwarmBuilder<S, SS, R, SR>
260where
261 S: Send + Sync + Clone + Component + 'static,
262 SS: Default + Send + Sync + Clone + Resource + 'static,
263 R: Send + 'static,
264 SR: Send + 'static,
265{
266 /// Add a vec of [`Account`]s to the swarm.
267 ///
268 /// Use [`Self::add_account`] to only add one account. If you want the
269 /// clients to have different default states, add them one at a time with
270 /// [`Self::add_account_with_state`].
271 ///
272 /// By default, every account will join at the same time, you can add a
273 /// delay with [`Self::join_delay`].
274 #[must_use]
275 pub fn add_accounts(mut self, accounts: Vec<Account>) -> Self
276 where
277 S: Default,
278 {
279 for account in accounts {
280 self = self.add_account(account);
281 }
282
283 self
284 }
285
286 /// Add a single new [`Account`] to the swarm.
287 ///
288 /// Use [`Self::add_accounts`] to add multiple accounts at a time.
289 ///
290 /// This will make the state for this client be the default, use
291 /// [`Self::add_account_with_state`] to avoid that.
292 #[must_use]
293 pub fn add_account(self, account: Account) -> Self
294 where
295 S: Default,
296 {
297 self.add_account_with_state_and_opts(account, S::default(), JoinOpts::default())
298 }
299
300 /// Add an account with a custom initial state.
301 ///
302 /// Use just [`Self::add_account`] to use the `Default` implementation for
303 /// the state.
304 #[must_use]
305 pub fn add_account_with_state(self, account: Account, state: S) -> Self {
306 self.add_account_with_state_and_opts(account, state, JoinOpts::default())
307 }
308
309 /// Add an account with a custom initial state.
310 ///
311 /// Use just [`Self::add_account`] to use the `Default` implementation for
312 /// the state.
313 #[must_use]
314 pub fn add_account_with_opts(self, account: Account, opts: JoinOpts) -> Self
315 where
316 S: Default,
317 {
318 self.add_account_with_state_and_opts(account, S::default(), opts)
319 }
320
321 /// Same as [`Self::add_account_with_state`], but allow passing in custom
322 /// join options.
323 #[must_use]
324 pub fn add_account_with_state_and_opts(
325 mut self,
326 account: Account,
327 state: S,
328 join_opts: JoinOpts,
329 ) -> Self {
330 self.accounts.push((account, join_opts));
331 self.states.push(state);
332 self
333 }
334
335 /// Set the swarm state instead of initializing defaults.
336 #[must_use]
337 pub fn set_swarm_state(mut self, swarm_state: SS) -> Self {
338 self.swarm_state = swarm_state;
339 self
340 }
341
342 /// Add one or more plugins to this swarm.
343 ///
344 /// See [`Self::new_without_plugins`] to learn how to disable default
345 /// plugins.
346 #[must_use]
347 pub fn add_plugins<M>(mut self, plugins: impl Plugins<M>) -> Self {
348 self.app.add_plugins(plugins);
349 self
350 }
351
352 /// Set how long we should wait between each bot joining the server.
353 ///
354 /// By default, every bot will connect at the same time. If you set this
355 /// field, however, the bots will wait for the previous one to have
356 /// connected and *then* they'll wait the given duration.
357 #[must_use]
358 pub fn join_delay(mut self, delay: Duration) -> Self {
359 self.join_delay = Some(delay);
360 self
361 }
362
363 /// Configures the auto-reconnection behavior for our bots.
364 ///
365 /// If this is `Some`, then it'll set the default reconnection delay for our
366 /// bots (how long they'll wait after being kicked before they try
367 /// rejoining). if it's `None`, then auto-reconnecting will be disabled.
368 ///
369 /// If this function isn't called, then our clients will reconnect after
370 /// [`DEFAULT_RECONNECT_DELAY`].
371 #[must_use]
372 pub fn reconnect_after(mut self, delay: impl Into<Option<Duration>>) -> Self {
373 self.reconnect_after = delay.into();
374 self
375 }
376
377 /// Build this `SwarmBuilder` into an actual [`Swarm`] and join the given
378 /// server.
379 ///
380 /// The `address` argument can be a `&str`, [`ServerAddr`],
381 /// [`ResolvedAddr`], or anything else that implements [`ResolvableAddr`].
382 ///
383 /// [`ServerAddr`]: ../../azalea_protocol/address/struct.ServerAddr.html
384 /// [`ResolvedAddr`]: ../../azalea_protocol/address/struct.ResolvedAddr.html
385 /// [`ResolvableAddr`]: ../../azalea_protocol/address/trait.ResolvableAddr.html
386 pub async fn start(self, address: impl ResolvableAddr) -> AppExit {
387 self.start_with_opts(address, JoinOpts::default()).await
388 }
389
390 #[doc(hidden)]
391 #[deprecated = "renamed to `start_with_opts`."]
392 pub async fn start_with_default_opts(
393 self,
394 address: impl ResolvableAddr,
395 default_join_opts: JoinOpts,
396 ) -> AppExit {
397 self.start_with_opts(address, default_join_opts).await
398 }
399
400 /// Do the same as [`Self::start`], but allow passing in default join
401 /// options for the bots.
402 pub async fn start_with_opts(
403 mut self,
404 address: impl ResolvableAddr,
405 join_opts: JoinOpts,
406 ) -> AppExit {
407 assert_eq!(
408 self.accounts.len(),
409 self.states.len(),
410 "There must be exactly one state per bot."
411 );
412
413 debug!("Starting Azalea {}", env!("CARGO_PKG_VERSION"));
414
415 let address = if let Some(socket_addr) = join_opts.custom_socket_addr {
416 let server_addr = if let Some(server_addr) = join_opts
417 .custom_server_addr
418 .clone()
419 .or_else(|| address.clone().server_addr().ok())
420 {
421 server_addr
422 } else {
423 error!(
424 "Failed to parse address: {address:?}. If this was expected, consider passing in a `ServerAddr` instead."
425 );
426 return AppExit::error();
427 };
428
429 ResolvedAddr {
430 server: server_addr,
431 socket: socket_addr,
432 }
433 } else {
434 let Ok(addr) = address.clone().resolve().await else {
435 error!(
436 "Failed to resolve address: {address:?}. If this was expected, consider resolving the address earlier with `ResolvableAddr::resolve`."
437 );
438 return AppExit::error();
439 };
440 addr
441 };
442
443 let worlds = Arc::new(RwLock::new(Worlds::default()));
444
445 // we can't modify the swarm plugins after this
446 let (bots_tx, mut bots_rx) = mpsc::unbounded_channel();
447 let (swarm_tx, mut swarm_rx) = mpsc::unbounded_channel();
448
449 swarm_tx.send(SwarmEvent::Init).unwrap();
450
451 let main_schedule_label = self.app.update_schedule.unwrap();
452
453 let local_set = task::LocalSet::new();
454
455 local_set.run_until(async move {
456 // start_ecs_runner must be run inside of the LocalSet
457 let (ecs_lock, start_running_systems, appexit_rx) = start_ecs_runner(&mut self.app);
458
459 let swarm = Swarm {
460 ecs: ecs_lock.clone(),
461
462 address: Arc::new(RwLock::new(address)),
463 worlds,
464
465 bots_tx,
466
467 swarm_tx: swarm_tx.clone(),
468 };
469
470 // run the main schedule so the startup systems run
471 {
472 let mut ecs = ecs_lock.write();
473 ecs.insert_resource(TokioRuntimeHandle(tokio::runtime::Handle::current()));
474 ecs.insert_resource(swarm.clone());
475 ecs.insert_resource(self.swarm_state.clone());
476 if let Some(reconnect_after) = self.reconnect_after {
477 ecs.insert_resource(AutoReconnectDelay {
478 delay: reconnect_after,
479 });
480 } else {
481 ecs.remove_resource::<AutoReconnectDelay>();
482 }
483 ecs.run_schedule(main_schedule_label);
484 ecs.clear_trackers();
485 }
486
487 // only do this after we inserted the Swarm and state resources to avoid errors
488 // where Res<Swarm> is inaccessible
489 start_running_systems();
490
491 // SwarmBuilder (self) isn't Send so we have to take all the things we need out
492 // of it
493 let swarm_clone = swarm.clone();
494 let join_delay = self.join_delay;
495 let accounts = self.accounts.clone();
496 let states = self.states.clone();
497
498 task::spawn_local(async move {
499 if let Some(join_delay) = join_delay {
500 // if there's a join delay, then join one by one
501 for ((account, bot_join_opts), state) in accounts.iter().zip(states) {
502 let mut join_opts = join_opts.clone();
503 join_opts.update(bot_join_opts);
504 let _ = swarm_clone.add_with_opts(account, state, &join_opts).await;
505 tokio::time::sleep(join_delay).await;
506 }
507 } else {
508 // otherwise, join all at once
509 let swarm_borrow = &swarm_clone;
510 join_all(accounts.iter().zip(states).map(
511 |((account, bot_join_opts), state)| async {
512 let mut join_opts = join_opts.clone();
513 join_opts.update(bot_join_opts);
514 let _ = swarm_borrow
515 .clone()
516 .add_with_opts(account, state, &join_opts)
517 .await;
518 },
519 ))
520 .await;
521 }
522
523 swarm_tx.send(SwarmEvent::Login).unwrap();
524 });
525
526 let swarm_state = self.swarm_state;
527
528 // Watch swarm_rx and send those events to the swarm_handle.
529 let swarm_clone = swarm.clone();
530 let swarm_handler_task = task::spawn_local(async move {
531 while let Some(event) = swarm_rx.recv().await {
532 if let Some(swarm_handler) = &self.swarm_handler {
533 task::spawn_local((swarm_handler)(
534 swarm_clone.clone(),
535 event,
536 swarm_state.clone(),
537 ));
538 }
539 }
540
541 unreachable!(
542 "The `Swarm` here contains a sender for the `SwarmEvent`s, so swarm_rx.recv() will never fail"
543 );
544 });
545
546 // bot events
547 let client_handler_task = task::spawn_local(async move {
548 while let Some((Some(first_event), first_bot)) = bots_rx.recv().await {
549 if bots_rx.len() > 1_000 {
550 static WARNED: AtomicBool = AtomicBool::new(false);
551 if !WARNED.swap(true, atomic::Ordering::Relaxed) {
552 warn!("the Client Event channel has more than 1000 items!")
553 }
554 }
555
556 if let Some(handler) = &self.handler {
557 let ecs_mutex = first_bot.ecs.clone();
558 let mut ecs = ecs_mutex.write();
559 let mut query = ecs.query::<Option<&S>>();
560 let Ok(Some(first_bot_state)) = query.get(&ecs, first_bot.entity) else {
561 error!(
562 "the first bot ({} / {}) is missing the required state component! none of the client handler functions will be called.",
563 first_bot.username(),
564 first_bot.entity
565 );
566 continue;
567 };
568 let first_bot_entity = first_bot.entity;
569 let first_bot_state = first_bot_state.clone();
570
571 task::spawn_local((handler)(first_bot, first_event, first_bot_state.clone()));
572
573 // this makes it not have to keep locking the ecs
574 let mut states = HashMap::new();
575 states.insert(first_bot_entity, first_bot_state);
576 while let Ok((Some(event), bot)) = bots_rx.try_recv() {
577 let state = match states.entry(bot.entity) {
578 hash_map::Entry::Occupied(e) => e.into_mut(),
579 hash_map::Entry::Vacant(e) => {
580 let Ok(Some(state)) = query.get(&ecs, bot.entity) else {
581 error!(
582 "one of our bots ({} / {}) is missing the required state component! its client handler function will not be called.",
583 bot.username(),
584 bot.entity
585 );
586 continue;
587 };
588 let state = state.clone();
589 e.insert(state)
590 }
591 };
592 task::spawn_local((handler)(bot, event, state.clone()));
593 }
594 }
595 }
596 });
597
598 let app_exit = appexit_rx
599 .await
600 .expect("appexit_tx shouldn't be dropped by the ECS runner before sending");
601
602 swarm_handler_task.abort();
603 client_handler_task.abort();
604
605 app_exit
606 }).await
607 }
608}
609
610impl Default for SwarmBuilder<NoState, NoSwarmState, (), ()> {
611 fn default() -> Self {
612 Self::new()
613 }
614}