aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorzeyla <[email protected]>2018-07-09 16:08:27 -0700
committerGitHub <[email protected]>2018-07-09 16:08:27 -0700
commite6026308b33c80aa33f0001c89cd271cc5cb6687 (patch)
tree2f8f50244c4d7a1234922fef61a72eb8dddbf4c5
parentRemove deprecated use of Colour associated methods (diff)
downloadserenity-e6026308b33c80aa33f0001c89cd271cc5cb6687.tar.xz
serenity-e6026308b33c80aa33f0001c89cd271cc5cb6687.zip
Add a message cache API (#345)
This adds an API for message caching. By default this caches 0 messages per channel. This can be customized when instantiating: ```rust use serenity::cache::{Cache, Settings}; let mut settings = Settings::new(); // Cache 10 messages per channel. settings.max_messages(10); let cache = Cache::new_with_settings(settings); ``` After instantiation: ```rust use serenity::cache::Cache; let mut cache = Cache::new(); cache.settings_mut().max_messages(10); ``` And during runtime through the global cache: ```rust use serenity::CACHE; CACHE.write().settings_mut().max_messages(10); ```
-rw-r--r--src/cache/mod.rs77
-rw-r--r--src/cache/settings.rs50
-rw-r--r--src/client/dispatch.rs8
-rw-r--r--src/lib.rs9
-rw-r--r--src/model/event.rs81
-rw-r--r--tests/test_cache.rs174
6 files changed, 397 insertions, 2 deletions
diff --git a/src/cache/mod.rs b/src/cache/mod.rs
index 3bf5a5a..6b6e1c4 100644
--- a/src/cache/mod.rs
+++ b/src/cache/mod.rs
@@ -47,7 +47,8 @@ use parking_lot::RwLock;
use std::collections::{
hash_map::Entry,
HashMap,
- HashSet
+ HashSet,
+ VecDeque,
};
use std::{
default::Default,
@@ -55,8 +56,12 @@ use std::{
};
mod cache_update;
+mod settings;
pub use self::cache_update::CacheUpdate;
+pub use self::settings::Settings;
+
+type MessageCache = HashMap<ChannelId, HashMap<MessageId, Message>>;
/// A cache of all events received over a [`Shard`], where storing at least
/// some data from the event is possible.
@@ -97,6 +102,12 @@ pub struct Cache {
/// [`Emoji`]: ../model/guild/struct.Emoji.html
/// [`Role`]: ../model/guild/struct.Role.html
pub guilds: HashMap<GuildId, Arc<RwLock<Guild>>>,
+ /// A map of channels to messages.
+ ///
+ /// This is a map of channel IDs to another map of message IDs to messages.
+ ///
+ /// This keeps only the ten most recent messages.
+ pub messages: MessageCache,
/// A map of notes that a user has made for individual users.
///
/// An empty note is equivalent to having no note, and creating an empty
@@ -160,6 +171,15 @@ pub struct Cache {
/// [`PresenceUpdateEvent`]: ../model/event/struct.PresenceUpdateEvent.html
/// [`ReadyEvent`]: ../model/event/struct.ReadyEvent.html
pub users: HashMap<UserId, Arc<RwLock<User>>>,
+ /// Queue of message IDs for each channel.
+ ///
+ /// This is simply a vecdeque so we can keep track of the order of messages
+ /// inserted into the cache. When a maximum number of messages are in a
+ /// channel's cache, we can pop the front and remove that ID from the cache.
+ pub(crate) message_queue: HashMap<ChannelId, VecDeque<MessageId>>,
+ /// The settings for the cache.
+ settings: Settings,
+ __nonexhaustive: (),
}
impl Cache {
@@ -169,6 +189,25 @@ impl Cache {
Self::default()
}
+ /// Creates a new cache instance with settings applied.
+ ///
+ /// # Examples
+ ///
+ /// ```rust
+ /// use serenity::cache::{Cache, Settings};
+ ///
+ /// let mut settings = Settings::new();
+ /// settings.max_messages(10);
+ ///
+ /// let cache = Cache::new_with_settings(settings);
+ /// ```
+ pub fn new_with_settings(settings: Settings) -> Self {
+ Self {
+ settings,
+ ..Default::default()
+ }
+ }
+
/// Fetches the number of [`Member`]s that have not had data received.
///
/// The important detail to note here is that this is the number of
@@ -622,6 +661,38 @@ impl Cache {
.and_then(|g| g.read().roles.get(&role_id).cloned())
}
+ /// Returns an immutable reference to the settings.
+ ///
+ /// # Examples
+ ///
+ /// Printing the maximum number of messages in a channel to be cached:
+ ///
+ /// ```rust
+ /// use serenity::cache::Cache;
+ ///
+ /// let mut cache = Cache::new();
+ /// println!("Max settings: {}", cache.settings().max_messages);
+ /// ```
+ pub fn settings(&self) -> &Settings {
+ &self.settings
+ }
+
+ /// Returns a mutable reference to the settings.
+ ///
+ /// # Examples
+ ///
+ /// Create a new cache and modify the settings afterwards:
+ ///
+ /// ```rust
+ /// use serenity::cache::Cache;
+ ///
+ /// let mut cache = Cache::new();
+ /// cache.settings_mut().max_messages(10);
+ /// ```
+ pub fn settings_mut(&mut self) -> &mut Settings {
+ &mut self.settings
+ }
+
/// Retrieves a `User` from the cache's [`users`] map, if it exists.
///
/// The only advantage of this method is that you can pass in anything that
@@ -704,13 +775,17 @@ impl Default for Cache {
categories: HashMap::default(),
groups: HashMap::with_capacity(128),
guilds: HashMap::default(),
+ messages: HashMap::default(),
notes: HashMap::default(),
presences: HashMap::default(),
private_channels: HashMap::with_capacity(128),
+ settings: Settings::default(),
shard_count: 1,
unavailable_guilds: HashSet::default(),
user: CurrentUser::default(),
users: HashMap::default(),
+ message_queue: HashMap::default(),
+ __nonexhaustive: (),
}
}
}
diff --git a/src/cache/settings.rs b/src/cache/settings.rs
new file mode 100644
index 0000000..e8e456b
--- /dev/null
+++ b/src/cache/settings.rs
@@ -0,0 +1,50 @@
+/// Settings for the cache.
+///
+/// # Examples
+///
+/// Create new settings, specifying the maximum number of messages:
+///
+/// ```rust
+/// use serenity::cache::Settings as CacheSettings;
+///
+/// let mut settings = CacheSettings::new();
+/// settings.max_messages(10);
+/// ```
+#[derive(Clone, Debug, Default)]
+pub struct Settings {
+ /// The maximum number of messages to store in a channel's message cache.
+ ///
+ /// Defaults to 0.
+ pub max_messages: usize,
+ __nonexhaustive: (),
+}
+
+impl Settings {
+ /// Creates new settings to be used with a cache.
+ #[inline]
+ pub fn new() -> Self {
+ Self::default()
+ }
+
+ /// Sets the maximum number of messages to cache in a channel.
+ ///
+ /// Refer to [`max_messages`] for more information.
+ ///
+ /// # Examples
+ ///
+ /// Set the maximum number of messages to cache:
+ ///
+ /// ```rust
+ /// use serenity::cache::Settings;
+ ///
+ /// let mut settings = Settings::new();
+ /// settings.max_messages(10);
+ /// ```
+ ///
+ /// [`max_messages`]: #structfield.max_messages
+ pub fn max_messages(&mut self, max: usize) -> &mut Self {
+ self.max_messages = max;
+
+ self
+ }
+}
diff --git a/src/client/dispatch.rs b/src/client/dispatch.rs
index b91b8d7..04bd8c9 100644
--- a/src/client/dispatch.rs
+++ b/src/client/dispatch.rs
@@ -68,7 +68,9 @@ pub(crate) fn dispatch<H: EventHandler + Send + Sync + 'static>(
shard_id: u64,
) {
match event {
- DispatchEvent::Model(Event::MessageCreate(event)) => {
+ DispatchEvent::Model(Event::MessageCreate(mut event)) => {
+ update!(event);
+
let context = context(data, runner_tx, shard_id);
dispatch_message(
context.clone(),
@@ -103,6 +105,8 @@ pub(crate) fn dispatch<H: EventHandler + Send + Sync + 'static>(
) {
match event {
DispatchEvent::Model(Event::MessageCreate(event)) => {
+ update!(event);
+
let context = context(data, runner_tx, shard_id);
dispatch_message(context, event.message, event_handler, threadpool);
},
@@ -506,6 +510,8 @@ fn handle_event<H: EventHandler + Send + Sync + 'static>(
});
},
DispatchEvent::Model(Event::MessageUpdate(mut event)) => {
+ update!(event);
+
let context = context(data, runner_tx, shard_id);
let event_handler = Arc::clone(event_handler);
diff --git a/src/lib.rs b/src/lib.rs
index 196a3e0..74f050a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -202,6 +202,15 @@ lazy_static! {
/// println!("{}", CACHE.read().user.id);
/// ```
///
+ /// Update the cache's settings to enable caching of channels' messages:
+ ///
+ /// ```rust
+ /// use serenity::CACHE;
+ ///
+ /// // Cache up to the 10 most recent messages per channel.
+ /// CACHE.write().settings_mut().max_messages(10);
+ /// ```
+ ///
/// [`CurrentUser`]: model/struct.CurrentUser.html
/// [`Cache`]: cache/struct.Cache.html
/// [cache module documentation]: cache/index.html
diff --git a/src/model/event.rs b/src/model/event.rs
index d78e7df..7a7e3de 100644
--- a/src/model/event.rs
+++ b/src/model/event.rs
@@ -153,6 +153,9 @@ impl CacheUpdate for ChannelDeleteEvent {
Channel::Private(_) | Channel::Group(_) => unreachable!(),
};
+ // Remove the cached messages for the channel.
+ cache.messages.remove(&self.channel.id());
+
None
}
}
@@ -400,7 +403,11 @@ impl CacheUpdate for GuildDeleteEvent {
// Remove channel entries for the guild if the guild is found.
cache.guilds.remove(&self.guild.id).map(|guild| {
for channel_id in guild.write().channels.keys() {
+ // Remove the channel from the cache.
cache.channels.remove(channel_id);
+
+ // Remove the channel's cached messages.
+ cache.messages.remove(channel_id);
}
guild
@@ -759,6 +766,40 @@ pub struct MessageCreateEvent {
pub message: Message,
}
+#[cfg(feature = "cache")]
+impl CacheUpdate for MessageCreateEvent {
+ /// The oldest message, if the channel's message cache was already full.
+ type Output = Message;
+
+ fn update(&mut self, cache: &mut Cache) -> Option<Self::Output> {
+ let max = cache.settings().max_messages;
+
+ if max == 0 {
+ return None;
+ }
+
+ let messages = cache.messages
+ .entry(self.message.channel_id)
+ .or_insert_with(Default::default);
+ let queue = cache.message_queue
+ .entry(self.message.channel_id)
+ .or_insert_with(Default::default);
+
+ let mut removed_msg = None;
+
+ if messages.len() == max {
+ if let Some(id) = queue.pop_front() {
+ removed_msg = messages.remove(&id);
+ }
+ }
+
+ queue.push_back(self.message.id);
+ messages.insert(self.message.id, self.message.clone());
+
+ removed_msg
+ }
+}
+
impl<'de> Deserialize<'de> for MessageCreateEvent {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> StdResult<Self, D::Error> {
Ok(Self {
@@ -805,6 +846,46 @@ pub struct MessageUpdateEvent {
pub embeds: Option<Vec<Value>>,
}
+#[cfg(feature = "cache")]
+impl CacheUpdate for MessageUpdateEvent {
+ type Output = ();
+
+ fn update(&mut self, cache: &mut Cache) -> Option<Self::Output> {
+ let messages = cache.messages.get_mut(&self.channel_id)?;
+ let message = messages.get_mut(&self.id)?;
+
+ if let Some(attachments) = self.attachments.clone() {
+ message.attachments = attachments;
+ }
+
+ if let Some(content) = self.content.clone() {
+ message.content = content;
+ }
+
+ if let Some(edited_timestamp) = self.edited_timestamp.clone() {
+ message.edited_timestamp = Some(edited_timestamp);
+ }
+
+ if let Some(mentions) = self.mentions.clone() {
+ message.mentions = mentions;
+ }
+
+ if let Some(mention_everyone) = self.mention_everyone {
+ message.mention_everyone = mention_everyone;
+ }
+
+ if let Some(mention_roles) = self.mention_roles.clone() {
+ message.mention_roles = mention_roles;
+ }
+
+ if let Some(pinned) = self.pinned {
+ message.pinned = pinned;
+ }
+
+ None
+ }
+}
+
#[derive(Clone, Debug, Serialize)]
pub struct PresenceUpdateEvent {
pub guild_id: Option<GuildId>,
diff --git a/tests/test_cache.rs b/tests/test_cache.rs
new file mode 100644
index 0000000..833af04
--- /dev/null
+++ b/tests/test_cache.rs
@@ -0,0 +1,174 @@
+#![cfg(feature = "cache")]
+
+extern crate chrono;
+extern crate serde_json;
+extern crate serenity;
+
+use chrono::DateTime;
+use serde_json::{Number, Value};
+use serenity::{
+ cache::{Cache, CacheUpdate, Settings},
+ model::prelude::*,
+ prelude::RwLock,
+};
+use std::{
+ collections::HashMap,
+ sync::Arc,
+};
+
+#[test]
+fn test_cache_messages() {
+ let mut settings = Settings::new();
+ settings.max_messages(2);
+ let mut cache = Cache::new_with_settings(settings);
+
+ // Test inserting one message into a channel's message cache.
+ let datetime = DateTime::parse_from_str(
+ "1983 Apr 13 12:09:14.274 +0000",
+ "%Y %b %d %H:%M:%S%.3f %z",
+ ).unwrap();
+ let mut event = MessageCreateEvent {
+ message: Message {
+ id: MessageId(3),
+ attachments: vec![],
+ author: User {
+ id: UserId(2),
+ avatar: None,
+ bot: false,
+ discriminator: 1,
+ name: "user 1".to_owned(),
+ },
+ channel_id: ChannelId(2),
+ guild_id: Some(GuildId(1)),
+ content: String::new(),
+ edited_timestamp: None,
+ embeds: vec![],
+ kind: MessageType::Regular,
+ member: None,
+ mention_everyone: false,
+ mention_roles: vec![],
+ mentions: vec![],
+ nonce: Value::Number(Number::from(1)),
+ pinned: false,
+ reactions: vec![],
+ timestamp: datetime.clone(),
+ tts: false,
+ webhook_id: None,
+ },
+ };
+ // Check that the channel cache doesn't exist.
+ assert!(!cache.messages.contains_key(&event.message.channel_id));
+ // Add first message, none because message ID 2 doesn't already exist.
+ assert!(event.update(&mut cache).is_none());
+ // None, it only returns the oldest message if the cache was already full.
+ assert!(event.update(&mut cache).is_none());
+ // Assert there's only 1 message in the channel's message cache.
+ assert_eq!(cache.messages.get(&event.message.channel_id).unwrap().len(), 1);
+
+ // Add a second message, assert that channel message cache length is 2.
+ event.message.id = MessageId(4);
+ assert!(event.update(&mut cache).is_none());
+ assert_eq!(cache.messages.get(&event.message.channel_id).unwrap().len(), 2);
+
+ // Add a third message, the first should now be removed.
+ event.message.id = MessageId(5);
+ assert!(event.update(&mut cache).is_some());
+
+ {
+ let channel = cache.messages.get(&event.message.channel_id).unwrap();
+
+ assert_eq!(channel.len(), 2);
+ // Check that the first message is now removed.
+ assert!(!channel.contains_key(&MessageId(3)));
+ }
+
+ let guild_channel = GuildChannel {
+ id: event.message.channel_id,
+ bitrate: None,
+ category_id: None,
+ guild_id: event.message.guild_id.unwrap(),
+ kind: ChannelType::Text,
+ last_message_id: None,
+ last_pin_timestamp: None,
+ name: String::new(),
+ permission_overwrites: vec![],
+ position: 0,
+ topic: None,
+ user_limit: None,
+ nsfw: false,
+ };
+
+ // Add a channel delete event to the cache, the cached messages for that
+ // channel should now be gone.
+ let mut delete = ChannelDeleteEvent {
+ channel: Channel::Guild(Arc::new(RwLock::new(guild_channel.clone()))),
+ };
+ assert!(cache.update(&mut delete).is_none());
+ assert!(!cache.messages.contains_key(&delete.channel.id()));
+
+ // Test deletion of a guild channel's message cache when a GuildDeleteEvent
+ // is received.
+ let mut guild_create = {
+ let mut channels = HashMap::new();
+ channels.insert(ChannelId(2), Arc::new(RwLock::new(guild_channel.clone())));
+
+ GuildCreateEvent {
+ guild: Guild {
+ id: GuildId(1),
+ afk_channel_id: None,
+ afk_timeout: 0,
+ application_id: None,
+ default_message_notifications: DefaultMessageNotificationLevel::All,
+ emojis: HashMap::new(),
+ explicit_content_filter: ExplicitContentFilter::None,
+ features: vec![],
+ icon: None,
+ joined_at: datetime,
+ large: false,
+ member_count: 0,
+ members: HashMap::new(),
+ mfa_level: MfaLevel::None,
+ name: String::new(),
+ owner_id: UserId(3),
+ presences: HashMap::new(),
+ region: String::new(),
+ roles: HashMap::new(),
+ splash: None,
+ system_channel_id: None,
+ verification_level: VerificationLevel::Low,
+ voice_states: HashMap::new(),
+ channels,
+ },
+ }
+ };
+ assert!(cache.update(&mut guild_create).is_none());
+ assert!(cache.update(&mut event).is_none());
+
+ let mut guild_delete = GuildDeleteEvent {
+ guild: PartialGuild {
+ id: GuildId(1),
+ afk_channel_id: None,
+ afk_timeout: 0,
+ default_message_notifications: DefaultMessageNotificationLevel::All,
+ embed_channel_id: None,
+ embed_enabled: false,
+ emojis: HashMap::new(),
+ features: vec![],
+ icon: None,
+ mfa_level: MfaLevel::None,
+ name: String::new(),
+ owner_id: UserId(3),
+ region: String::new(),
+ roles: HashMap::new(),
+ splash: None,
+ verification_level: VerificationLevel::Low,
+ },
+ };
+
+ // The guild existed in the cache, so the cache's guild is returned by the
+ // update.
+ assert!(cache.update(&mut guild_delete).is_some());
+
+ // Assert that the channel's message cache no longer exists.
+ assert!(!cache.messages.contains_key(&ChannelId(2)));
+}