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