Core: simplify + asyncify soft_bans

youmubot-core/src/admin/soft_ban.rs

@@ -1,13 +1,15 @@
 use crate::db::{ServerSoftBans, SoftBans};
 use chrono::offset::Utc;
+use futures_util::{stream, TryStreamExt};
 use serenity::{
-    framework::standard::{macros::command, Args, CommandError as Error, CommandResult},
+    framework::standard::{macros::command, Args, CommandResult},
     model::{
         channel::Message,
-        id::{RoleId, UserId},
+        id::{GuildId, RoleId, UserId},
     },
+    CacheAndHttp,
 };
-use std::cmp::max;
+use std::sync::Arc;
 use youmubot_prelude::*;
 
 #[command]
@@ -18,57 +20,56 @@ use youmubot_prelude::*;
 #[min_args(1)]
 #[max_args(2)]
 #[only_in("guilds")]
-pub fn soft_ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
+pub async fn soft_ban(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
-    let user = args.single::<UserId>()?.to_user(&ctx)?;
+    let user = args.single::<UserId>()?.to_user(&ctx).await?;
+    let data = ctx.data.read().await;
     let duration = if args.is_empty() {
         None
     } else {
-        Some(
+        Some(args.single::<args::Duration>()?)
-            args.single::<args::Duration>()
-                .map_err(|e| Error::from(&format!("{:?}", e)))?,
-        )
     };
-    let guild = msg.guild_id.ok_or(Error::from("Command is guild only"))?;
+    let guild = msg.guild_id.ok_or(Error::msg("Command is guild only"))?;
 
-    let db = SoftBans::open(&*ctx.data.read());
+    let db = SoftBans::open(&*data);
-    let mut db = db.borrow_mut()?;
+    let val = db
-    let mut server_ban = db.get_mut(&guild).and_then(|v| match v {
+        .borrow()?
-        ServerSoftBans::Unimplemented => None,
+        .get(&guild)
-        ServerSoftBans::Implemented(ref mut v) => Some(v),
+        .map(|v| (v.role, v.periodical_bans.get(&user.id).cloned()));
-    });
+    let (role, current_ban_deadline) = match val {
-
-    match server_ban {
         None => {
-            println!("get here");
+            msg.reply(&ctx, format!("⚠ This server has not enabled the soft-ban feature. Check out `y!a soft-ban-init`.")).await?;
-            msg.reply(&ctx, format!("⚠ This server has not enabled the soft-ban feature. Check out `y!a soft-ban-init`."))?;
+            return Ok(());
         }
-        Some(ref mut server_ban) => {
+        Some(v) => v,
-            let mut member = guild.member(&ctx, &user)?;
+    };
-            match duration {
+
-                None if member.roles.contains(&server_ban.role) => {
+    let mut member = guild.member(&ctx, &user).await?;
-                    msg.reply(&ctx, format!("⛓ Lifting soft-ban for user {}.", user.tag()))?;
+    match duration {
-                    member.remove_role(&ctx, server_ban.role)?;
+        None if member.roles.contains(&role) => {
-                    return Ok(());
+            msg.reply(&ctx, format!("⛓ Lifting soft-ban for user {}.", user.tag()))
-                }
+                .await?;
-                None => {
+            member.remove_role(&ctx, role).await?;
-                    msg.reply(&ctx, format!("⛓ Soft-banning user {}.", user.tag()))?;
+            return Ok(());
-                }
+        }
-                Some(v) => {
+        None => {
-                    let until = Utc::now() + chrono::Duration::from_std(v.0)?;
+            msg.reply(&ctx, format!("⛓ Soft-banning user {}.", user.tag()))
-                    let until = server_ban
+                .await?;
-                        .periodical_bans
+        }
-                        .entry(user.id)
+        Some(v) => {
-                        .and_modify(|v| *v = max(*v, until))
+            // Add the duration into the ban timeout.
-                        .or_insert(until);
+            let until =
-                    msg.reply(
+                current_ban_deadline.unwrap_or(Utc::now()) + chrono::Duration::from_std(v.0)?;
-                        &ctx,
+            msg.reply(
-                        format!("⛓ Soft-banning user {} until {}.", user.tag(), until),
+                &ctx,
-                    )?;
+                format!("⛓ Soft-banning user {} until {}.", user.tag(), until),
-                }
+            )
-            }
+            .await?;
-            member.add_role(&ctx, server_ban.role)?;
+            db.borrow_mut()?
+                .get_mut(&guild)
+                .map(|v| v.periodical_bans.insert(user.id, until));
         }
     }
+    member.add_role(&ctx, role).await?;
 
     Ok(())
 }
@@ -79,86 +80,89 @@ pub fn soft_ban(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResu
 #[usage = "{soft_ban_role_id}"]
 #[num_args(1)]
 #[only_in("guilds")]
-pub fn soft_ban_init(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
+pub async fn soft_ban_init(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let role_id = args.single::<RoleId>()?;
-    let guild = msg.guild(&ctx).ok_or(Error::from("Guild-only command"))?;
+    let guild = msg.guild(&ctx).await.unwrap();
-    let guild = guild.read();
     // Check whether the role_id is the one we wanted
     if !guild.roles.contains_key(&role_id) {
-        return Err(Error::from(format!(
+        Err(Error::msg(format!(
             "{} is not a role in this server.",
             role_id
-        )));
+        )))?;
     }
     // Check if we already set up
-    let db = SoftBans::open(&*ctx.data.read());
+    let db = SoftBans::open(&*ctx.data.read().await);
-    let mut db = db.borrow_mut()?;
+    let set_up = db.borrow()?.contains_key(&guild.id);
-    let server = db
-        .get(&guild.id)
-        .map(|v| match v {
-            ServerSoftBans::Unimplemented => false,
-            _ => true,
-        })
-        .unwrap_or(false);
 
-    if !server {
+    if !set_up {
-        db.insert(guild.id, ServerSoftBans::new_implemented(role_id));
+        db.borrow_mut()?
-        msg.react(&ctx, "👌")?;
+            .insert(guild.id, ServerSoftBans::new(role_id));
-        Ok(())
+        msg.react(&ctx, '👌').await?;
     } else {
-        Err(Error::from("Server already set up soft-bans."))
+        Err(Error::msg("Server already set up soft-bans."))?
     }
+    Ok(())
 }
 
 // Watch the soft bans.
-pub fn watch_soft_bans(client: &serenity::Client) -> impl FnOnce() -> () + 'static {
+pub async fn watch_soft_bans(cache_http: Arc<CacheAndHttp>, data: AppData) {
-    let cache_http = {
+    loop {
-        let cache_http = client.cache_and_http.clone();
+        // Scope so that locks are released
-        let cache: serenity::cache::CacheRwLock = cache_http.cache.clone().into();
+        {
-        (cache, cache_http.http.clone())
+            // Poll the data for any changes.
-    };
+            let db = data.read().await;
-    let data = client.data.clone();
+            let db = SoftBans::open(&*db);
-    return move || {
+            let mut db = db.borrow().unwrap().clone();
-        let cache_http = (&cache_http.0, &*cache_http.1);
+            let now = Utc::now();
-        loop {
+            for (server_id, bans) in db.iter_mut() {
-            // Scope so that locks are released
+                let server_name: String = match server_id.to_partial_guild(&*cache_http.http).await
-            {
+                {
-                // Poll the data for any changes.
+                    Err(_) => continue,
-                let db = data.read();
+                    Ok(v) => v.name,
-                let db = SoftBans::open(&*db);
+                };
-                let mut db = db.borrow_mut().expect("Borrowable");
+                let to_remove: Vec<_> = bans
-                let now = Utc::now();
+                    .periodical_bans
-                for (server_id, soft_bans) in db.iter_mut() {
+                    .iter()
-                    let server_name: String = match server_id.to_partial_guild(cache_http) {
+                    .filter_map(|(user, time)| if time <= &now { Some(user) } else { None })
-                        Err(_) => continue,
+                    .cloned()
-                        Ok(v) => v.name,
+                    .collect();
-                    };
+                if let Err(e) = to_remove
-                    if let ServerSoftBans::Implemented(ref mut bans) = soft_bans {
+                    .into_iter()
-                        let to_remove: Vec<_> = bans
+                    .map(|user_id| {
-                            .periodical_bans
+                        bans.periodical_bans.remove(&user_id);
-                            .iter()
+                        lift_soft_ban_for(
-                            .filter_map(|(user, time)| if time <= &now { Some(user) } else { None })
+                            &*cache_http,
-                            .cloned()
+                            *server_id,
-                            .collect();
+                            &server_name[..],
-                        for user_id in to_remove {
+                            bans.role,
-                            server_id
+                            user_id,
-                                .member(cache_http, user_id)
+                        )
-                                .and_then(|mut m| {
+                    })
-                                    println!(
+                    .collect::<stream::FuturesUnordered<_>>()
-                                        "Soft-ban for `{}` in server `{}` unlifted.",
+                    .try_collect::<()>()
-                                        m.user.read().name,
+                    .await
-                                        server_name
+                {
-                                    );
+                    eprintln!("Error while scanning soft-bans list: {}", e)
-                                    m.remove_role(cache_http, bans.role)
-                                })
-                                .unwrap_or(());
-                            bans.periodical_bans.remove(&user_id);
-                        }
-                    }
                 }
             }
-            // Sleep the thread for a minute
-            std::thread::sleep(std::time::Duration::from_secs(60))
         }
-    };
+        // Sleep the thread for a minute
+        tokio::time::delay_for(std::time::Duration::from_secs(60)).await
+    }
+}
+
+async fn lift_soft_ban_for(
+    cache_http: &CacheAndHttp,
+    server_id: GuildId,
+    server_name: &str,
+    ban_role: RoleId,
+    user_id: UserId,
+) -> Result<()> {
+    let mut m = server_id.member(cache_http, user_id).await?;
+    println!(
+        "Soft-ban for `{}` in server `{}` unlifted.",
+        m.user.name, server_name
+    );
+    m.remove_role(&cache_http.http, ban_role).await?;
+    Ok(())
 }

youmubot-core/src/db.rs

@@ -14,30 +14,25 @@ pub type Roles = DB<GuildMap<HashMap<RoleId, Role>>>;
 /// For the admin commands:
 ///  - Each server might have a `soft ban` role implemented.
 ///  - We allow periodical `soft ban` applications.
-#[derive(Serialize, Deserialize, Debug, Clone)]
-pub enum ServerSoftBans {
-    Implemented(ImplementedSoftBans),
-    Unimplemented,
-}
-
-impl ServerSoftBans {
-    // Create a new, implemented role.
-    pub fn new_implemented(role: RoleId) -> ServerSoftBans {
-        ServerSoftBans::Implemented(ImplementedSoftBans {
-            role,
-            periodical_bans: HashMap::new(),
-        })
-    }
-}
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
-pub struct ImplementedSoftBans {
+pub struct ServerSoftBans {
     /// The soft-ban role.
     pub role: RoleId,
     /// List of all to-unban people.
     pub periodical_bans: HashMap<UserId, DateTime<Utc>>,
 }
 
+impl ServerSoftBans {
+    // Create a new, implemented role.
+    pub fn new(role: RoleId) -> Self {
+        Self {
+            role,
+            periodical_bans: HashMap::new(),
+        }
+    }
+}
+
 /// Role represents an assignable role.
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct Role {