Buffer user moves - limit to 10 at a time

Hopefully mitigates the issues with hangs on larger calls
This commit is contained in:
bluepython508
2025-01-30 15:07:31 +00:00
parent 994babaacf
commit 0e4df70f39

View File

@@ -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 eyre::{Context as _, Error, OptionExt, Result};
use futures::FutureExt; use futures::{FutureExt, StreamExt, TryStreamExt};
use poise::serenity_prelude as serenity; use poise::serenity_prelude as serenity;
use rand::seq::SliceRandom; use rand::seq::SliceRandom;
use rusqlite::OptionalExtension; use rusqlite::OptionalExtension;
@@ -90,6 +91,19 @@ async fn channel_children(
.collect()) .collect())
} }
fn move_users<'a>(
ctx: Context<'a>,
guild: GuildId,
users: impl Iterator<Item = (ChannelId, UserId)> + Send + 'a,
) -> impl Future<Output = Result<()>> + 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( async fn spread(
ctx: Context<'_>, ctx: Context<'_>,
from: serenity::ChannelId, from: serenity::ChannelId,
@@ -107,17 +121,18 @@ async fn spread(
.guild() .guild()
.ok_or_eyre("This bot only works in servers")?; .ok_or_eyre("This bot only works in servers")?;
let mut to = channel_children(&ctx, to).await?; let mut to = channel_children(&ctx, to).await?;
let mut users = from.members(ctx)?; let mut users = from.members(ctx)?;
users.shuffle(&mut rand::thread_rng()); users.shuffle(&mut rand::thread_rng());
to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty())); to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty()));
to.shuffle(&mut rand::thread_rng()); to.shuffle(&mut rand::thread_rng());
futures::future::try_join_all( move_users(
ctx,
guild,
to.into_iter() to.into_iter()
.zip(&users) .zip(&users)
.map(|(channel, user)| guild.move_member(ctx, user.user.id, channel)), .map(|(channel, user)| (channel.id, user.user.id)),
) )
.await?; .await?;
@@ -143,6 +158,7 @@ async fn collect(
.flatten() .flatten()
.collect::<Vec<_>>(); .collect::<Vec<_>>();
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))) futures::future::try_join_all(users.iter().map(|user| guild.move_member(&ctx, user, to)))
.await?; .await?;
@@ -227,7 +243,10 @@ async fn st(ctx: Context<'_>, mut spectators: Vec<serenity::UserId>) -> Result<(
|member| { |member| {
let spectators = &spectators; let spectators = &spectators;
async move { 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, true) | (false, false) => Ok(()),
(true, false) => member.add_role(&ctx, st).await, (true, false) => member.add_role(&ctx, st).await,
(false, true) => member.remove_role(&ctx, st).await, (false, true) => member.remove_role(&ctx, st).await,