slash commands

Introduces a `/ping` slash command and a `/restart` admin slash command

Author: vilgotf

.github/template-values.toml

@@ -0,0 +1,3 @@
+[values]
+admin_guild_id = "123"
+application_id = "456"

.github/workflows/check.yml

@@ -24,7 +24,7 @@ jobs:
         with:
           tool: cargo-generate
 
-      - run: cargo generate --path . --name instance
+      - run: cargo generate --path . --name instance --template-values-file .github/template-values.toml
 
       - run: cargo clippy
         working-directory: instance
@@ -45,7 +45,7 @@ jobs:
         with:
           tool: cargo-generate
 
-      - run: cargo generate --path . --name instance
+      - run: cargo generate --path . --name instance --template-values-file .github/template-values.toml
 
       - run: cargo fmt -- --check
         working-directory: instance

Cargo.toml

@@ -5,6 +5,7 @@ edition = "2024"
 
 [dependencies]
 anyhow = "1"
+dashmap = "6"
 rustls = "0.23"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
@@ -15,3 +16,5 @@ tracing-subscriber = "0.3"
 twilight-gateway = "0.17"
 twilight-http = "0.17"
 twilight-model = "0.17"
+twilight-standby = "0.17"
+twilight-util = { version = "0.17", features = ["builder"] }

cargo-generate.toml

@@ -3,3 +3,7 @@ ignore = [".github", "Cargo.lock", "target"]
 
 [hooks]
 init = ["init-script.rhai"]
+
+[placeholders]
+admin_guild_id = { prompt = "Enter admin guild ID (where admin commands are available)", type = "string" }
+application_id = { prompt = "Enter the application ID (available in the application's dashboard)", type = "string" }

src/command.rs

@@ -0,0 +1,64 @@
+mod ping;
+mod restart;
+
+use twilight_model::{
+    application::{
+        command::Command,
+        interaction::{InteractionData, InteractionType},
+    },
+    gateway::payload::incoming::InteractionCreate,
+};
+
+pub fn admin_commands(shards: u32) -> [Command; 1] {
+    [restart::command(shards)]
+}
+
+pub fn global_commands() -> [Command; 1] {
+    [ping::command()]
+}
+
+#[derive(Clone, Copy, Debug)]
+enum Type {
+    Ping,
+    Restart,
+}
+
+impl From<&str> for Type {
+    fn from(value: &str) -> Self {
+        match value {
+            ping::NAME => Type::Ping,
+            restart::NAME => Type::Restart,
+            _ => panic!("unknown command name: '{value}'"),
+        }
+    }
+}
+
+pub async fn interaction(mut event: Box<InteractionCreate>) -> anyhow::Result<()> {
+    match event.kind {
+        InteractionType::ApplicationCommandAutocomplete => {
+            let InteractionData::ApplicationCommand(data) = event.data.take().unwrap() else {
+                unreachable!();
+            };
+            let kind = data.name.as_str().into();
+
+            match kind {
+                Type::Ping => ping::autocomplete(event, data).await?,
+                Type::Restart => restart::autocomplete(event, data).await?,
+            }
+        }
+        InteractionType::ApplicationCommand => {
+            let InteractionData::ApplicationCommand(data) = event.data.take().unwrap() else {
+                unreachable!();
+            };
+            let kind = data.name.as_str().into();
+
+            match kind {
+                Type::Ping => ping::run(event, data).await?,
+                Type::Restart => restart::run(event, data).await?,
+            }
+        }
+        _ => {}
+    }
+
+    Ok(())
+}

src/command/ping.rs

@@ -0,0 +1,44 @@
+use crate::{APPLICATION_ID, CONTEXT};
+use twilight_model::{
+    application::{
+        command::{Command, CommandType},
+        interaction::application_command::CommandData,
+    },
+    channel::message::MessageFlags,
+    gateway::payload::incoming::InteractionCreate,
+    http::interaction::{InteractionResponse, InteractionResponseData, InteractionResponseType},
+};
+use twilight_util::builder::command::CommandBuilder;
+
+pub const NAME: &str = "ping";
+
+pub fn command() -> Command {
+    CommandBuilder::new(NAME, "Ping the bot", CommandType::ChatInput).build()
+}
+
+pub async fn autocomplete(
+    _event: Box<InteractionCreate>,
+    _data: Box<CommandData>,
+) -> anyhow::Result<()> {
+    Ok(())
+}
+
+pub async fn run(event: Box<InteractionCreate>, _data: Box<CommandData>) -> anyhow::Result<()> {
+    let data = InteractionResponseData {
+        content: Some("Pong!".to_owned()),
+        flags: Some(MessageFlags::EPHEMERAL),
+        ..Default::default()
+    };
+
+    let response = InteractionResponse {
+        kind: InteractionResponseType::ChannelMessageWithSource,
+        data: Some(data),
+    };
+    CONTEXT
+        .http
+        .interaction(APPLICATION_ID)
+        .create_response(event.id, &event.token, &response)
+        .await?;
+
+    Ok(())
+}

src/command/restart.rs

@@ -0,0 +1,191 @@
+use crate::{APPLICATION_ID, CONTEXT, ShardRestartType};
+use std::iter;
+use twilight_gateway::Event;
+use twilight_model::{
+    application::{
+        command::{Command, CommandOptionChoice, CommandOptionChoiceValue, CommandType},
+        interaction::application_command::{CommandData, CommandDataOption, CommandOptionValue},
+    },
+    gateway::payload::incoming::InteractionCreate,
+    guild::Permissions,
+    http::interaction::{InteractionResponse, InteractionResponseData, InteractionResponseType},
+};
+use twilight_util::builder::command::{BooleanBuilder, CommandBuilder, IntegerBuilder};
+
+pub const NAME: &str = "restart";
+
+pub fn command(shards: u32) -> Command {
+    CommandBuilder::new(NAME, "Restart a shard", CommandType::ChatInput)
+        .default_member_permissions(Permissions::empty())
+        .option(
+            IntegerBuilder::new("id", "Shard ID")
+                .autocomplete(true)
+                .max_value(shards as i64)
+                .min_value(0)
+                .required(true),
+        )
+        .option(BooleanBuilder::new(
+            "resume",
+            "Resume ression? [default: false]",
+        ))
+        .build()
+}
+
+pub async fn autocomplete(
+    event: Box<InteractionCreate>,
+    mut data: Box<CommandData>,
+) -> anyhow::Result<()> {
+    let choice = |shard_id: u32| CommandOptionChoice {
+        name: shard_id.to_string(),
+        name_localizations: None,
+        value: CommandOptionChoiceValue::Integer(shard_id.into()),
+    };
+
+    let mut options = data.options.drain(..);
+    let CommandOptionValue::Focused(value, _) = options.next().unwrap().value else {
+        unreachable!()
+    };
+
+    let choices = match value.parse() {
+        Ok(shard_id) if shard_id == 0 => vec![choice(shard_id)],
+        Ok(shard_id) => starts_with(shard_id, CONTEXT.shard_handles.len() as u32)
+            .take(25)
+            .map(choice)
+            .collect::<Vec<_>>(),
+        Err(_) => Vec::new(),
+    };
+    let data = InteractionResponseData {
+        choices: Some(choices),
+        ..Default::default()
+    };
+
+    let response = InteractionResponse {
+        kind: InteractionResponseType::ApplicationCommandAutocompleteResult,
+        data: Some(data),
+    };
+    CONTEXT
+        .http
+        .interaction(APPLICATION_ID)
+        .create_response(event.id, &event.token, &response)
+        .await?;
+
+    Ok(())
+}
+
+pub async fn run(event: Box<InteractionCreate>, mut data: Box<CommandData>) -> anyhow::Result<()> {
+    let mut options = data.options.drain(..);
+    let CommandOptionValue::Integer(shard_id) = options.next().unwrap().value else {
+        unreachable!()
+    };
+    let kind = match options.next() {
+        Some(CommandDataOption {
+            name: _,
+            value: CommandOptionValue::Boolean(resume),
+        }) => match resume {
+            true => ShardRestartType::Resume,
+            false => ShardRestartType::Normal,
+        },
+        None => ShardRestartType::Normal,
+        Some(_) => unreachable!(),
+    };
+
+    let shard_handle = CONTEXT
+        .shard_handles
+        .get(&(shard_id as u32))
+        .unwrap()
+        .clone();
+    let is_force_restarted = shard_handle.restart(kind).is_forced();
+
+    let response = if is_force_restarted {
+        tracing::debug!(shard_id, "force restarting shard");
+        let data = InteractionResponseData {
+            content: Some("Force restarted shard".to_owned()),
+            ..Default::default()
+        };
+        InteractionResponse {
+            kind: InteractionResponseType::ChannelMessageWithSource,
+            data: Some(data),
+        }
+    } else {
+        tracing::debug!(shard_id, type = ?kind, "restarting shard");
+        InteractionResponse {
+            kind: InteractionResponseType::DeferredChannelMessageWithSource,
+            data: None,
+        }
+    };
+
+    CONTEXT
+        .http
+        .interaction(APPLICATION_ID)
+        .create_response(event.id, &event.token, &response)
+        .await?;
+    if is_force_restarted {
+        return Ok(());
+    }
+
+    shard_handle.restarted().await;
+    let is_shutdown = !CONTEXT
+        .shard_handles
+        .get(&(shard_id as u32))
+        .unwrap()
+        .is_valid();
+
+    let content = match is_shutdown {
+        true => "Bot shut down",
+        false => {
+            tracing::debug!(shard_id, "shard restarted");
+            "Shard restarted"
+        }
+    };
+    CONTEXT
+        .http
+        .interaction(APPLICATION_ID)
+        .update_response(&event.token)
+        .content(Some(content))
+        .await?;
+    if is_shutdown {
+        return Ok(());
+    }
+
+    CONTEXT.standby.wait_for_event(
+        move |event: &Event| matches!(event, Event::Ready(ready) if ready.shard.is_some_and(|id| id.number() == shard_id as u32)),
+    ).await?;
+    tracing::debug!(shard_id, "shard ready");
+    CONTEXT
+        .http
+        .interaction(APPLICATION_ID)
+        .create_followup(&event.token)
+        .content("Shard ready")
+        .await?;
+
+    Ok(())
+}
+
+/// Creates an iterator which computes values up to `max` whose string
+/// representanion starts with `value`.
+///
+/// Produces invalid values if `value` is 0.
+///
+/// # Example
+///
+/// ```
+/// let values = starts_with(5, 100);
+/// assert!(values.eq([5..6, 50..60].into_iter().flatten()));
+///
+/// let values = starts_with(50, 1000);
+/// assert!(values.eq([50..51, 500..510].into_iter().flatten()));
+///
+/// let values = starts_with(5, 1000);
+/// assert!(values.eq([5..6, 50..60, 500..600].into_iter().flatten()));
+/// ```
+fn starts_with(value: u32, max: u32) -> impl Iterator<Item = u32> {
+    debug_assert_ne!(value, 0);
+
+    iter::successors(Some(1_u32), |n| n.checked_mul(10))
+        .take_while(move |&n| value * n <= max)
+        .flat_map(move |n| {
+            let start = value * n;
+            let end = ((value + 1) * n - 1).min(max);
+            start..=end
+        })
+}

src/context.rs

@@ -1,15 +1,24 @@
+use crate::ShardHandle;
+use dashmap::DashMap;
 use std::{ops::Deref, sync::OnceLock};
 use twilight_http::Client;
+use twilight_standby::Standby;
 
 pub static CONTEXT: Ref = Ref(OnceLock::new());
 
 #[derive(Debug)]
 pub struct Context {
     pub http: Client,
+    pub shard_handles: DashMap<u32, ShardHandle>,
+    pub standby: Standby,
 }
 
-pub fn initialize(http: Client) {
-    let context = Context { http };
+pub fn initialize(http: Client, shard_handles: DashMap<u32, ShardHandle>, standby: Standby) {
+    let context = Context {
+        http,
+        shard_handles,
+        standby,
+    };
     assert!(CONTEXT.0.set(context).is_ok());
 }