Fix ShardQueue batch size bug (#3121)

If the last batch to be started was not of maximum size, it wouldn't get popped
off the queue. Instead, we should just pop off the queue periodically to ensure
that we fully flush it. 

On client startup, we queue all shards before popping the first batch, so
suboptimal batch sizes are not a concern. Now, the queue also doesn't have to
worry about "concurrent mode" and saturating the size of its batches.

Author: mkrasnitski

src/gateway/sharding/shard_manager.rs

@@ -141,16 +141,18 @@ impl ShardManager {
         shard_init: u16,
         shard_total: NonZeroU16,
     ) -> Result<(), GatewayError> {
-        self.initialize(shard_index, shard_init, shard_total).await;
+        self.initialize(shard_index, shard_init, shard_total);
         loop {
             if let Ok(Some(msg)) =
                 timeout(self.wait_time_between_shard_start, self.manager_rx.next()).await
             {
                 match msg {
-                    ShardManagerMessage::Boot(shard_id) => self.boot(shard_id, false).await,
+                    ShardManagerMessage::Boot(shard_id) => self.queue_for_start(shard_id),
                     ShardManagerMessage::Quit(err) => return Err(err),
                 }
             }
+            let batch = self.queue.pop_batch();
+            self.checked_start(batch).await;
         }
     }
 
@@ -159,24 +161,20 @@ impl ShardManager {
     /// Note that this queues all shards but does not actually start them. To start the manager's
     /// event loop and dispatch [`ShardRunner`]s as they get queued, call [`Self::run`] instead.
     #[cfg_attr(feature = "tracing_instrument", instrument(skip(self)))]
-    pub async fn initialize(&mut self, shard_index: u16, shard_init: u16, shard_total: NonZeroU16) {
+    pub fn initialize(&mut self, shard_index: u16, shard_init: u16, shard_total: NonZeroU16) {
         let shard_to = shard_index + shard_init;
 
         self.shard_total = shard_total;
         for shard_id in shard_index..shard_to {
-            self.boot(ShardId(shard_id), true).await;
+            self.queue_for_start(ShardId(shard_id));
         }
     }
 
     #[cfg_attr(feature = "tracing_instrument", instrument(skip(self)))]
-    async fn boot(&mut self, shard_id: ShardId, concurrent: bool) {
+    fn queue_for_start(&mut self, shard_id: ShardId) {
         info!("Queueing shard {shard_id} for starting");
 
         self.queue.push_back(shard_id);
-        self.queue.set_concurrent(concurrent);
-        if let Some(batch) = self.queue.pop_batch() {
-            self.checked_start(batch).await;
-        }
     }
 
     /// Restarts a shard runner.

src/gateway/sharding/shard_queue.rs

@@ -9,7 +9,6 @@ use crate::internal::prelude::*;
 #[must_use]
 pub struct ShardQueue {
     buckets: FixedArray<VecDeque<ShardId>, u16>,
-    concurrent: bool,
 }
 
 impl ShardQueue {
@@ -19,7 +18,6 @@ impl ShardQueue {
 
         Self {
             buckets,
-            concurrent: false,
         }
     }
 
@@ -40,21 +38,7 @@ impl ShardQueue {
     }
 
     /// Pops a `ShardId` from every bucket containing at least one and returns them all as a `Vec`.
-    ///
-    /// If the queue is in concurrent mode, calling this will only return `Some` if all buckets are
-    /// filled and a maximally-sized batch can be popped, otherwise will return `None`.
-    ///
-    /// If the queue is _not_ in concurrent mode, this will return `Some` as long as at least one
-    /// bucket is filled.
-    pub fn pop_batch(&mut self) -> Option<Vec<ShardId>> {
-        if self.concurrent && !self.buckets.iter().all(|b| !b.is_empty()) {
-            None
-        } else {
-            Some(self.buckets.iter_mut().filter_map(VecDeque::pop_front).collect())
-        }
-    }
-
-    pub fn set_concurrent(&mut self, concurrent: bool) {
-        self.concurrent = concurrent;
+    pub fn pop_batch(&mut self) -> Vec<ShardId> {
+        self.buckets.iter_mut().filter_map(VecDeque::pop_front).collect()
     }
 }