Fix CacheUpdate implementations for extra FullEvents (#3066)

Author: mkrasnitski

src/cache/event.rs

@@ -1,4 +1,5 @@
 use std::collections::VecDeque;
+use std::num::NonZeroU16;
 
 use extract_map::entry::Entry;
 
@@ -46,6 +47,7 @@ use crate::model::guild::{
     Role,
     ScheduledEvent,
 };
+use crate::model::id::GuildId;
 use crate::model::user::{CurrentUser, OnlineStatus};
 use crate::model::voice::VoiceState;
 
@@ -107,15 +109,20 @@ impl CacheUpdate for ChannelPinsUpdateEvent {
 }
 
 impl CacheUpdate for GuildCreateEvent {
-    type Output = std::convert::Infallible;
+    type Output = Vec<GuildId>;
 
     fn update(&mut self, cache: &Cache) -> Option<Self::Output> {
         cache.unavailable_guilds.remove(&self.guild.id);
         let guild = self.guild.clone();
 
         cache.guilds.insert(self.guild.id, guild);
 
-        None
+        if cache.unavailable_guilds.len() == 0 {
+            cache.unavailable_guilds.shrink_to_fit();
+            Some(cache.guilds.iter().map(|i| *i.key()).collect())
+        } else {
+            None
+        }
     }
 }
 
@@ -462,24 +469,28 @@ impl CacheUpdate for PresenceUpdateEvent {
 }
 
 impl CacheUpdate for ReadyEvent {
-    type Output = std::convert::Infallible;
+    type Output = NonZeroU16;
 
     fn update(&mut self, cache: &Cache) -> Option<Self::Output> {
         for unavailable in &self.ready.guilds {
             cache.guilds.remove(&unavailable.id);
             cache.unavailable_guilds.insert(unavailable.id, ());
         }
 
-        let shard_data = self.ready.shard.unwrap_or_default();
+        let shard_info = self.ready.shard.unwrap_or_default();
 
-        {
-            let mut cached_shard_data = cache.shard_data.write();
-            cached_shard_data.total = shard_data.total;
-            cached_shard_data.connected.insert(shard_data.id);
-        }
         cache.user.write().clone_from(&self.ready.user);
 
-        None
+        let mut shards = cache.shard_data.write();
+        shards.total = shard_info.total;
+        shards.connected.insert(shard_info.id);
+
+        if shards.connected.len() == shards.total.get() as usize && !shards.has_sent_shards_ready {
+            shards.has_sent_shards_ready = true;
+            Some(shards.total)
+        } else {
+            None
+        }
     }
 }
 

src/cache/mod.rs

@@ -119,10 +119,10 @@ pub type ChannelMessagesRef<'a> = CacheRef<'a, GenericChannelId, VecDeque<Messag
 
 #[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 #[derive(Debug)]
-pub(crate) struct CachedShardData {
-    pub total: NonZeroU16,
-    pub connected: HashSet<ShardId>,
-    pub has_sent_shards_ready: bool,
+struct CachedShardData {
+    total: NonZeroU16,
+    connected: HashSet<ShardId>,
+    has_sent_shards_ready: bool,
 }
 
 /// A cache containing data received from [`Shard`]s.
@@ -179,28 +179,28 @@ pub struct Cache {
     // ---
     /// A map of guilds with full data available. This includes data like [`Role`]s and [`Emoji`]s
     /// that are not available through the REST API.
-    pub(crate) guilds: MaybeMap<GuildId, Guild>,
+    guilds: MaybeMap<GuildId, Guild>,
     /// A list of guilds which are "unavailable".
     ///
     /// Additionally, guilds are always unavailable for bot users when a Ready is received. Guilds
     /// are "sent in" over time through the receiving of [`Event::GuildCreate`]s.
-    pub(crate) unavailable_guilds: MaybeMap<GuildId, ()>,
+    unavailable_guilds: MaybeMap<GuildId, ()>,
 
     // Messages cache:
     // ---
-    pub(crate) messages: DashMap<GenericChannelId, VecDeque<Message>, BuildHasher>,
+    messages: DashMap<GenericChannelId, VecDeque<Message>, BuildHasher>,
 
     // Miscellanous fixed-size data
     // ---
     /// Information about running shards
-    pub(crate) shard_data: RwLock<CachedShardData>,
+    shard_data: RwLock<CachedShardData>,
     /// The current user "logged in" and for which events are being received for.
     ///
     /// The current user contains information that a regular [`User`] does not, such as whether it
     /// is a bot, whether the user is verified, etc.
     ///
     /// Refer to the documentation for [`CurrentUser`] for more information.
-    pub(crate) user: RwLock<CurrentUser>,
+    user: RwLock<CurrentUser>,
     /// The settings for the cache.
     settings: RwLock<Settings>,
 }
@@ -595,7 +595,7 @@ mod test {
                 ..Default::default()
             },
         };
-        assert!(cache.update(&mut guild_create).is_none());
+        assert!(cache.update(&mut guild_create).is_some());
         assert!(cache.update(&mut event).is_none());
 
         let mut guild_delete = GuildDeleteEvent {

src/cache/wrappers.rs

@@ -26,10 +26,6 @@ impl<K: Eq + Hash, V> MaybeMap<K, V> {
         self.0.as_ref()?.get_mut(k)
     }
 
-    pub fn contains(&self, k: &K) -> bool {
-        self.0.as_ref().is_some_and(|m| m.contains_key(k))
-    }
-
     pub fn insert(&self, k: K, v: V) -> Option<V> {
         self.0.as_ref()?.insert(k, v)
     }
@@ -82,6 +78,10 @@ impl<K: Eq + Hash, V> ReadOnlyMapRef<'_, K, V> {
     pub fn len(&self) -> usize {
         self.0.map_or(0, DashMap::len)
     }
+
+    pub fn contains(&self, k: &K) -> bool {
+        self.0.is_some_and(|m| m.contains_key(k))
+    }
 }
 pub struct Hasher(foldhash::fast::FoldHasher);
 impl std::hash::Hasher for Hasher {

src/gateway/client/dispatch.rs

@@ -11,8 +11,6 @@ use crate::internal::tokio::spawn_named;
 use crate::model::channel::ChannelType;
 use crate::model::event::Event;
 use crate::model::guild::Member;
-#[cfg(feature = "cache")]
-use crate::model::id::GuildId;
 
 #[cfg(feature = "cache")]
 macro_rules! if_cache {
@@ -144,7 +142,7 @@ fn update_cache_with_event(
             }
         },
         Event::ChannelDelete(mut event) => {
-            let cached_messages = if_cache!(event.update(cache));
+            let cached_messages = if_cache!(update_cache!(cache, event));
 
             let channel = event.channel;
             if channel.base.kind == ChannelType::Category {
@@ -162,7 +160,7 @@ fn update_cache_with_event(
             pin: event,
         },
         Event::ChannelUpdate(mut event) => {
-            let old_channel = if_cache!(event.update(cache));
+            let old_channel = if_cache!(update_cache!(cache, event));
 
             FullEvent::ChannelUpdate {
                 old: old_channel,
@@ -182,22 +180,13 @@ fn update_cache_with_event(
             unbanned_user: event.user,
         },
         Event::GuildCreate(mut event) => {
-            let is_new = if_cache!(Some(!cache.unavailable_guilds.contains(&event.guild.id)));
-
-            update_cache!(cache, event);
+            let is_new = if_cache!(Some(!cache.unavailable_guilds().contains(&event.guild.id)));
 
             #[cfg(feature = "cache")]
-            {
-                if cache.unavailable_guilds.len() == 0 {
-                    cache.unavailable_guilds.shrink_to_fit();
-
-                    let guild_amount =
-                        cache.guilds.iter().map(|i| *i.key()).collect::<Vec<GuildId>>();
-
-                    extra_event = Some(FullEvent::CacheReady {
-                        guilds: guild_amount,
-                    });
-                }
+            if let Some(guilds) = update_cache!(cache, event) {
+                extra_event = Some(FullEvent::CacheReady {
+                    guilds,
+                });
             }
 
             FullEvent::GuildCreate {
@@ -206,7 +195,7 @@ fn update_cache_with_event(
             }
         },
         Event::GuildDelete(mut event) => {
-            let full = if_cache!(event.update(cache));
+            let full = if_cache!(update_cache!(cache, event));
 
             FullEvent::GuildDelete {
                 incomplete: event.guild,
@@ -232,7 +221,7 @@ fn update_cache_with_event(
             }
         },
         Event::GuildMemberRemove(mut event) => {
-            let member = if_cache!(event.update(cache));
+            let member = if_cache!(update_cache!(cache, event));
 
             FullEvent::GuildMemberRemoval {
                 guild_id: event.guild_id,
@@ -241,7 +230,7 @@ fn update_cache_with_event(
             }
         },
         Event::GuildMemberUpdate(mut event) => {
-            let before = if_cache!(event.update(cache));
+            let before = if_cache!(update_cache!(cache, event));
             let after: Option<Member> = if_cache!({
                 let guild = cache.guild(event.guild_id);
                 guild.and_then(|g| g.members.get(&event.user.id).cloned())
@@ -268,7 +257,7 @@ fn update_cache_with_event(
             }
         },
         Event::GuildRoleDelete(mut event) => {
-            let role = if_cache!(event.update(cache));
+            let role = if_cache!(update_cache!(cache, event));
 
             FullEvent::GuildRoleDelete {
                 guild_id: event.guild_id,
@@ -277,7 +266,7 @@ fn update_cache_with_event(
             }
         },
         Event::GuildRoleUpdate(mut event) => {
-            let before = if_cache!(event.update(cache));
+            let before = if_cache!(update_cache!(cache, event));
 
             FullEvent::GuildRoleUpdate {
                 old_data_if_available: before,
@@ -324,15 +313,15 @@ fn update_cache_with_event(
             guild_id: event.guild_id,
         },
         Event::MessageUpdate(mut event) => {
-            let before = if_cache!(event.update(cache));
+            let before = if_cache!(update_cache!(cache, event));
 
             FullEvent::MessageUpdate {
                 old_if_available: before,
                 event,
             }
         },
         Event::PresenceUpdate(mut event) => {
-            let old_data = if_cache!(event.update(cache));
+            let old_data = if_cache!(update_cache!(cache, event));
 
             FullEvent::PresenceUpdate {
                 old_data,
@@ -354,19 +343,11 @@ fn update_cache_with_event(
             removed_reactions: event.reaction,
         },
         Event::Ready(mut event) => {
-            update_cache!(cache, event);
-
             #[cfg(feature = "cache")]
-            {
-                let mut shards = cache.shard_data.write();
-                if shards.connected.len() == shards.total.get() as usize
-                    && !shards.has_sent_shards_ready
-                {
-                    shards.has_sent_shards_ready = true;
-                    extra_event = Some(FullEvent::ShardsReady {
-                        total_shards: shards.total,
-                    });
-                }
+            if let Some(total_shards) = update_cache!(cache, event) {
+                extra_event = Some(FullEvent::ShardsReady {
+                    total_shards,
+                });
             }
 
             FullEvent::Ready {
@@ -380,7 +361,7 @@ fn update_cache_with_event(
             event,
         },
         Event::UserUpdate(mut event) => {
-            let before = if_cache!(event.update(cache));
+            let before = if_cache!(update_cache!(cache, event));
 
             FullEvent::UserUpdate {
                 old_data: before,
@@ -391,7 +372,7 @@ fn update_cache_with_event(
             event,
         },
         Event::VoiceStateUpdate(mut event) => {
-            let before = if_cache!(event.update(cache));
+            let before = if_cache!(update_cache!(cache, event));
 
             FullEvent::VoiceStateUpdate {
                 old: before,
@@ -445,15 +426,15 @@ fn update_cache_with_event(
             }
         },
         Event::ThreadUpdate(mut event) => {
-            let old = if_cache!(event.update(cache));
+            let old = if_cache!(update_cache!(cache, event));
 
             FullEvent::ThreadUpdate {
                 old,
                 new: event.thread,
             }
         },
         Event::ThreadDelete(mut event) => {
-            let full_thread_data = if_cache!(event.update(cache));
+            let full_thread_data = if_cache!(update_cache!(cache, event));
 
             FullEvent::ThreadDelete {
                 thread: event.thread,

src/model/channel/message.rs

@@ -1137,7 +1137,6 @@ pub struct PollAnswerCount {
 // all tests here require cache, move if non-cache test is added
 #[cfg(all(test, feature = "cache"))]
 mod tests {
-    use dashmap::DashMap;
     use extract_map::ExtractMap;
     use small_fixed_array::FixedArray;
 
@@ -1153,7 +1152,7 @@ mod tests {
         UserId,
     };
     use crate::cache::Cache;
-    use crate::cache::wrappers::MaybeMap;
+    use crate::model::event::GuildCreateEvent;
 
     fn new_extract_map<K, T>(val: T) -> ExtractMap<K, T>
     where
@@ -1204,12 +1203,10 @@ mod tests {
         };
 
         // Cache, with the guild setup.
-        let mut cache = Cache::new();
-        cache.guilds = MaybeMap(Some({
-            let guilds = DashMap::default();
-            guilds.insert(guild.id, guild);
-            guilds
-        }));
+        let cache = Cache::new();
+        cache.update(&mut GuildCreateEvent {
+            guild,
+        });
 
         // The author should only have the one permission, SEND_MESSAGES.
         assert_eq!(message.author_permissions(&cache), Some(Permissions::SEND_MESSAGES));