Skip to content

Commit

Permalink
Fix CacheUpdate implementations for extra FullEvents
Browse files Browse the repository at this point in the history
  • Loading branch information
@mkrasnitski
mkrasnitski committed Nov 28, 2024
1 parent 7827e5a commit 793e787
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 69 deletions.
37 changes: 24 additions & 13 deletions src/cache/event.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::collections::{HashSet, VecDeque};
use std::num::NonZeroU16;

use super::{Cache, CacheUpdate};
use crate::internal::prelude::*;
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -90,15 +92,20 @@ impl CacheUpdate for ChannelPinsUpdateEvent {
}

impl CacheUpdate for GuildCreateEvent {
type Output = ();
type Output = Vec<GuildId>;

fn update(&mut self, cache: &Cache) -> Option<()> {
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
}
}
}

Expand Down Expand Up @@ -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<Self::Output> {
for unavailable in &self.ready.guilds {
cache.guilds.remove(&unavailable.id);
cache.unavailable_guilds.insert(unavailable.id, ());
Expand All @@ -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::<HashSet<_>>();
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);
Expand All @@ -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
}
}
}

Expand Down
18 changes: 9 additions & 9 deletions src/cache/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,10 +117,10 @@ pub type ChannelMessagesRef<'a> = CacheRef<'a, ChannelId, VecDeque<Message>, Nev

#[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.
Expand Down Expand Up @@ -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<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<ChannelId, VecDeque<Message>, BuildHasher>,
messages: DashMap<ChannelId, 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>,
}
Expand Down
8 changes: 4 additions & 4 deletions src/cache/wrappers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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(fxhash::FxHasher);
impl std::hash::Hasher for Hasher {
Expand Down
65 changes: 22 additions & 43 deletions src/gateway/client/dispatch.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
Expand All @@ -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::<Vec<GuildId>>();
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 {
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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<Member> = if_cache!({
let guild = cache.guild(event.guild_id);
guild.and_then(|g| g.members.get(&event.user.id).cloned())
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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 {
Expand All @@ -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,
Expand All @@ -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 {
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -423,15 +402,15 @@ 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,
new: event.thread,
}
},
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,
Expand Down

0 comments on commit 793e787

Please sign in to comment.