diff --git a/src/main.rs b/src/main.rs index 00775ba..0dcdab4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,8 @@ -use std::path::Path; +use std::{future::Future, path::Path}; +use ::serenity::all::{ChannelId, GuildId, UserId}; use eyre::{Context as _, Error, OptionExt, Result}; -use futures::FutureExt; +use futures::{FutureExt, StreamExt, TryStreamExt}; use poise::serenity_prelude as serenity; use rand::seq::SliceRandom; use rusqlite::OptionalExtension; @@ -90,6 +91,19 @@ async fn channel_children( .collect()) } +fn move_users<'a>( + ctx: Context<'a>, + guild: GuildId, + users: impl Iterator + Send + 'a, +) -> impl Future> + Send + 'a { + futures::stream::iter(users.map(move |(channel, user)| async move { + guild.move_member(ctx, user, channel).await?; + Ok(()) + })) + .buffer_unordered(10) + .try_collect::<()>() +} + async fn spread( ctx: Context<'_>, from: serenity::ChannelId, @@ -107,17 +121,18 @@ async fn spread( .guild() .ok_or_eyre("This bot only works in servers")?; let mut to = channel_children(&ctx, to).await?; - let mut users = from.members(ctx)?; users.shuffle(&mut rand::thread_rng()); to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty())); to.shuffle(&mut rand::thread_rng()); - futures::future::try_join_all( + move_users( + ctx, + guild, to.into_iter() .zip(&users) - .map(|(channel, user)| guild.move_member(ctx, user.user.id, channel)), + .map(|(channel, user)| (channel.id, user.user.id)), ) .await?; @@ -143,6 +158,7 @@ async fn collect( .flatten() .collect::>(); + move_users(ctx, guild, users.iter().map(|user| (to, user.user.id))).await?; futures::future::try_join_all(users.iter().map(|user| guild.move_member(&ctx, user, to))) .await?; @@ -227,7 +243,10 @@ async fn st(ctx: Context<'_>, mut spectators: Vec) -> Result<( |member| { let spectators = &spectators; async move { - match (spectators.contains(&member.user.id), member.roles.contains(&st)) { + match ( + spectators.contains(&member.user.id), + member.roles.contains(&st), + ) { (true, true) | (false, false) => Ok(()), (true, false) => member.add_role(&ctx, st).await, (false, true) => member.remove_role(&ctx, st).await,