Implement beatmap metadata caching

youmubot-osu/src/discord/announcer.rs

@@ -1,8 +1,9 @@
 use super::db::{OsuSavedUsers, OsuUser};
 use super::{embeds::score_embed, BeatmapWithMode, OsuClient};
 use crate::{
+    discord::beatmap_cache::BeatmapMetaCache,
     models::{Mode, Score},
-    request::{BeatmapRequestKind, UserID},
+    request::UserID,
     Client as Osu,
 };
 use announcer::MemberToChannels;
@@ -22,6 +23,7 @@ pub const ANNOUNCER_KEY: &'static str = "osu";
 /// Announce osu! top scores.
 pub fn updates(c: Arc<CacheAndHttp>, d: AppData, channels: MemberToChannels) -> CommandResult {
     let osu = d.get_cloned::<OsuClient>();
+    let cache = d.get_cloned::<BeatmapMetaCache>();
     // For each user...
     let mut data = OsuSavedUsers::open(&*d.read()).borrow()?.clone();
     for (user_id, osu_user) in data.iter_mut() {
@@ -31,7 +33,17 @@ pub fn updates(c: Arc<CacheAndHttp>, d: AppData, channels: MemberToChannels) ->
         }
         osu_user.pp = match (&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
             .par_iter()
-            .map(|m| handle_user_mode(c.clone(), &osu, &osu_user, *user_id, &channels[..], *m))
+            .map(|m| {
+                handle_user_mode(
+                    c.clone(),
+                    &osu,
+                    &cache,
+                    &osu_user,
+                    *user_id,
+                    &channels[..],
+                    *m,
+                )
+            })
             .collect::<Result<_, _>>()
         {
             Ok(v) => v,
@@ -51,6 +63,7 @@ pub fn updates(c: Arc<CacheAndHttp>, d: AppData, channels: MemberToChannels) ->
 fn handle_user_mode(
     c: Arc<CacheAndHttp>,
     osu: &Osu,
+    cache: &BeatmapMetaCache,
     osu_user: &OsuUser,
     user_id: UserId,
     channels: &[ChannelId],
@@ -63,9 +76,9 @@ fn handle_user_mode(
     scores
         .into_par_iter()
         .filter_map(|(rank, score)| {
-            let beatmap = osu
-                .beatmaps(BeatmapRequestKind::Beatmap(score.beatmap_id), |f| f)
-                .map(|v| BeatmapWithMode(v.into_iter().next().unwrap(), mode));
+            let beatmap = cache
+                .get_beatmap_default(score.beatmap_id)
+                .map(|v| BeatmapWithMode(v, mode));
             match beatmap {
                 Ok(v) => Some((rank, score, v)),
                 Err(e) => {

youmubot-osu/src/discord/beatmap_cache.rs

@@ -0,0 +1,70 @@
+use crate::{
+    models::{ApprovalStatus, Beatmap, Mode},
+    Client,
+};
+use dashmap::DashMap;
+use serenity::framework::standard::CommandError;
+use std::sync::Arc;
+use youmubot_prelude::TypeMapKey;
+
+/// BeatmapMetaCache intercepts beatmap-by-id requests and caches them for later recalling.
+/// Does not cache non-Ranked beatmaps.
+#[derive(Clone, Debug)]
+pub struct BeatmapMetaCache {
+    client: Client,
+    cache: Arc<DashMap<(u64, Mode), Beatmap>>,
+}
+
+impl TypeMapKey for BeatmapMetaCache {
+    type Value = BeatmapMetaCache;
+}
+
+impl BeatmapMetaCache {
+    /// Create a new beatmap cache.
+    pub fn new(client: Client) -> Self {
+        BeatmapMetaCache {
+            client,
+            cache: Arc::new(DashMap::new()),
+        }
+    }
+    fn insert_if_possible(&self, id: u64, mode: Option<Mode>) -> Result<Beatmap, CommandError> {
+        let beatmap = self
+            .client
+            .beatmaps(crate::BeatmapRequestKind::Beatmap(id), |f| {
+                if let Some(mode) = mode {
+                    f.mode(mode, true);
+                }
+                f
+            })
+            .and_then(|v| {
+                v.into_iter()
+                    .next()
+                    .ok_or(CommandError::from("beatmap not found"))
+            })?;
+        if let ApprovalStatus::Ranked(_) = beatmap.approval {
+            self.cache.insert((id, beatmap.mode), beatmap.clone());
+        };
+        Ok(beatmap)
+    }
+    /// Get the given beatmap
+    pub fn get_beatmap(&self, id: u64, mode: Mode) -> Result<Beatmap, CommandError> {
+        self.cache
+            .get(&(id, mode))
+            .map(|b| Ok(b.clone()))
+            .unwrap_or_else(|| self.insert_if_possible(id, Some(mode)))
+    }
+
+    /// Get a beatmap without a mode...
+    pub fn get_beatmap_default(&self, id: u64) -> Result<Beatmap, CommandError> {
+        (&[Mode::Std, Mode::Taiko, Mode::Catch, Mode::Mania])
+            .iter()
+            .filter_map(|&mode| {
+                self.cache
+                    .get(&(id, mode))
+                    .filter(|b| b.mode == mode)
+                    .map(|b| Ok(b.clone()))
+            })
+            .next()
+            .unwrap_or_else(|| self.insert_if_possible(id, None))
+    }
+}

youmubot-osu/src/discord/mod.rs

@@ -1,4 +1,5 @@
 use crate::{
+    discord::beatmap_cache::BeatmapMetaCache,
     discord::oppai_cache::BeatmapCache,
     models::{Beatmap, Mode, Mods, Score, User},
     request::{BeatmapRequestKind, UserID},
@@ -17,6 +18,7 @@ use std::str::FromStr;
 use youmubot_prelude::*;
 
 mod announcer;
+pub(crate) mod beatmap_cache;
 mod cache;
 mod db;
 pub(crate) mod embeds;
@@ -59,11 +61,15 @@ pub fn setup(
 
     // API client
     let http_client = data.get_cloned::<HTTPClient>();
-    data.insert::<OsuClient>(OsuHttpClient::new(
+    let osu_client = OsuHttpClient::new(
         http_client.clone(),
         std::env::var("OSU_API_KEY").expect("Please set OSU_API_KEY as osu! api key."),
-    ));
+    );
+    data.insert::<OsuClient>(osu_client.clone());
     data.insert::<oppai_cache::BeatmapCache>(oppai_cache::BeatmapCache::new(http_client));
+    data.insert::<beatmap_cache::BeatmapMetaCache>(beatmap_cache::BeatmapMetaCache::new(
+        osu_client,
+    ));
 
     // Announcer
     announcers.add(announcer::ANNOUNCER_KEY, announcer::updates);
@@ -224,7 +230,7 @@ impl FromStr for Nth {
 
 fn list_plays(plays: &[Score], mode: Mode, ctx: Context, m: &Message) -> CommandResult {
     let watcher = ctx.data.get_cloned::<ReactionWatcher>();
-    let osu = ctx.data.get_cloned::<OsuClient>();
+    let osu = ctx.data.get_cloned::<BeatmapMetaCache>();
     let beatmap_cache = ctx.data.get_cloned::<BeatmapCache>();
 
     if plays.is_empty() {
@@ -254,11 +260,7 @@ fn list_plays(plays: &[Score], mode: Mode, ctx: Context, m: &Message) -> Command
                     .enumerate()
                     .map(|(i, v)| {
                         v.get_or_insert_with(|| {
-                            if let Some(b) = osu
-                                .beatmaps(BeatmapRequestKind::Beatmap(plays[i].beatmap_id), |f| f)
-                                .ok()
-                                .and_then(|v| v.into_iter().next())
-                            {
+                            if let Some(b) = osu.get_beatmap(plays[i].beatmap_id, mode).ok() {
                                 let stars = beatmap_cache
                                     .get_beatmap(b.beatmap_id)
                                     .ok()
@@ -350,6 +352,7 @@ pub fn recent(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult
     let user = to_user_id_query(args.single::<UsernameArg>().ok(), &*ctx.data.read(), msg)?;
 
     let osu = ctx.data.get_cloned::<OsuClient>();
+    let meta_cache = ctx.data.get_cloned::<BeatmapMetaCache>();
     let user = osu
         .user(user, |f| f.mode(mode))?
         .ok_or(Error::from("User not found"))?;
@@ -360,12 +363,8 @@ pub fn recent(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult
                 .into_iter()
                 .last()
                 .ok_or(Error::from("No such play"))?;
-            let beatmap = osu
-                .beatmaps(BeatmapRequestKind::Beatmap(recent_play.beatmap_id), |f| {
-                    f.mode(mode, true)
-                })?
-                .into_iter()
-                .next()
+            let beatmap = meta_cache
+                .get_beatmap(recent_play.beatmap_id, mode)
                 .map(|v| BeatmapWithMode(v, mode))
                 .unwrap();
 

youmubot-osu/src/lib.rs

@@ -14,7 +14,7 @@ use std::{convert::TryInto, sync::Arc};
 
 /// Client is the client that will perform calls to the osu! api server.
 /// It's cheap to clone, so do it.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Client {
     key: Arc<String>,
     client: HTTPClient,

youmubot-osu/src/models/mod.rs

@@ -94,7 +94,7 @@ impl fmt::Display for Language {
     }
 }
 
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, std::hash::Hash)]
 pub enum Mode {
     Std,
     Taiko,

youmubot-osu/src/models/raw.rs

@@ -1,6 +1,6 @@
 use serde::Deserialize;
 
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub(crate) struct Beatmap {
     pub approved: String,
     pub submit_date: String,