diff --git a/src/gateway/client/mod.rs b/src/gateway/client/mod.rs index ca84b3c3cf4..fc3722c3d67 100644 --- a/src/gateway/client/mod.rs +++ b/src/gateway/client/mod.rs @@ -33,6 +33,7 @@ use std::ops::Range; use std::sync::Arc; #[cfg(feature = "framework")] use std::sync::OnceLock; +use std::time::Duration; use futures::channel::mpsc::UnboundedReceiver as Receiver; use futures::future::BoxFuture; @@ -49,7 +50,14 @@ use crate::cache::Settings as CacheSettings; use crate::framework::Framework; #[cfg(feature = "voice")] use crate::gateway::VoiceGatewayManager; -use crate::gateway::{ActivityData, GatewayError, PresenceData, ShardManager, ShardManagerOptions}; +use crate::gateway::{ + ActivityData, + GatewayError, + PresenceData, + ShardManager, + ShardManagerOptions, + DEFAULT_WAIT_BETWEEN_SHARD_START, +}; use crate::http::Http; use crate::internal::prelude::*; use crate::internal::tokio::spawn_named; @@ -73,6 +81,7 @@ pub struct ClientBuilder { event_handler: Option>, raw_event_handler: Option>, presence: PresenceData, + wait_time_between_shard_start: Duration, } impl ClientBuilder { @@ -106,6 +115,7 @@ impl ClientBuilder { event_handler: None, raw_event_handler: None, presence: PresenceData::default(), + wait_time_between_shard_start: DEFAULT_WAIT_BETWEEN_SHARD_START, } } @@ -153,6 +163,19 @@ impl ClientBuilder { self.framework.as_deref() } + /// Sets the time to wait between starting shards. + /// + /// This should only be used when using a gateway proxy, such as [Sandwich] or [Twilight Gateway + /// Proxy], as otherwise this will lead to gateway disconnects if the shard start rate limit is + /// not respected. + /// + /// [Sandwich]: https://github.com/WelcomerTeam/Sandwich-Daemon + /// [Twilight Gateway Proxy]: https://github.com/Gelbpunkt/gateway-proxy + pub fn wait_time_between_shard_start(mut self, wait_time: Duration) -> Self { + self.wait_time_between_shard_start = wait_time; + self + } + /// Sets the voice gateway handler to be used. It will receive voice events sent over the /// gateway and then consider - based on its settings - whether to dispatch a command. #[cfg(feature = "voice")] @@ -318,6 +341,7 @@ impl IntoFuture for ClientBuilder { intents, presence: Some(presence), max_concurrency, + wait_time_between_shard_start: self.wait_time_between_shard_start, }); let client = Client { diff --git a/src/gateway/sharding/mod.rs b/src/gateway/sharding/mod.rs index cdf755fb9b4..9ce61ce27a5 100644 --- a/src/gateway/sharding/mod.rs +++ b/src/gateway/sharding/mod.rs @@ -50,7 +50,11 @@ use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame; use tracing::{debug, error, info, trace, warn}; use url::Url; -pub use self::shard_manager::{ShardManager, ShardManagerOptions}; +pub use self::shard_manager::{ + ShardManager, + ShardManagerOptions, + DEFAULT_WAIT_BETWEEN_SHARD_START, +}; pub use self::shard_messenger::ShardMessenger; pub use self::shard_queuer::{ShardQueue, ShardQueuer, ShardQueuerMessage}; pub use self::shard_runner::{ShardRunner, ShardRunnerMessage, ShardRunnerOptions}; diff --git a/src/gateway/sharding/shard_manager.rs b/src/gateway/sharding/shard_manager.rs index dba609ea3a2..471896d2846 100644 --- a/src/gateway/sharding/shard_manager.rs +++ b/src/gateway/sharding/shard_manager.rs @@ -25,6 +25,9 @@ use crate::internal::prelude::*; use crate::internal::tokio::spawn_named; use crate::model::gateway::GatewayIntents; +/// The default time to wait between starting each shard or set of shards. +pub const DEFAULT_WAIT_BETWEEN_SHARD_START: Duration = Duration::from_secs(5); + /// A manager for handling the status of shards by starting them, restarting them, and stopping /// them when required. /// @@ -50,7 +53,7 @@ use crate::model::gateway::GatewayIntents; /// use std::sync::{Arc, OnceLock}; /// /// use serenity::gateway::client::EventHandler; -/// use serenity::gateway::{ShardManager, ShardManagerOptions}; +/// use serenity::gateway::{ShardManager, ShardManagerOptions, DEFAULT_WAIT_BETWEEN_SHARD_START}; /// use serenity::http::Http; /// use serenity::model::gateway::GatewayIntents; /// use serenity::prelude::*; @@ -84,6 +87,7 @@ use crate::model::gateway::GatewayIntents; /// intents: GatewayIntents::non_privileged(), /// presence: None, /// max_concurrency, +/// wait_time_between_shard_start: DEFAULT_WAIT_BETWEEN_SHARD_START, /// }); /// # Ok(()) /// # } @@ -146,6 +150,7 @@ impl ShardManager { http: opt.http, intents: opt.intents, presence: opt.presence, + wait_time_between_shard_start: opt.wait_time_between_shard_start, }; spawn_named("shard_queuer::run", async move { @@ -372,4 +377,6 @@ pub struct ShardManagerOptions { pub intents: GatewayIntents, pub presence: Option, pub max_concurrency: NonZeroU16, + /// Number of seconds to wait between starting each shard/set of shards start + pub wait_time_between_shard_start: Duration, } diff --git a/src/gateway/sharding/shard_queuer.rs b/src/gateway/sharding/shard_queuer.rs index b16738db5ca..94cad7dc607 100644 --- a/src/gateway/sharding/shard_queuer.rs +++ b/src/gateway/sharding/shard_queuer.rs @@ -31,8 +31,6 @@ use crate::internal::prelude::*; use crate::internal::tokio::spawn_named; use crate::model::gateway::{GatewayIntents, ShardInfo}; -const WAIT_BETWEEN_BOOTS_IN_SECONDS: u64 = 5; - /// The shard queuer is a simple loop that runs indefinitely to manage the startup of shards. /// /// A shard queuer instance _should_ be run in its own thread, due to the blocking nature of the @@ -68,6 +66,8 @@ pub struct ShardQueuer { pub ws_url: Arc, /// The total amount of shards to start. pub shard_total: NonZeroU16, + /// Number of seconds to wait between each start + pub wait_time_between_shard_start: Duration, #[cfg(feature = "cache")] pub cache: Arc, pub http: Arc, @@ -94,14 +94,14 @@ impl ShardQueuer { /// **Note**: This should be run in its own thread due to the blocking nature of the loop. #[cfg_attr(feature = "tracing_instrument", instrument(skip(self)))] pub async fn run(&mut self) { - // We read from the Rx channel in a loop, and use a timeout of 5 seconds so that we don't + // We read from the Rx channel in a loop, and use a timeout of + // {self.WAIT_TIME_BETWEEN_SHARD_START} (5 seconds normally) seconds so that we don't // hang forever. When we receive a command to start a shard, we append it to our queue. The // queue is popped in batches of shards, which are started in parallel. A batch is fired - // every 5 seconds at minimum in order to avoid being ratelimited. - const TIMEOUT: Duration = Duration::from_secs(WAIT_BETWEEN_BOOTS_IN_SECONDS); + // every WAIT_TIME_BETWEEN_SHARD_START at minimum in order to avoid being ratelimited. loop { - if let Ok(msg) = timeout(TIMEOUT, self.rx.next()).await { + if let Ok(msg) = timeout(self.wait_time_between_shard_start, self.rx.next()).await { match msg { Some(ShardQueuerMessage::SetShardTotal(shard_total)) => { self.shard_total = shard_total; @@ -157,14 +157,13 @@ impl ShardQueuer { let Some(instant) = self.last_start else { return }; // We must wait 5 seconds between IDENTIFYs to avoid session invalidations. - let duration = Duration::from_secs(WAIT_BETWEEN_BOOTS_IN_SECONDS); let elapsed = instant.elapsed(); - if elapsed >= duration { + if elapsed >= self.wait_time_between_shard_start { return; } - let to_sleep = duration - elapsed; + let to_sleep = self.wait_time_between_shard_start - elapsed; sleep(to_sleep).await; }