Merge ShardQueuer functionality into ShardManager (#3108)

Untangles the spaghetti of shard management code by removing the
middleman `ShardQueuer` struct, which used a disorienting mess of mpsc
channels to facilitate message passing between the `ShardManager` and
each of its `ShardRunner`s.

The new `ShardManager::run` function takes care of creating/spawning new
runners as needed. This is fine to call without spawning a separate task
(which was being done for `ShardQueuer`) because
`Client::start_connection` previously already waited for the
`ShardQueuer` loop to (essentially) finish before returning. In other
words, the Client would start the `ShardManager` and then block, but the
waiting was just being facilitated through channels, which wasn't
obvious.

The only messages passed across channels now are `ShardRunnerMessage`
(sent from the manager to a given runner), and a new
`ShardManagerMessage` (sent from a runner back to the manager). For
example, the manager can send a message to a runner telling it to
restart, then the runner will shut itself down and ask the manager to
reboot the shard by sending it a message back.

The shard queue now keeps track of if it should dispatch in concurrent
batches or not. Since we only use concurrent startup when first
initializing all shards, and then switch it off once and don't re-enable
it, this should be fine.

Author: mkrasnitski

examples/e07_shard_manager/src/main.rs

@@ -52,21 +52,22 @@ async fn main() {
     let mut client =
         Client::builder(token, intents).event_handler(Handler).await.expect("Err creating client");
 
-    // Here we clone a lock to the Shard Manager, and then move it into a new thread. The thread
-    // will unlock the manager and print shards' status on a loop.
-    let manager = client.shard_manager.clone();
+    // Here we get a HashMap of of the shards' status that we move into a new thread. A separate
+    // tokio task holds the ownership to each entry, so each one will require acquiring a lock
+    // before reading.
+    let runners = client.shard_manager.runner_info();
 
     tokio::spawn(async move {
         loop {
             sleep(Duration::from_secs(30)).await;
 
-            let shard_runners = manager.runners.lock().await;
-
-            for (id, runner) in shard_runners.iter() {
-                println!(
-                    "Shard ID {} is {} with a latency of {:?}",
-                    id, runner.stage, runner.latency,
-                );
+            for (id, runner) in &runners {
+                if let Ok(runner) = runner.lock() {
+                    println!(
+                        "Shard ID {} is {} with a latency of {:?}",
+                        id, runner.stage, runner.latency,
+                    );
+                }
             }
         }
     });

src/gateway/client/context.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 #[cfg(feature = "cache")]
 pub use crate::cache::Cache;
-use crate::gateway::{ActivityData, ShardMessenger, ShardRunner};
+use crate::gateway::{ActivityData, ShardMessenger};
 use crate::http::{CacheHttp, Http};
 use crate::model::prelude::*;
 
@@ -62,15 +62,15 @@ impl Context {
     /// Create a new Context to be passed to an event handler.
     pub(crate) fn new(
         data: Arc<dyn std::any::Any + Send + Sync>,
-        runner: &ShardRunner,
+        shard_messenger: ShardMessenger,
         shard_id: ShardId,
         http: Arc<Http>,
         #[cfg(feature = "cache")] cache: Arc<Cache>,
     ) -> Context {
         Context {
-            shard: ShardMessenger::new(runner),
-            shard_id,
             data,
+            shard: shard_messenger,
+            shard_id,
             http,
             #[cfg(feature = "cache")]
             cache,

src/gateway/client/mod.rs

@@ -35,32 +35,29 @@ use std::sync::Arc;
 use std::sync::OnceLock;
 use std::time::Duration;
 
-use futures::channel::mpsc::UnboundedReceiver as Receiver;
 use futures::future::BoxFuture;
-use futures::StreamExt as _;
 #[cfg(feature = "tracing_instrument")]
 use tracing::instrument;
 use tracing::{debug, warn};
 
 pub use self::context::Context;
 pub use self::event_handler::{EventHandler, FullEvent, RawEventHandler};
-use super::TransportCompression;
-#[cfg(feature = "cache")]
-use crate::cache::Cache;
-#[cfg(feature = "cache")]
-use crate::cache::Settings as CacheSettings;
-#[cfg(feature = "framework")]
-use crate::framework::Framework;
 #[cfg(feature = "voice")]
-use crate::gateway::VoiceGatewayManager;
-use crate::gateway::{
+use super::VoiceGatewayManager;
+use super::{
     ActivityData,
-    GatewayError,
     PresenceData,
     ShardManager,
     ShardManagerOptions,
+    TransportCompression,
     DEFAULT_WAIT_BETWEEN_SHARD_START,
 };
+#[cfg(feature = "cache")]
+use crate::cache::Cache;
+#[cfg(feature = "cache")]
+use crate::cache::Settings as CacheSettings;
+#[cfg(feature = "framework")]
+use crate::framework::Framework;
 use crate::http::Http;
 use crate::internal::prelude::*;
 use crate::internal::tokio::spawn_named;
@@ -316,9 +313,6 @@ impl IntoFuture for ClientBuilder {
             }
         }
 
-        #[cfg(feature = "voice")]
-        let voice_manager = self.voice_manager;
-
         #[cfg(feature = "cache")]
         let cache = Arc::new(Cache::new_with_settings(self.cache_settings));
 
@@ -337,33 +331,33 @@ impl IntoFuture for ClientBuilder {
 
             #[cfg(feature = "framework")]
             let framework_cell = Arc::new(OnceLock::new());
-            let (shard_manager, shard_manager_ret_value) = ShardManager::new(ShardManagerOptions {
+
+            let shard_manager = ShardManager::new(ShardManagerOptions {
                 token: self.token,
                 data: Arc::clone(&data),
                 event_handler: self.event_handler,
                 raw_event_handler: self.raw_event_handler,
                 #[cfg(feature = "framework")]
                 framework: Arc::clone(&framework_cell),
                 #[cfg(feature = "voice")]
-                voice_manager: voice_manager.clone(),
+                voice_manager: self.voice_manager.clone(),
                 ws_url: Arc::clone(&ws_url),
+                compression: self.compression,
                 shard_total,
+                max_concurrency,
                 #[cfg(feature = "cache")]
                 cache: Arc::clone(&cache),
                 http: Arc::clone(&http),
                 intents,
                 presence: Some(presence),
-                max_concurrency,
                 wait_time_between_shard_start: self.wait_time_between_shard_start,
-                compression: self.compression,
             });
 
             let client = Client {
                 data,
                 shard_manager,
-                shard_manager_return_value: shard_manager_ret_value,
                 #[cfg(feature = "voice")]
-                voice_manager,
+                voice_manager: self.voice_manager,
                 ws_url,
                 #[cfg(feature = "cache")]
                 cache,
@@ -449,43 +443,17 @@ pub struct Client {
     /// # use std::time::Duration;
     /// #
     /// # fn run(client: Client) {
-    /// // Create a clone of the `Arc` containing the shard manager.
-    /// let shard_manager = client.shard_manager.clone();
-    ///
     /// tokio::spawn(async move {
     ///     loop {
-    ///         let count = shard_manager.shards_instantiated().await.len();
+    ///         let count = client.shard_manager.shards_instantiated().len();
     ///         println!("Shard count instantiated: {}", count);
     ///
     ///         tokio::time::sleep(Duration::from_millis(5000)).await;
     ///     }
     /// });
     /// # }
     /// ```
-    ///
-    /// Shutting down all connections after one minute of operation:
-    ///
-    /// ```rust,no_run
-    /// # use serenity::prelude::*;
-    /// # use std::time::Duration;
-    /// #
-    /// # fn run(client: Client) {
-    /// // Create a clone of the `Arc` containing the shard manager.
-    /// let shard_manager = client.shard_manager.clone();
-    ///
-    /// // Create a thread which will sleep for 60 seconds and then have the shard manager
-    /// // shutdown.
-    /// tokio::spawn(async move {
-    ///     tokio::time::sleep(Duration::from_secs(60)).await;
-    ///
-    ///     shard_manager.shutdown_all().await;
-    ///
-    ///     println!("Shutdown shard manager!");
-    /// });
-    /// # }
-    /// ```
-    pub shard_manager: Arc<ShardManager>,
-    shard_manager_return_value: Receiver<Result<(), GatewayError>>,
+    pub shard_manager: ShardManager,
     /// The voice manager for the client.
     ///
     /// This is an ergonomic structure for interfacing over shards' voice
@@ -796,12 +764,7 @@ impl Client {
 
         debug!("Initializing shard info: {} - {}/{}", start_shard, init, total_shards);
 
-        self.shard_manager.initialize(start_shard, init, total_shards);
-        if let Some(Err(err)) = self.shard_manager_return_value.next().await {
-            return Err(Error::Gateway(err));
-        }
-
-        Ok(())
+        self.shard_manager.run(start_shard, init, total_shards).await.map_err(Error::Gateway)
     }
 }
 

src/gateway/sharding/mod.rs

@@ -16,16 +16,9 @@
 //!
 //! ### [`ShardManager`]
 //!
-//! The shard manager provides a clean interface for communicating with shard runners either
-//! individually or collectively, with functions such as [`ShardManager::shutdown`] and
-//! [`ShardManager::restart`] to manage shards in a fine-grained way.
-//!
-//! For most use cases, the [`ShardManager`] will fit all your low-level sharding needs.
-//!
-//! ### [`ShardQueuer`]
-//!
-//! A light wrapper around an mpsc receiver that receives [`ShardQueuerMessage`]s. It should be run
-//! in its own thread so it can receive messages to start shards concurrently in a queue.
+//! The shard manager provides an interface for managing the starting, restarting, and shutdown of
+//! shards by spawning tasks, each containing a [`ShardRunner`], and communicating with each task
+//! using message passing.
 //!
 //! ### [`ShardRunner`]
 //!
@@ -37,7 +30,7 @@
 
 mod shard_manager;
 mod shard_messenger;
-mod shard_queuer;
+mod shard_queue;
 mod shard_runner;
 
 use std::fmt;
@@ -56,11 +49,12 @@ use url::Url;
 
 pub use self::shard_manager::{
     ShardManager,
+    ShardManagerMessage,
     ShardManagerOptions,
     DEFAULT_WAIT_BETWEEN_SHARD_START,
 };
 pub use self::shard_messenger::ShardMessenger;
-pub use self::shard_queuer::{ShardQueue, ShardQueuer, ShardQueuerMessage};
+pub use self::shard_queue::ShardQueue;
 pub use self::shard_runner::{ShardRunner, ShardRunnerMessage, ShardRunnerOptions};
 use super::{ActivityData, ChunkGuildFilter, GatewayError, PresenceData, WsClient};
 use crate::constants::{self, CloseCode};
@@ -779,9 +773,6 @@ pub enum ShardAction {
 pub struct ShardRunnerInfo {
     /// The latency between when a heartbeat was sent and when the acknowledgement was received.
     pub latency: Option<StdDuration>,
-    /// The channel used to communicate with the shard runner, telling it what to do with regards
-    /// to its status.
-    pub runner_tx: ShardMessenger,
     /// The current connection stage of the shard.
     pub stage: ConnectionStage,
 }