Merge pull request #240 from Rustmail/239-ping-command

feat(commands): add ping command

Author: Akinator31

rustmail/src/bot.rs

@@ -7,21 +7,28 @@ use crate::prelude::panel_commands::*;
 use crate::prelude::types::*;
 use base64::Engine;
 use rand::RngCore;
-use serenity::all::{ClientBuilder, GatewayIntents};
+use serenity::all::{ClientBuilder, GatewayIntents, ShardManager};
 use serenity::cache::Settings as CacheSettings;
+use serenity::prelude::TypeMapKey;
 use std::collections::HashMap;
 use std::process;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::Mutex;
 use tokio::{select, spawn};
 
+pub struct ShardManagerKey;
+
+impl TypeMapKey for ShardManagerKey {
+    type Value = Arc<ShardManager>;
+}
+
 pub async fn init_bot_state() -> Arc<Mutex<BotState>> {
     let pool = init_database().await.expect("An error occured!");
     println!("Database connected!");
 
     let mut bytes = [0u8; 32];
-    rand::thread_rng().fill_bytes(&mut bytes);
+    rand::rng().fill_bytes(&mut bytes);
     let token = base64::engine::general_purpose::STANDARD.encode(&bytes);
 
     let config = load_config("config.toml");
@@ -136,6 +143,7 @@ pub async fn run_bot(
     registry.register_command(LogsCommand);
     registry.register_command(TakeCommand);
     registry.register_command(ReleaseCommand);
+    registry.register_command(PingCommand);
 
     let registry = Arc::new(registry);
 
@@ -174,6 +182,11 @@ pub async fn run_bot(
         state_lock.bot_http = Some(client.http.clone());
     }
 
+    {
+        let mut data = client.data.write().await;
+        data.insert::<ShardManagerKey>(client.shard_manager.clone());
+    }
+
     if let Err(e) = config.validate_servers(&client.http).await {
         eprintln!("Configuration validation error: {}", e);
         eprintln!(

rustmail/src/commands/mod.rs

@@ -22,6 +22,7 @@ pub mod id;
 pub mod logs;
 pub mod move_thread;
 pub mod new_thread;
+pub mod ping;
 pub mod recover;
 pub mod release;
 pub mod remove_reminder;
@@ -42,6 +43,7 @@ pub use id::*;
 pub use logs::*;
 pub use move_thread::*;
 pub use new_thread::*;
+pub use ping::*;
 pub use recover::*;
 pub use release::*;
 pub use remove_reminder::*;

rustmail/src/commands/ping/mod.rs

@@ -0,0 +1,5 @@
+mod slash_command;
+mod text_command;
+
+pub use slash_command::*;
+pub use text_command::*;

rustmail/src/commands/ping/slash_command/mod.rs

@@ -0,0 +1,3 @@
+pub mod ping;
+
+pub use ping::*;

rustmail/src/commands/ping/slash_command/ping.rs

@@ -0,0 +1,121 @@
+use crate::bot::ShardManagerKey;
+use crate::commands::{BoxFuture, CommunityRegistrable, RegistrableCommand};
+use crate::config::Config;
+use crate::errors::{DiscordError, ModmailError, ModmailResult};
+use crate::handlers::InteractionHandler;
+use crate::i18n::get_translated_message;
+use crate::utils::{MessageBuilder, defer_response};
+use chrono::Utc;
+use serenity::FutureExt;
+use serenity::all::{CommandInteraction, Context, CreateCommand, ResolvedOption};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::time::Instant;
+
+pub struct PingCommand;
+
+#[async_trait::async_trait]
+impl RegistrableCommand for PingCommand {
+    fn as_community(&self) -> Option<&dyn CommunityRegistrable> {
+        None
+    }
+
+    fn name(&self) -> &'static str {
+        "ping"
+    }
+
+    fn doc<'a>(&self, config: &'a Config) -> BoxFuture<'a, String> {
+        async move { get_translated_message(config, "help.ping", None, None, None, None).await }
+            .boxed()
+    }
+
+    fn register(&self, config: &Config) -> BoxFuture<'_, Vec<CreateCommand>> {
+        let config = config.clone();
+
+        Box::pin(async move {
+            let cmd_desc = get_translated_message(
+                &config,
+                "slash_command.ping_command_desc",
+                None,
+                None,
+                None,
+                None,
+            )
+            .await;
+
+            vec![CreateCommand::new(self.name()).description(cmd_desc)]
+        })
+    }
+
+    fn run(
+        &self,
+        ctx: &Context,
+        command: &CommandInteraction,
+        _options: &[ResolvedOption<'_>],
+        config: &Config,
+        _handler: Arc<InteractionHandler>,
+    ) -> BoxFuture<'_, ModmailResult<()>> {
+        let ctx = ctx.clone();
+        let command = command.clone();
+        let config = config.clone();
+
+        Box::pin(async move {
+            defer_response(&ctx, &command).await?;
+            let response = MessageBuilder::system_message(&ctx, &config)
+                .content("...")
+                .to_channel(command.channel_id)
+                .build_interaction_message_followup()
+                .await;
+
+            let shard_manager = ctx
+                .data
+                .read()
+                .await
+                .get::<ShardManagerKey>()
+                .cloned()
+                .ok_or(ModmailError::Discord(DiscordError::ShardManagerNotFound))?;
+
+            let time_before = Instant::now();
+            let mut res = command.create_followup(&ctx.http, response).await?;
+            let msg_send_ping = time_before.elapsed().as_millis();
+
+            let gateway_ping = {
+                let runners = shard_manager.runners.lock().await;
+                runners.get(&ctx.shard_id).and_then(|runner| runner.latency)
+            };
+
+            let start = Instant::now();
+            ctx.http.get_gateway().await?;
+            let api_ping = start.elapsed();
+
+            let mut params = HashMap::new();
+            params.insert(
+                "gateway_latency".to_string(),
+                format!(
+                    "{:?}",
+                    gateway_ping.unwrap_or(Duration::default()).as_millis()
+                ),
+            );
+            params.insert(
+                "api_latency".to_string(),
+                format!("{:?}", api_ping.as_millis()),
+            );
+            params.insert(
+                "message_latency".to_string(),
+                format!("{:?}", msg_send_ping),
+            );
+
+            let response = MessageBuilder::system_message(&ctx, &config)
+                .translated_content("slash_command.ping_command", Some(&params), None, None)
+                .await
+                .to_channel(command.channel_id)
+                .build_edit_message()
+                .await;
+
+            res.edit(&ctx.http, response).await?;
+
+            Ok(())
+        })
+    }
+}

rustmail/src/commands/ping/text_command/mod.rs

@@ -0,0 +1,3 @@
+pub mod ping;
+
+pub use ping::*;

rustmail/src/commands/ping/text_command/ping.rs

@@ -0,0 +1,67 @@
+use crate::bot::ShardManagerKey;
+use crate::prelude::config::*;
+use crate::prelude::errors::*;
+use crate::prelude::handlers::*;
+use crate::utils::MessageBuilder;
+use chrono::Utc;
+use serenity::all::{Context, Message};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::time::Instant;
+
+pub async fn ping(
+    ctx: Context,
+    msg: Message,
+    config: &Config,
+    _handler: Arc<GuildMessagesHandler>,
+) -> ModmailResult<()> {
+    let shard_manager = ctx
+        .data
+        .read()
+        .await
+        .get::<ShardManagerKey>()
+        .cloned()
+        .ok_or(ModmailError::Discord(DiscordError::ShardManagerNotFound))?;
+
+    let time_before = Instant::now();
+    let mut res = msg.reply(&ctx.http, "...").await?;
+    let msg_send_ping = time_before.elapsed().as_millis();
+
+    let gateway_ping = {
+        let runners = shard_manager.runners.lock().await;
+        runners.get(&ctx.shard_id).and_then(|runner| runner.latency)
+    };
+
+    let start = Instant::now();
+    ctx.http.get_gateway().await?;
+    let api_ping = start.elapsed();
+
+    let mut params = HashMap::new();
+    params.insert(
+        "gateway_latency".to_string(),
+        format!(
+            "{:?}",
+            gateway_ping.unwrap_or(Duration::default()).as_millis()
+        ),
+    );
+    params.insert(
+        "api_latency".to_string(),
+        format!("{:?}", api_ping.as_millis()),
+    );
+    params.insert(
+        "message_latency".to_string(),
+        format!("{:?}", msg_send_ping),
+    );
+
+    let edited_msg = MessageBuilder::system_message(&ctx, &config)
+        .translated_content("slash_command.ping_command", Some(&params), None, None)
+        .await
+        .to_channel(msg.channel_id)
+        .build_edit_message()
+        .await;
+
+    res.edit(&ctx.http, edited_msg).await?;
+
+    Ok(())
+}

rustmail/src/config.rs

@@ -21,7 +21,8 @@ pub struct Config {
 
     pub db_pool: Option<SqlitePool>,
     pub error_handler: Option<Arc<ErrorHandler>>,
-    pub thread_locks: Arc<std::sync::Mutex<std::collections::HashMap<u64, Arc<tokio::sync::Mutex<()>>>>>,
+    pub thread_locks:
+        Arc<std::sync::Mutex<std::collections::HashMap<u64, Arc<tokio::sync::Mutex<()>>>>>,
 }
 
 fn get_local_ip() -> Option<String> {
@@ -50,21 +51,30 @@ pub fn load_config(path: &str) -> Option<Config> {
     }
 
     if u64::from_str_radix(&config_response.thread.user_message_color, 16).is_err() {
-        eprintln!("Incorrect user message color in the config.toml! Please put a color in hex format!");
+        eprintln!(
+            "Incorrect user message color in the config.toml! Please put a color in hex format!"
+        );
         return None;
     }
 
     if u64::from_str_radix(&config_response.thread.staff_message_color, 16).is_err() {
-        eprintln!("Incorrect staff message color in the config.toml! Please put a color in hex format!");
+        eprintln!(
+            "Incorrect staff message color in the config.toml! Please put a color in hex format!"
+        );
         return None;
     }
 
     if u64::from_str_radix(&config_response.reminders.embed_color, 16).is_err() {
-        eprintln!("Incorrect reminder embed color in the config.toml! Please put a color in hex format!");
+        eprintln!(
+            "Incorrect reminder embed color in the config.toml! Please put a color in hex format!"
+        );
         return None;
     }
 
-    if !config_response.language.is_language_supported(config_response.language.get_default_language()) {
+    if !config_response
+        .language
+        .is_language_supported(config_response.language.get_default_language())
+    {
         eprintln!(
             "Warning: Default language '{}' is not in supported languages list",
             config_response.language.default_language

rustmail/src/errors/dictionary.rs

@@ -188,6 +188,9 @@ impl DictionaryManager {
                 DiscordError::UserIsABot => ("discord.user_is_a_bot".to_string(), None),
                 DiscordError::PermissionDenied => ("discord.permission_denied".to_string(), None),
                 DiscordError::DmCreationFailed => ("discord.dm_creation_failed".to_string(), None),
+                DiscordError::ShardManagerNotFound => {
+                    ("discord.shard_manager_not_found".to_string(), None)
+                }
                 _ => ("discord.api_error".to_string(), None),
             },
             ModmailError::Command(cmd_err) => match cmd_err {

rustmail/src/errors/types.rs

@@ -56,6 +56,7 @@ pub enum DiscordError {
     DmCreationFailed,
     FailedToFetchCategories,
     FailedToMoveChannel,
+    ShardManagerNotFound,
 }
 
 #[derive(Debug, Clone)]
@@ -198,6 +199,7 @@ impl fmt::Display for DiscordError {
             DiscordError::DmCreationFailed => write!(f, "Failed to create DM channel"),
             DiscordError::FailedToFetchCategories => write!(f, "Failed to fetch categories"),
             DiscordError::FailedToMoveChannel => write!(f, "Failed to move_thread channel"),
+            DiscordError::ShardManagerNotFound => write!(f, "Shard manager not found"),
         }
     }
 }