From 26e8cf9d3af4ec7198e43802de9bb560b85361a4 Mon Sep 17 00:00:00 2001 From: Michael Krasnitski Date: Tue, 26 Nov 2024 01:01:07 -0500 Subject: [PATCH] Fix CacheUpdate implementations for extra FullEvents --- src/cache/event.rs | 37 ++++++++++++------- src/cache/mod.rs | 18 +++++----- src/cache/wrappers.rs | 8 ++--- src/gateway/client/dispatch.rs | 65 ++++++++++++---------------------- 4 files changed, 59 insertions(+), 69 deletions(-) diff --git a/src/cache/event.rs b/src/cache/event.rs index 64cdeb3f016..bebedca6dad 100644 --- a/src/cache/event.rs +++ b/src/cache/event.rs @@ -1,4 +1,5 @@ use std::collections::{HashSet, VecDeque}; +use std::num::NonZeroU16; use super::{Cache, CacheUpdate}; use crate::internal::prelude::*; @@ -33,6 +34,7 @@ use crate::model::event::{ }; use crate::model::gateway::Presence; use crate::model::guild::{Guild, GuildMemberFlags, Member, MemberGeneratedFlags, Role}; +use crate::model::id::GuildId; use crate::model::user::{CurrentUser, OnlineStatus}; use crate::model::voice::VoiceState; @@ -90,15 +92,20 @@ impl CacheUpdate for ChannelPinsUpdateEvent { } impl CacheUpdate for GuildCreateEvent { - type Output = (); + type Output = Vec; - fn update(&mut self, cache: &Cache) -> Option<()> { + fn update(&mut self, cache: &Cache) -> Option { 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 + } } } @@ -435,9 +442,9 @@ impl CacheUpdate for PresenceUpdateEvent { } impl CacheUpdate for ReadyEvent { - type Output = (); + type Output = NonZeroU16; - fn update(&mut self, cache: &Cache) -> Option<()> { + fn update(&mut self, cache: &Cache) -> Option { for unavailable in &self.ready.guilds { cache.guilds.remove(&unavailable.id); cache.unavailable_guilds.insert(unavailable.id, ()); @@ -447,12 +454,12 @@ impl CacheUpdate for ReadyEvent { let mut guilds_to_remove = vec![]; let ready_guilds_hashset = self.ready.guilds.iter().map(|status| status.id).collect::>(); - let shard_data = self.ready.shard.unwrap_or_default(); + let shard_info = self.ready.shard.unwrap_or_default(); for guild_entry in cache.guilds.iter() { let guild = guild_entry.key(); // Only handle data for our shard. - if crate::utils::shard_id(*guild, shard_data.total) == shard_data.id.0 + if crate::utils::shard_id(*guild, shard_info.total) == shard_info.id.0 && !ready_guilds_hashset.contains(guild) { guilds_to_remove.push(*guild); @@ -464,14 +471,18 @@ impl CacheUpdate for ReadyEvent { } } - { - 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 + } } } diff --git a/src/cache/mod.rs b/src/cache/mod.rs index 453dea5c17f..81781138d16 100644 --- a/src/cache/mod.rs +++ b/src/cache/mod.rs @@ -117,10 +117,10 @@ pub type ChannelMessagesRef<'a> = CacheRef<'a, ChannelId, VecDeque, Nev #[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))] #[derive(Debug)] -pub(crate) struct CachedShardData { - pub total: NonZeroU16, - pub connected: HashSet, - pub has_sent_shards_ready: bool, +struct CachedShardData { + total: NonZeroU16, + connected: HashSet, + has_sent_shards_ready: bool, } /// A cache containing data received from [`Shard`]s. @@ -172,28 +172,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, + guilds: MaybeMap, /// 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, + unavailable_guilds: MaybeMap, // Messages cache: // --- - pub(crate) messages: DashMap, BuildHasher>, + messages: DashMap, BuildHasher>, // Miscellanous fixed-size data // --- /// Information about running shards - pub(crate) shard_data: RwLock, + shard_data: RwLock, /// 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, + user: RwLock, /// The settings for the cache. settings: RwLock, } diff --git a/src/cache/wrappers.rs b/src/cache/wrappers.rs index 8fe87ce931b..9a82785dd55 100644 --- a/src/cache/wrappers.rs +++ b/src/cache/wrappers.rs @@ -26,10 +26,6 @@ impl MaybeMap { 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 { self.0.as_ref()?.insert(k, v) } @@ -82,6 +78,10 @@ impl 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(fxhash::FxHasher); impl std::hash::Hasher for Hasher { diff --git a/src/gateway/client/dispatch.rs b/src/gateway/client/dispatch.rs index ae16e8d04f0..fc4bc9bc089 100644 --- a/src/gateway/client/dispatch.rs +++ b/src/gateway/client/dispatch.rs @@ -10,8 +10,6 @@ use crate::internal::prelude::*; 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 { @@ -121,7 +119,7 @@ fn update_cache_with_event( } }, Event::ChannelDelete(mut event) => { - let cached_messages = if_cache!(event.update(cache)); + let cached_messages = update_cache!(cache, event); let channel = event.channel; if channel.kind == ChannelType::Category { @@ -139,7 +137,7 @@ fn update_cache_with_event( pin: event, }, Event::ChannelUpdate(mut event) => { - let old_channel = if_cache!(event.update(cache)); + let old_channel = update_cache!(cache, event); FullEvent::ChannelUpdate { old: old_channel, @@ -159,22 +157,12 @@ 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); - - #[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::>(); + let is_new = if_cache!(Some(!cache.unavailable_guilds().contains(&event.guild.id))); - extra_event = Some(FullEvent::CacheReady { - guilds: guild_amount, - }); - } + if let Some(guilds) = update_cache!(cache, event) { + extra_event = Some(FullEvent::CacheReady { + guilds, + }); } FullEvent::GuildCreate { @@ -183,7 +171,7 @@ fn update_cache_with_event( } }, Event::GuildDelete(mut event) => { - let full = if_cache!(event.update(cache)); + let full = update_cache!(cache, event); FullEvent::GuildDelete { incomplete: event.guild, @@ -209,7 +197,7 @@ fn update_cache_with_event( } }, Event::GuildMemberRemove(mut event) => { - let member = if_cache!(event.update(cache)); + let member = update_cache!(cache, event); FullEvent::GuildMemberRemoval { guild_id: event.guild_id, @@ -218,7 +206,7 @@ fn update_cache_with_event( } }, Event::GuildMemberUpdate(mut event) => { - let before = if_cache!(event.update(cache)); + let before = update_cache!(cache, event); let after: Option = if_cache!({ let guild = cache.guild(event.guild_id); guild.and_then(|g| g.members.get(&event.user.id).cloned()) @@ -245,7 +233,7 @@ fn update_cache_with_event( } }, Event::GuildRoleDelete(mut event) => { - let role = if_cache!(event.update(cache)); + let role = update_cache!(cache, event); FullEvent::GuildRoleDelete { guild_id: event.guild_id, @@ -254,7 +242,7 @@ fn update_cache_with_event( } }, Event::GuildRoleUpdate(mut event) => { - let before = if_cache!(event.update(cache)); + let before = update_cache!(cache, event); FullEvent::GuildRoleUpdate { old_data_if_available: before, @@ -301,7 +289,7 @@ fn update_cache_with_event( guild_id: event.guild_id, }, Event::MessageUpdate(mut event) => { - let before = if_cache!(event.update(cache)); + let before = update_cache!(cache, event); let after = if_cache!(cache.message(event.channel_id, event.id).map(|m| m.clone())); FullEvent::MessageUpdate { @@ -311,7 +299,7 @@ fn update_cache_with_event( } }, Event::PresenceUpdate(mut event) => { - let old_data = if_cache!(event.update(cache)); + let old_data = update_cache!(cache, event); FullEvent::PresenceUpdate { old_data, @@ -333,19 +321,10 @@ 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 { @@ -359,7 +338,7 @@ fn update_cache_with_event( event, }, Event::UserUpdate(mut event) => { - let before = if_cache!(event.update(cache)); + let before = update_cache!(cache, event); FullEvent::UserUpdate { old_data: before, @@ -370,7 +349,7 @@ fn update_cache_with_event( event, }, Event::VoiceStateUpdate(mut event) => { - let before = if_cache!(event.update(cache)); + let before = update_cache!(cache, event); FullEvent::VoiceStateUpdate { old: before, @@ -423,7 +402,7 @@ fn update_cache_with_event( } }, Event::ThreadUpdate(mut event) => { - let old = if_cache!(event.update(cache)); + let old = update_cache!(cache, event); FullEvent::ThreadUpdate { old, @@ -431,7 +410,7 @@ fn update_cache_with_event( } }, Event::ThreadDelete(mut event) => { - let full_thread_data = if_cache!(event.update(cache)); + let full_thread_data = update_cache!(cache, event); FullEvent::ThreadDelete { thread: event.thread,