diff --git a/src/cli.rs b/src/cli.rs index cac6dee..7a757c3 100644 --- a/src/cli.rs +++ b/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 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 { + 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 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 diff --git a/src/framework/bot.rs b/src/framework/bot.rs index 4a445be..c9beffd 100644 --- a/src/framework/bot.rs +++ b/src/framework/bot.rs @@ -89,6 +89,7 @@ where } pub struct Bot { + shard: ShardId, state: Arc, prefixes: Arc + Send + Sync>, intents: Intents, @@ -108,6 +109,7 @@ pub struct Bot { impl Bot { pub fn new(state: T, prefixes: impl PrefixGetter + 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 Bot { } } + 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 Bot { 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 { diff --git a/src/main.rs b/src/main.rs index 7d964b4..b16fcc1 100644 --- a/src/main.rs +++ b/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 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)