From 6bf2779d61e32a32531233ddc609d5254b87ae99 Mon Sep 17 00:00:00 2001 From: Natsu Kagami Date: Mon, 23 Nov 2020 02:26:18 -0500 Subject: [PATCH] `osu updatelb` command (#8) * Make `paginate` take a Paginate trait impl, while `paginate_fn` takes a function * Add `updatelb` command * Implement a member cache * Update member queries to use member cache * Allow everyone to updatelb --- Cargo.lock | 1 + youmubot-cf/src/lib.rs | 11 +- youmubot-core/src/community/roles.rs | 2 +- youmubot-osu/src/discord/mod.rs | 12 +- youmubot-osu/src/discord/server_rank.rs | 168 ++++++++++++++++++++---- youmubot-prelude/Cargo.toml | 1 + youmubot-prelude/src/announcer.rs | 15 ++- youmubot-prelude/src/lib.rs | 4 +- youmubot-prelude/src/member_cache.rs | 45 +++++++ youmubot-prelude/src/pagination.rs | 26 ++-- youmubot-prelude/src/setup.rs | 3 + 11 files changed, 241 insertions(+), 47 deletions(-) create mode 100644 youmubot-prelude/src/member_cache.rs diff --git a/Cargo.lock b/Cargo.lock index 42bf36f..1d7f716 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1900,6 +1900,7 @@ dependencies = [ "anyhow", "async-trait", "chrono", + "dashmap", "flume", "futures-util", "reqwest", diff --git a/youmubot-cf/src/lib.rs b/youmubot-cf/src/lib.rs index da182df..e038d8c 100644 --- a/youmubot-cf/src/lib.rs +++ b/youmubot-cf/src/lib.rs @@ -172,7 +172,7 @@ pub async fn ranks(ctx: &Context, m: &Message) -> CommandResult { let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE; let last_updated = ranks.iter().map(|(_, cfu)| cfu.last_update).min().unwrap(); - paginate( + paginate_fn( move |page, ctx, msg| { let ranks = ranks.clone(); Box::pin(async move { @@ -255,16 +255,17 @@ pub async fn contestranks(ctx: &Context, m: &Message, mut args: Args) -> Command let data = ctx.data.read().await; let contest_id: u64 = args.single()?; let guild = m.guild_id.unwrap(); // Guild-only command + let member_cache = data.get::().unwrap(); let members = CfSavedUsers::open(&*data).borrow()?.clone(); let members = members .into_iter() .map(|(user_id, cf_user)| { - guild - .member(&ctx, user_id) + member_cache + .query(&ctx, user_id, guild) .map(|v| v.map(|v| (cf_user.handle, v))) }) .collect::>() - .filter_map(|v| future::ready(v.ok())) + .filter_map(|v| future::ready(v)) .collect::>() .await; let http = data.get::().unwrap(); @@ -301,7 +302,7 @@ pub async fn contestranks(ctx: &Context, m: &Message, mut args: Args) -> Command const ITEMS_PER_PAGE: usize = 10; let total_pages = (ranks.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE; - paginate( + paginate_fn( move |page, ctx, msg| { let contest = contest.clone(); let problems = problems.clone(); diff --git a/youmubot-core/src/community/roles.rs b/youmubot-core/src/community/roles.rs index c04fbca..79e0042 100644 --- a/youmubot-core/src/community/roles.rs +++ b/youmubot-core/src/community/roles.rs @@ -33,7 +33,7 @@ async fn list(ctx: &Context, m: &Message, _: Args) -> CommandResult { const ROLES_PER_PAGE: usize = 8; let pages = (roles.len() + ROLES_PER_PAGE - 1) / ROLES_PER_PAGE; - paginate( + paginate_fn( |page, ctx, msg| { let roles = roles.clone(); Box::pin(async move { diff --git a/youmubot-osu/src/discord/mod.rs b/youmubot-osu/src/discord/mod.rs index bb45194..5ecfb4a 100644 --- a/youmubot-osu/src/discord/mod.rs +++ b/youmubot-osu/src/discord/mod.rs @@ -29,7 +29,7 @@ use db::OsuUser; use db::{OsuLastBeatmap, OsuSavedUsers, OsuUserBests}; use embeds::{beatmap_embed, score_embed, user_embed}; pub use hook::hook; -use server_rank::{LEADERBOARD_COMMAND, SERVER_RANK_COMMAND}; +use server_rank::{LEADERBOARD_COMMAND, SERVER_RANK_COMMAND, UPDATE_LEADERBOARD_COMMAND}; /// The osu! client. pub(crate) struct OsuClient; @@ -59,6 +59,11 @@ pub fn setup( OsuLastBeatmap::insert_into(&mut *data, &path.join("last_beatmaps.yaml"))?; OsuUserBests::insert_into(&mut *data, &path.join("osu_user_bests.yaml"))?; + // Locks + data.insert::( + server_rank::update_lock::UpdateLock::default(), + ); + // API client let http_client = data.get::().unwrap().clone(); let osu_client = Arc::new(OsuHttpClient::new( @@ -89,7 +94,8 @@ pub fn setup( check, top, server_rank, - leaderboard + leaderboard, + update_leaderboard )] #[default_command(std)] struct Osu; @@ -247,7 +253,7 @@ async fn list_plays<'a>( const ITEMS_PER_PAGE: usize = 5; let total_pages = (plays.len() + ITEMS_PER_PAGE - 1) / ITEMS_PER_PAGE; - paginate( + paginate_fn( move |page, ctx, msg| { let plays = plays.clone(); Box::pin(async move { diff --git a/youmubot-osu/src/discord/server_rank.rs b/youmubot-osu/src/discord/server_rank.rs index 357b78d..4e29915 100644 --- a/youmubot-osu/src/discord/server_rank.rs +++ b/youmubot-osu/src/discord/server_rank.rs @@ -4,12 +4,13 @@ use super::{ ModeArg, OsuClient, }; use crate::{ + discord::BeatmapWithMode, models::{Mode, Score}, request::UserID, }; use serenity::{ framework::standard::{macros::command, Args, CommandResult}, - model::channel::Message, + model::{channel::Message, id::UserId}, utils::MessageBuilder, }; use youmubot_prelude::*; @@ -23,18 +24,22 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR let data = ctx.data.read().await; let mode = args.single::().map(|v| v.0).unwrap_or(Mode::Std); let guild = m.guild_id.expect("Guild-only command"); + let member_cache = data.get::().unwrap(); let users = OsuSavedUsers::open(&*data).borrow()?.clone(); let users = users .into_iter() .map(|(user_id, osu_user)| async move { - guild.member(&ctx, user_id).await.ok().and_then(|member| { - osu_user - .pp - .get(mode as usize) - .cloned() - .and_then(|pp| pp) - .map(|pp| (pp, member.distinct(), osu_user.last_update.clone())) - }) + member_cache + .query(&ctx, user_id, guild) + .await + .and_then(|member| { + osu_user + .pp + .get(mode as usize) + .cloned() + .and_then(|pp| pp) + .map(|pp| (pp, member.distinct(), osu_user.last_update.clone())) + }) }) .collect::>() .filter_map(|v| future::ready(v)) @@ -55,7 +60,7 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR let users = std::sync::Arc::new(users); let last_update = last_update.unwrap(); - paginate( + paginate_fn( move |page: u8, ctx: &Context, m: &mut Message| { const ITEMS_PER_PAGE: usize = 10; let users = users.clone(); @@ -101,14 +106,54 @@ pub async fn server_rank(ctx: &Context, m: &Message, mut args: Args) -> CommandR Ok(()) } -#[command("leaderboard")] -#[aliases("lb", "bmranks", "br", "cc")] -#[description = "See the server's ranks on the last seen beatmap"] +pub(crate) mod update_lock { + use serenity::{model::id::GuildId, prelude::TypeMapKey}; + use std::collections::HashSet; + use std::sync::Mutex; + #[derive(Debug, Default)] + pub struct UpdateLock(Mutex>); + + pub struct UpdateLockGuard<'a>(&'a UpdateLock, GuildId); + + impl TypeMapKey for UpdateLock { + type Value = UpdateLock; + } + + impl UpdateLock { + pub fn get(&self, guild: GuildId) -> Option { + let mut set = self.0.lock().unwrap(); + if set.contains(&guild) { + None + } else { + set.insert(guild); + Some(UpdateLockGuard(self, guild)) + } + } + } + + impl<'a> Drop for UpdateLockGuard<'a> { + fn drop(&mut self) { + let mut set = self.0 .0.lock().unwrap(); + set.remove(&self.1); + } + } +} + +#[command("updatelb")] +#[description = "Update the leaderboard on the last seen beatmap"] #[max_args(0)] #[only_in(guilds)] -pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> CommandResult { +pub async fn update_leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> CommandResult { + let guild = m.guild_id.unwrap(); let data = ctx.data.read().await; - let mut osu_user_bests = OsuUserBests::open(&*data); + let update_lock = data.get::().unwrap(); + let update_lock = match update_lock.get(guild) { + None => { + m.reply(&ctx, "Another update is running.").await?; + return Ok(()); + } + Some(v) => v, + }; let bm = match get_beatmap(&*data, m.channel_id)? { Some(bm) => bm, None => { @@ -116,6 +161,86 @@ pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> Command return Ok(()); } }; + let member_cache = data.get::().unwrap(); + // Signal that we are running. + let running_reaction = m.react(&ctx, '⌛').await?; + + // Run a check on everyone in the server basically. + let all_server_users: Vec<(UserId, Vec)> = { + let osu = data.get::().unwrap(); + let osu_users = OsuSavedUsers::open(&*data); + let osu_users = osu_users + .borrow()? + .iter() + .map(|(&user_id, osu_user)| (user_id, osu_user.id)) + .collect::>(); + let beatmap_id = bm.0.beatmap_id; + osu_users + .into_iter() + .map(|(user_id, osu_id)| { + member_cache + .query(&ctx, user_id, guild) + .map(move |t| t.map(|_| (user_id, osu_id))) + }) + .collect::>() + .filter_map(future::ready) + .filter_map(|(member, osu_id)| async move { + let scores = osu + .scores(beatmap_id, |f| f.user(UserID::ID(osu_id))) + .await + .ok(); + scores + .filter(|s| !s.is_empty()) + .map(|scores| (member, scores)) + }) + .collect::>() + .await + }; + let updated_users = all_server_users.len(); + // Update everything. + { + let mut osu_user_bests = OsuUserBests::open(&*data); + let mut osu_user_bests = osu_user_bests.borrow_mut()?; + let user_bests = osu_user_bests.entry((bm.0.beatmap_id, bm.1)).or_default(); + all_server_users.into_iter().for_each(|(member, scores)| { + user_bests.insert(member, scores); + }) + } + // Signal update complete. + running_reaction.delete(&ctx).await.ok(); + m.reply( + &ctx, + format!( + "update for beatmap ({}, {}) complete, {} users updated.", + bm.0.beatmap_id, bm.1, updated_users + ), + ) + .await + .ok(); + drop(update_lock); + show_leaderboard(ctx, m, bm).await +} + +#[command("leaderboard")] +#[aliases("lb", "bmranks", "br", "cc")] +#[description = "See the server's ranks on the last seen beatmap"] +#[max_args(0)] +#[only_in(guilds)] +pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> CommandResult { + let data = ctx.data.read().await; + let bm = match get_beatmap(&*data, m.channel_id)? { + Some(bm) => bm, + None => { + m.reply(&ctx, "No beatmap queried on this channel.").await?; + return Ok(()); + } + }; + show_leaderboard(ctx, m, bm).await +} + +async fn show_leaderboard(ctx: &Context, m: &Message, bm: BeatmapWithMode) -> CommandResult { + let data = ctx.data.read().await; + let mut osu_user_bests = OsuUserBests::open(&*data); // Run a check on the user once too! { @@ -139,6 +264,7 @@ pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> Command } let guild = m.guild_id.expect("Guild-only command"); + let member_cache = data.get::().unwrap(); let scores = { const NO_SCORES: &'static str = "No scores have been recorded for this beatmap. Run `osu check` to scan for yours!"; @@ -161,12 +287,10 @@ pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> Command let mut scores: Vec<(f64, String, Score)> = users .into_iter() - .map(|(user_id, scores)| async move { - guild - .member(&ctx, user_id) - .await - .ok() - .and_then(|m| Some((m.distinct(), scores))) + .map(|(user_id, scores)| { + member_cache + .query(&ctx, user_id, guild) + .map(|m| m.map(move |m| (m.distinct(), scores))) }) .collect::>() .filter_map(|v| future::ready(v)) @@ -192,7 +316,7 @@ pub async fn leaderboard(ctx: &Context, m: &Message, mut _args: Args) -> Command .await?; return Ok(()); } - paginate( + paginate_fn( move |page: u8, ctx: &Context, m: &mut Message| { const ITEMS_PER_PAGE: usize = 5; let start = (page as usize) * ITEMS_PER_PAGE; diff --git a/youmubot-prelude/Cargo.toml b/youmubot-prelude/Cargo.toml index 2fd42a0..3873dfe 100644 --- a/youmubot-prelude/Cargo.toml +++ b/youmubot-prelude/Cargo.toml @@ -15,6 +15,7 @@ youmubot-db = { path = "../youmubot-db" } reqwest = "0.10" chrono = "0.4" flume = "0.9" +dashmap = "3" [dependencies.serenity] version = "0.9" diff --git a/youmubot-prelude/src/announcer.rs b/youmubot-prelude/src/announcer.rs index ef6693e..aed2707 100644 --- a/youmubot-prelude/src/announcer.rs +++ b/youmubot-prelude/src/announcer.rs @@ -1,4 +1,4 @@ -use crate::{AppData, Result}; +use crate::{AppData, MemberCache, Result}; use async_trait::async_trait; use futures_util::{ future::{join_all, ready, FutureExt}, @@ -46,7 +46,7 @@ pub trait Announcer: Send { } /// A simple struct that allows looking up the relevant channels to an user. -pub struct MemberToChannels(Vec<(GuildId, ChannelId)>); +pub struct MemberToChannels(Vec<(GuildId, ChannelId)>, AppData); impl MemberToChannels { /// Gets the channel list of an user related to that channel. @@ -56,13 +56,14 @@ impl MemberToChannels { u: impl Into, ) -> Vec { let u: UserId = u.into(); + let member_cache = self.1.read().await.get::().unwrap().clone(); self.0 .clone() .into_iter() - .map(|(guild, channel): (GuildId, ChannelId)| { - guild - .member(http.clone(), u) - .map(move |v| v.ok().map(|_| channel.clone())) + .map(|(guild, channel)| { + member_cache + .query(http.clone(), u.into(), guild) + .map(move |t| t.map(|_| channel)) }) .collect::>() .filter_map(|v| ready(v)) @@ -137,7 +138,7 @@ impl AnnouncerHandler { key: &'static str, announcer: &'_ RwLock>, ) -> Result<()> { - let channels = MemberToChannels(Self::get_guilds(&data, key).await?); + let channels = MemberToChannels(Self::get_guilds(&data, key).await?, data.clone()); announcer .write() .await diff --git a/youmubot-prelude/src/lib.rs b/youmubot-prelude/src/lib.rs index 4d7a8ac..6e93af5 100644 --- a/youmubot-prelude/src/lib.rs +++ b/youmubot-prelude/src/lib.rs @@ -6,6 +6,7 @@ use std::sync::Arc; pub mod announcer; pub mod args; pub mod hook; +pub mod member_cache; pub mod pagination; pub mod ratelimit; pub mod setup; @@ -13,7 +14,8 @@ pub mod setup; pub use announcer::{Announcer, AnnouncerHandler}; pub use args::{Duration, UsernameArg}; pub use hook::Hook; -pub use pagination::paginate; +pub use member_cache::MemberCache; +pub use pagination::{paginate, paginate_fn}; /// Re-exporting async_trait helps with implementing Announcer. pub use async_trait::async_trait; diff --git a/youmubot-prelude/src/member_cache.rs b/youmubot-prelude/src/member_cache.rs new file mode 100644 index 0000000..a89d3bb --- /dev/null +++ b/youmubot-prelude/src/member_cache.rs @@ -0,0 +1,45 @@ +use chrono::{DateTime, Utc}; +use dashmap::DashMap; +use serenity::model::{ + guild::Member, + id::{GuildId, UserId}, +}; +use serenity::{http::CacheHttp, prelude::*}; +use std::sync::Arc; + +const VALID_CACHE_SECONDS: i64 = 15 * 60; // 15 minutes + +/// MemberCache resolves `does User belong to Guild` requests, and store them in a cache. +#[derive(Debug, Default)] +pub struct MemberCache(DashMap<(UserId, GuildId), (Option, DateTime)>); + +impl TypeMapKey for MemberCache { + type Value = Arc; +} + +impl MemberCache { + pub async fn query( + &self, + cache_http: impl CacheHttp, + user_id: UserId, + guild_id: GuildId, + ) -> Option { + let now = Utc::now(); + // Check cache + if let Some(r) = self.0.get(&(user_id, guild_id)) { + if &r.1 > &now { + return r.0.clone(); + } + } + // Query + let t = guild_id.member(&cache_http, user_id).await.ok(); + self.0.insert( + (user_id, guild_id), + ( + t.clone(), + now + chrono::Duration::seconds(VALID_CACHE_SECONDS), + ), + ); + t + } +} diff --git a/youmubot-prelude/src/pagination.rs b/youmubot-prelude/src/pagination.rs index c16802e..2c3023b 100644 --- a/youmubot-prelude/src/pagination.rs +++ b/youmubot-prelude/src/pagination.rs @@ -14,7 +14,7 @@ const ARROW_RIGHT: &'static str = "➡️"; const ARROW_LEFT: &'static str = "⬅️"; #[async_trait::async_trait] -pub trait Paginate { +pub trait Paginate: Send { async fn render(&mut self, page: u8, ctx: &Context, m: &mut Message) -> Result; } @@ -36,12 +36,7 @@ where // Paginate! with a pager function. /// If awaited, will block until everything is done. pub async fn paginate( - mut pager: impl for<'m> FnMut( - u8, - &'m Context, - &'m mut Message, - ) -> std::pin::Pin> + Send + 'm>> - + Send, + mut pager: impl Paginate, ctx: &Context, channel: ChannelId, timeout: std::time::Duration, @@ -56,7 +51,7 @@ pub async fn paginate( message .react(&ctx, ReactionType::try_from(ARROW_RIGHT)?) .await?; - pager(0, ctx, &mut message).await?; + pager.render(0, ctx, &mut message).await?; // Build a reaction collector let mut reaction_collector = message.await_reactions(&ctx).removed(true).await; let mut page = 0; @@ -80,6 +75,21 @@ pub async fn paginate( res } +/// Same as `paginate`, but for function inputs, especially anonymous functions. +pub async fn paginate_fn( + pager: impl for<'m> FnMut( + u8, + &'m Context, + &'m mut Message, + ) -> std::pin::Pin> + Send + 'm>> + + Send, + ctx: &Context, + channel: ChannelId, + timeout: std::time::Duration, +) -> Result<()> { + paginate(pager, ctx, channel, timeout).await +} + // Handle the reaction and return a new page number. async fn handle_reaction( page: u8, diff --git a/youmubot-prelude/src/setup.rs b/youmubot-prelude/src/setup.rs index 2d3704c..081289f 100644 --- a/youmubot-prelude/src/setup.rs +++ b/youmubot-prelude/src/setup.rs @@ -11,4 +11,7 @@ pub fn setup_prelude(db_path: &Path, data: &mut TypeMap) { // Set up the HTTP client. data.insert::(reqwest::Client::new()); + + // Set up the member cache. + data.insert::(std::sync::Arc::new(crate::MemberCache::default())); }