feat(gateway): custom shard presences on identify (#1474)

Provides a function that takes a closure which allows setting a shard's
presence based on its ID.

Closes #1472.

Co-authored-by: BlackHoleFox <[email protected]>

gateway/src/cluster/builder.rs

@@ -47,9 +47,10 @@ impl ClusterBuilder {
     pub fn new(token: impl Into<String>, intents: Intents) -> Self {
         Self(
             ClusterConfig {
-                shard_scheme: ShardScheme::Auto,
                 queue: Arc::new(LocalQueue::new()),
                 resume_sessions: HashMap::new(),
+                shard_presence: None,
+                shard_scheme: ShardScheme::Auto,
             },
             ShardBuilder::new(token, intents),
         )
@@ -184,6 +185,22 @@ impl ClusterBuilder {
         self
     }
 
+    /// Set specific shard presences to use when identifying with the gateway.
+    ///
+    /// Accepts a closure. The closure accepts a [`u64`] and returns an
+    /// [`Option<UpdatePresencePayload>`]. This presence will override any set
+    /// by [`presence`], even if the provided closure returns [`None`].
+    ///
+    /// [`presence`]: Self::presence
+    pub fn shard_presence<F>(mut self, shard_presence: F) -> Self
+    where
+        F: Fn(u64) -> Option<UpdatePresencePayload> + Send + Sync + 'static,
+    {
+        self.0.shard_presence = Some(Box::new(shard_presence));
+
+        self
+    }
+
     /// Set the scheme to use for shard managing.
     ///
     /// For example, [`ShardScheme::Auto`] means that the cluster will

gateway/src/cluster/config.rs

@@ -1,16 +1,21 @@
-use super::scheme::ShardScheme;
-use crate::shard::ResumeSession;
-use std::{collections::HashMap, sync::Arc};
+use crate::{cluster::ShardScheme, shard::ResumeSession};
+use std::{
+    collections::HashMap,
+    fmt::{Debug, Formatter, Result as FmtResult},
+    sync::Arc,
+};
 use twilight_gateway_queue::Queue;
+use twilight_model::gateway::payload::outgoing::update_presence::UpdatePresencePayload;
 
 /// Built configuration for a [`Cluster`].
 ///
 /// [`Cluster`]: crate::Cluster
-#[derive(Debug)]
 pub struct Config {
-    pub(super) shard_scheme: ShardScheme,
     pub(super) queue: Arc<dyn Queue>,
     pub(super) resume_sessions: HashMap<u64, ResumeSession>,
+    pub(super) shard_presence:
+        Option<Box<dyn Fn(u64) -> Option<UpdatePresencePayload> + Send + Sync + 'static>>,
+    pub(super) shard_scheme: ShardScheme,
 }
 
 impl Config {
@@ -30,6 +35,17 @@ impl Config {
     }
 }
 
+impl Debug for Config {
+    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
+        f.debug_struct("Config")
+            .field("queue", &self.queue)
+            .field("resume_sessions", &self.resume_sessions)
+            .field("shard_presence", &"<Fn>")
+            .field("shard_scheme", &self.shard_scheme)
+            .finish()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::Config;

gateway/src/cluster/impl.rs

@@ -311,6 +311,10 @@ impl Cluster {
                 shard_config.sequence = Some(data.sequence);
             }
 
+            if let Some(shard_presence) = &config.shard_presence {
+                shard_config.presence = shard_presence(idx)
+            }
+
             let (shard, stream) = Shard::new_with_config(shard_config);
 
             fold.shards.insert(idx, shard);

gateway/src/shard/config.rs

@@ -21,7 +21,7 @@ pub struct Config {
     pub(super) identify_properties: Option<IdentifyProperties>,
     pub(super) intents: Intents,
     pub(super) large_threshold: u64,
-    pub(super) presence: Option<UpdatePresencePayload>,
+    pub(crate) presence: Option<UpdatePresencePayload>,
     pub(super) queue: Arc<dyn Queue>,
     pub(crate) shard: [u64; 2],
     pub(super) token: Box<str>,