Add support for custom shard ID and count

src/cli.rs

@@ -59,7 +59,9 @@
 //! starting the server. You can also use the `--bind` argument to specify the address to bind the
 //! server to, which defaults to `0.0.0.0:3000`.
 
-use std::net::SocketAddr;
+use std::{net::SocketAddr, str::FromStr};
+
+use twilight_gateway::ShardId;
 
 /// Payphone's argument parser.
 #[derive(clap::Parser)]
@@ -122,6 +124,52 @@ impl Into<log::LevelFilter> for LevelFilter {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Shard {
+    pub id: u32,
+    pub total: u32,
+}
+
+impl Shard {
+    fn new(id: u32, total: u32) -> Self {
+        Shard { id, total }
+    }
+}
+
+impl Default for Shard {
+    fn default() -> Self {
+        Shard::new(0, 1)
+    }
+}
+
+impl FromStr for Shard {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.split_once(":") {
+            Some((id, total)) => {
+                let id = id.parse().map_err(|_| "Invalid shard ID")?;
+                let total = total.parse().map_err(|_| "Invalid shard total")?;
+
+                Ok(Shard { id, total })
+            }
+            None => Err("Invalid shard format, expected `id:total`"),
+        }
+    }
+}
+
+impl ToString for Shard {
+    fn to_string(&self) -> String {
+        format!("{}:{}", self.id, self.total)
+    }
+}
+
+impl Into<ShardId> for Shard {
+    fn into(self) -> ShardId {
+        ShardId::new(self.id, self.total)
+    }
+}
+
 /// The available commands for the Payphone CLI and their available arguments.
 #[derive(clap::Subcommand, Clone)]
 pub enum Commands {
@@ -140,6 +188,15 @@ pub enum Commands {
             default_value_t = String::from("127.0.0.1:9042"))
         ]
         database_addr: String,
+
+        #[arg(
+            long,
+            short,
+            env = "PAYPHONE_SHARD",
+            value_name = "ID:TOTAL",
+            default_value_t = Shard::default())
+        ]
+        shard: Shard,
     },
 
     /// Starts a documentation server for the project which serves `cargo doc` and refreshes on

src/framework/bot.rs

@@ -89,6 +89,7 @@ where
 }
 
 pub struct Bot<T: Send + Sync = ()> {
+    shard: ShardId,
     state: Arc<T>,
     prefixes: Arc<dyn PrefixGetter<T> + Send + Sync>,
     intents: Intents,
@@ -108,6 +109,7 @@ pub struct Bot<T: Send + Sync = ()> {
 impl<T: 'static + Send + Sync> Bot<T> {
     pub fn new(state: T, prefixes: impl PrefixGetter<T> + Send + Sync + 'static) -> Self {
         Self {
+            shard: ShardId::new(0, 1),
             state: Arc::new(state),
             prefixes: Arc::new(prefixes),
             intents: Intents::empty(),
@@ -125,6 +127,11 @@ impl<T: 'static + Send + Sync> Bot<T> {
         }
     }
 
+    pub fn shard(mut self, shard: ShardId) -> Self {
+        self.shard = shard;
+        self
+    }
+
     pub fn intent(mut self, intent: Intents) -> Self {
         self.intents.insert(intent);
         self
@@ -221,7 +228,7 @@ impl<T: 'static + Send + Sync> Bot<T> {
             commands: bot.commands.clone(),
         });
 
-        let mut gateway = Shard::new(ShardId::new(0, 1), token.to_string(), bot.intents);
+        let mut gateway = Shard::new(bot.shard, token.to_string(), bot.intents);
 
         loop {
             if let Some(Ok(event)) = gateway.next_event(EventTypeFlags::all()).await {

src/main.rs

@@ -113,7 +113,7 @@ use clap::Parser;
 use notify::{Event, EventKind, Watcher};
 use tokio::net::TcpListener;
 use tower_http::services::ServeDir;
-use twilight_gateway::Intents;
+use twilight_gateway::{Intents, ShardId};
 
 use crate::{
     cli::{Args, Commands},
@@ -165,8 +165,9 @@ pub async fn main() {
             Commands::Start {
                 token,
                 database_addr,
+                shard,
             } => {
-                start(token, database_addr).await;
+                start(token, database_addr, shard.into()).await;
             }
             Commands::Docs { bind, open } => {
                 docs(bind, open).await;
@@ -202,7 +203,7 @@ impl PrefixGetter<State> for ServerPrefixOrPayphoneMention {
 /// Starts a Payphone shard.
 ///
 /// TODO: Support multiple shards.
-async fn start(token: String, database_addr: String) {
+async fn start(token: String, database_addr: String, shard: ShardId) {
     log::info!("Welcome to Payphone, Nyeki!");
     log::info!(
         "Payphone is running on version v{}",
@@ -218,6 +219,7 @@ async fn start(token: String, database_addr: String) {
     let state = State { db };
 
     let bot = Bot::new(state, ServerPrefixOrPayphoneMention)
+        .shard(shard)
         .intent(Intents::GUILDS)
         .intent(Intents::GUILD_MESSAGES)
         .intent(Intents::MESSAGE_CONTENT)