From 4c35fe5327a87cbf942b8aefea824bfee0a1592a Mon Sep 17 00:00:00 2001
From: Ax333l <main@axelen.xyz>
Date: Tue, 20 Dec 2022 13:01:01 +0100
Subject: [PATCH] refactor(moderation): improve unmute scheduling

---
 src/commands/moderation.rs | 18 +++++++------
 src/events/ready.rs        |  5 ++--
 src/main.rs                |  2 +-
 src/utils/moderation.rs    | 53 ++++++++++++++------------------------
 4 files changed, 33 insertions(+), 45 deletions(-)

diff --git a/src/commands/moderation.rs b/src/commands/moderation.rs
index a3c8782..514367c 100644
--- a/src/commands/moderation.rs
+++ b/src/commands/moderation.rs
@@ -176,16 +176,16 @@ pub async fn unmute(
     let author = ctx.author();
 
     let queue = queue_unmute_member(
-        &ctx.discord().http,
-        &ctx.discord().cache,
-        &data.database,
+        ctx.discord().clone(),
+        data.database.clone(),
         ctx.guild_id().unwrap(),
         id,
         configuration.general.mute.role,
         0,
     )
     .await
-    .unwrap();
+    .unwrap()
+    .err();
 
     respond_moderation(
         &ctx,
@@ -243,6 +243,9 @@ pub async fn mute(
     let mute = &configuration.general.mute;
     let guild_id = ctx.guild_id().unwrap();
 
+    let discord = ctx.discord();
+    let cache = &discord.cache;
+
     if let Some(pending_unmute) = data.pending_unmutes.get(&id.0) {
         trace!("Cancelling pending unmute for {}", id.0);
         pending_unmute.abort();
@@ -252,9 +255,8 @@ pub async fn mute(
         data.pending_unmutes.insert(
             id.0,
             queue_unmute_member(
-                &ctx.discord().http,
-                &ctx.discord().cache,
-                &data.database,
+                discord.clone(),
+                data.database.clone(),
                 guild_id,
                 id,
                 mute.role,
@@ -274,7 +276,7 @@ pub async fn mute(
     };
 
     let result = async {
-        if let Some(mut member) = ctx.discord().cache.member(guild_id, id) {
+        if let Some(mut member) = cache.member(guild_id, id) {
             let (is_currently_muted, removed_roles) =
                 crate::utils::moderation::mute_moderation(&ctx, &mut member, mute).await?;
             // Prevent the bot from overriding the "take" field.
diff --git a/src/events/ready.rs b/src/events/ready.rs
index 85de711..b3a6ad8 100644
--- a/src/events/ready.rs
+++ b/src/events/ready.rs
@@ -39,9 +39,8 @@ pub async fn load_muted_members(ctx: &serenity::Context, _: &serenity::Ready) {
                 data.pending_unmutes.insert(
                     member.user.id.0,
                     queue_unmute_member(
-                        &ctx.http,
-                        &ctx.cache,
-                        &data.database,
+                        ctx.clone(),
+                        data.database.clone(),
                         member.guild_id,
                         member.user.id,
                         mute_role_id,
diff --git a/src/main.rs b/src/main.rs
index 5a50e91..2f9a8d8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -31,7 +31,7 @@ impl serenity::TypeMapKey for Data {
 pub struct Data {
     configuration: Configuration,
     database: Arc<Database>,
-    pending_unmutes: HashMap<u64, JoinHandle<Option<Error>>>,
+    pending_unmutes: HashMap<u64, JoinHandle<Result<(), Error>>>,
     api: Api,
 }
 
diff --git a/src/utils/moderation.rs b/src/utils/moderation.rs
index a555d8a..ec4747d 100644
--- a/src/utils/moderation.rs
+++ b/src/utils/moderation.rs
@@ -2,9 +2,7 @@ use std::cmp;
 use std::sync::Arc;
 
 use mongodb::options::FindOptions;
-use poise::serenity_prelude::{
-    Cache, ChannelId, GuildChannel, GuildId, Http, Mentionable, User, UserId,
-};
+use poise::serenity_prelude::{ChannelId, GuildChannel, GuildId, Mentionable, User, UserId};
 use tokio::task::JoinHandle;
 use tracing::{debug, error, warn};
 
@@ -73,21 +71,16 @@ pub async fn mute_on_join(ctx: &serenity::Context, new_member: &mut serenity::Me
 }
 
 pub fn queue_unmute_member(
-    http: &Arc<Http>,
-    cache: &Arc<Cache>,
-    database: &Arc<Database>,
+    ctx: serenity::Context,
+    database: Arc<Database>,
     guild_id: GuildId,
     user_id: UserId,
     mute_role_id: u64,
     mute_duration: u64,
-) -> JoinHandle<Option<Error>> {
-    let cache = cache.clone();
-    let http = http.clone();
-    let database = database.clone();
-
+) -> JoinHandle<Result<(), Error>> {
     tokio::spawn(async move {
         tokio::time::sleep(std::time::Duration::from_secs(mute_duration)).await;
-        let delete_result = database
+        let db_result = database
             .find_and_delete::<Muted>(
                 "muted",
                 Muted {
@@ -97,30 +90,24 @@ pub fn queue_unmute_member(
                 .into(),
                 None,
             )
-            .await;
+            .await?
+            .ok_or("User was not muted.")?;
 
-        if let Err(database_remove_result) = delete_result {
-            Some(database_remove_result)
-        } else if let Some(find_result) = delete_result.unwrap() {
-            let taken_roles = find_result
-                .taken_roles
-                .unwrap()
-                .into_iter()
-                .map(|r| RoleId::from(r.parse::<u64>().unwrap()))
-                .collect::<Vec<_>>();
+        let taken_roles = db_result
+            .taken_roles
+            .unwrap()
+            .into_iter()
+            .map(|r| RoleId::from(r.parse::<u64>().unwrap()))
+            .collect::<Vec<_>>();
 
-            // Update roles if they didn't leave the guild.
-            let mut member = cache.member(guild_id, user_id)?;
-            if let Err(add_role_result) = member.add_roles(&http, &taken_roles).await {
-                Some(Error::from(add_role_result))
-            } else if let Err(remove_result) = member.remove_role(http, mute_role_id).await {
-                Some(Error::from(remove_result))
-            } else {
-                None
-            }
-        } else {
-            None
+        let http = &ctx.http;
+
+        // Update roles if they didn't leave the guild.
+        if let Some(mut member) = ctx.cache.member(guild_id, user_id) {
+            member.add_roles(http, &taken_roles).await?;
+            member.remove_role(http, mute_role_id).await?;
         }
+        Ok(())
     })
 }