diff --git a/Cargo.lock b/Cargo.lock index 26d1cad..4840710 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -102,6 +102,17 @@ dependencies = [ "tungstenite", ] +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi 0.3.9", +] + [[package]] name = "autocfg" version = "1.0.1" @@ -330,6 +341,19 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "env_logger" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54532e3223c5af90a6a757c90b5c5521564b07e5e7a958681bcd2afad421cdcd" +dependencies = [ + "atty", + "humantime", + "log", + "regex", + "termcolor", +] + [[package]] name = "fake-simd" version = "0.1.2" @@ -592,6 +616,12 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "494b4d60369511e7dea41cf646832512a94e542f68bb9c49e54518e0f468eb47" +[[package]] +name = "humantime" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c1ad908cc71012b7bea4d0c53ba96a8cba9962f048fa68d143376143d863b7a" + [[package]] name = "hyper" version = "0.13.8" @@ -1143,6 +1173,7 @@ dependencies = [ "chrono-tz", "custom_error", "dotenv", + "env_logger", "lazy_static", "log", "num-integer", @@ -1529,6 +1560,15 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "termcolor" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb6bfa289a4d7c5766392812c0a1f4c1ba45afa1ad47803c11e1f407d846d75f" +dependencies = [ + "winapi-util", +] + [[package]] name = "thiserror" version = "1.0.21" @@ -1968,6 +2008,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +dependencies = [ + "winapi 0.3.9", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" diff --git a/Cargo.toml b/Cargo.toml index 6484622..fb69aaf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ sqlx = { version = "0.3.5", default-features = false, features = ["runtime-tokio regex = "1.3.9" async-trait = "0.1.36" log = "0.4.11" +env_logger = "0.8.1" chrono = "0.4" chrono-tz = "0.5" lazy_static = "1.4.0" diff --git a/src/main.rs b/src/main.rs index 22a2a52..d580cb6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,6 +36,8 @@ use crate::{ use serenity::futures::TryFutureExt; +use log::info; + struct SQLPool; impl TypeMapKey for SQLPool { @@ -56,6 +58,8 @@ impl TypeMapKey for FrameworkCtx { #[tokio::main] async fn main() -> Result<(), Box> { + env_logger::init(); + dotenv()?; let token = env::var("DISCORD_TOKEN").expect("Missing DISCORD_TOKEN from environment"); @@ -137,7 +141,45 @@ async fn main() -> Result<(), Box> { data.insert::(framework_arc); } - client.start_autosharded().await?; + if let Ok((Some(lower), Some(upper))) = env::var("SHARD_RANGE").map(|sr| { + let mut split = sr + .split(',') + .map(|val| val.parse::().expect("SHARD_RANGE not an integer")); + + (split.next(), split.next()) + }) { + let total_shards = env::var("SHARD_COUNT") + .map(|shard_count| shard_count.parse::().ok()) + .ok() + .flatten() + .expect("No SHARD_COUNT provided, but SHARD_RANGE was provided"); + + assert!( + lower < upper, + "SHARD_RANGE lower limit is not less than the upper limit" + ); + + info!( + "Starting client fragment with shards {}-{}/{}", + lower, upper, total_shards + ); + + client + .start_shard_range([lower, upper], total_shards) + .await?; + } else if let Ok(total_shards) = env::var("SHARD_COUNT").map(|shard_count| { + shard_count + .parse::() + .expect("SHARD_COUNT not an integer") + }) { + info!("Starting client with {} shards", total_shards); + + client.start_shards(total_shards).await?; + } else { + info!("Starting client as autosharded"); + + client.start_autosharded().await?; + } Ok(()) }