From 6b7103d2e143d3ef22b25bb52e6b751dcb598ce1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mladen=20Brankovi=C4=87?= Date: Sat, 31 Aug 2024 18:23:15 +0200 Subject: [PATCH] wip: delay ready role removal --- Cargo.lock | 8 +-- Cargo.toml | 3 +- src/commands/about.rs | 4 +- src/events/mod.rs | 31 +++++----- src/events/presence_update.rs | 107 +++++++++++++++++++++++++++++++--- src/main.rs | 11 +++- src/util/lunaro_tracking.rs | 16 ++--- src/util/play.rs | 12 +++- 8 files changed, 150 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2c1ed09..749fdb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1036,6 +1036,7 @@ dependencies = [ "serde_json", "thiserror", "tokio", + "tokio-util", "toml", "uuid", ] @@ -2204,16 +2205,15 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.8" +version = "0.7.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "806fe8c2c87eccc8b3267cbae29ed3ab2d0bd37fca70ab622e46aaa9375ddb7d" +checksum = "9cf6b47b3771c49ac75ad09a6162f53ad4b8088b76ac60e8ec1455b31a189fe1" dependencies = [ "bytes", "futures-core", "futures-sink", "pin-project-lite", "tokio", - "tracing", ] [[package]] @@ -2639,7 +2639,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 0115187..ce3d698 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ dotenv = { version = "0.15.0" } futures = { version = "0.3.30" } log = { version = "0.4.22" } log4rs = { version = "1.3.0" } -poise = { version = "0.6.1" } +poise = { version = "0.6.0" } regex = { version = "1.10.6" } reqwest = { version = "0.12.7" } rustc_version_runtime = { version = "0.3.0" } @@ -18,5 +18,6 @@ serde = { version = "1.0.209" } serde_json = { version = "1.0.127" } thiserror = { version = "1.0.63" } tokio = { version = "1.40.0", features = ["rt-multi-thread"] } +tokio-util = { version = "0.7.11" } toml = { version = "0.8.19" } uuid = { version = "1.10.0", features = ["v4"] } diff --git a/src/commands/about.rs b/src/commands/about.rs index ac4cc5d..338c0e0 100644 --- a/src/commands/about.rs +++ b/src/commands/about.rs @@ -38,7 +38,7 @@ struct Committer { pub async fn run(context: PoiseContext<'_>) -> Result<(), Error> { let env = Environment::instance(); - let tracking_config = lunaro_tracking::Config::instance().await; + let tracking_config = lunaro_tracking::LunaroTrackingConfig::instance().await; let member_count = context .partial_guild() .await @@ -99,7 +99,7 @@ pub async fn run(context: PoiseContext<'_>) -> Result<(), Error> { format!("⚙️ [Serenity]() {serenity_version} / [Poise]() {poise_version}"), format!("🚧 [Last updated](<{commit_url}>) {commit_date}",), String::new(), - format!("⏱️ Started online {uptime}"), + format!("⏱️ Instance went online {uptime}"), ] .join("\n"), ), diff --git a/src/events/mod.rs b/src/events/mod.rs index 087d18a..0472d2f 100644 --- a/src/events/mod.rs +++ b/src/events/mod.rs @@ -1,21 +1,24 @@ -use poise::serenity_prelude::{async_trait, Context, Event, RawEventHandler}; +use poise::{ + serenity_prelude::{Context, FullEvent}, + FrameworkContext, +}; + +use crate::types::{error::Error, poise::PoiseData}; mod presence_update; mod ready; -pub struct EventHandlers; - -#[async_trait] -impl RawEventHandler for EventHandlers { - async fn raw_event(&self, context: Context, raw_event: Event) { - let handler = match &raw_event { - Event::Ready(event) => ready::handle(context, &event.ready).await, - Event::PresenceUpdate(event) => presence_update::handle(context, &event.presence).await, - _ => Ok(()), - }; +pub async fn event_handler( + context: &Context, + event: &FullEvent, + _framework: FrameworkContext<'_, PoiseData, Error>, + _data: &PoiseData, +) -> Result<(), Error> { + log::debug!("Received event: {:#?}", event); - if let Err(error) = handler { - log::warn!("Error while handling {:?}: {error:?}: {error}", raw_event); - } + match event { + FullEvent::Ready { data_about_bot } => ready::handle(context, data_about_bot).await, + FullEvent::PresenceUpdate { new_data } => presence_update::handle(context, new_data).await, + _ => Ok(()), } } diff --git a/src/events/presence_update.rs b/src/events/presence_update.rs index 41178d7..02c3274 100644 --- a/src/events/presence_update.rs +++ b/src/events/presence_update.rs @@ -1,4 +1,8 @@ -use poise::serenity_prelude::{Context, Presence}; +use std::time::Duration; + +use poise::serenity_prelude::{prelude::TypeMapKey, Context, Member, Presence, UserId}; +use tokio::time::sleep; +use tokio_util::sync::CancellationToken; use crate::{ env::Environment, @@ -7,10 +11,24 @@ use crate::{ util::{lunaro_tracking, play}, }; +struct ScheduledReadyRemovals { + pub scheduled: Vec<(UserId, CancellationToken)>, +} + +impl Default for ScheduledReadyRemovals { + fn default() -> Self { + Self { + scheduled: Vec::new(), + } + } +} + /// Handles the presence update event. -pub async fn handle(context: Context, presence: &Presence) -> Result<(), Error> { +pub async fn handle(context: &Context, presence: &Presence) -> Result<(), Error> { + log::debug!("Presence update"); + let env = Environment::instance(); - let tracking_config = lunaro_tracking::Config::instance().await; + let tracking_config = lunaro_tracking::LunaroTrackingConfig::instance().await; if tracking_config.is_blocked(&presence.user.id) { return Ok(()); @@ -21,18 +39,89 @@ pub async fn handle(context: Context, presence: &Presence) -> Result<(), Error> .get_member(presence.guild_id.unwrap(), presence.user.id) .await?; - let is_playing = member + let has_playing_role = member .user .has_role(&context, env.home_guild_id, env.playing_role_id) .await?; - let is_playing_lunaro = lunaro_tracking::is_playing_lunaro(presence).is_ok_and(|value| value); + let is_in_game = lunaro_tracking::is_playing_lunaro(presence).is_ok_and(|value| value); - if is_playing_lunaro && !is_playing { - play::add(member, &context).await?; - } else if !is_playing_lunaro && is_playing { - play::remove(member, &context).await?; + if is_in_game && !has_playing_role { + add_ready_role(member, &context).await?; + } else if !is_in_game && has_playing_role { + schedule_ready_removal(member.clone(), &context).await; } Ok(()) } + +async fn add_ready_role(member: &mut Member, context: &Context) -> Result<(), Error> { + let cancellation_schedules = context + .data + .read() + .await + .get::() + .unwrap_or(&ScheduledReadyRemovals::default()) + .scheduled + .clone(); + + if !cancellation_schedules.is_empty() { + let cancellation_token = cancellation_schedules + .iter() + .find(|(user_id, _)| user_id == &member.user.id) + .map(|(_, token)| token.clone()); + + if let Some(token) = cancellation_token { + token.cancel(); + } + + let uncancelled_tokens = cancellation_schedules + .iter() + .filter(|(_, token)| !token.is_cancelled()) + .map(|(user_id, token)| (*user_id, token.clone())) + .collect::>(); + + context + .data + .write() + .await + .get_mut::() + .unwrap() + .scheduled = uncancelled_tokens; + } + + play::add(member, context).await +} + +async fn schedule_ready_removal(mut member: Member, context: &Context) { + let cancellation_token = CancellationToken::new(); + let entry = (member.user.id, cancellation_token.clone()); + + context + .data + .write() + .await + .get_mut::() + .unwrap() + .scheduled + .push(entry); + + let context = context.clone(); + + tokio::spawn(async move { + tokio::select! { + _ = cancellation_token.cancelled() => { + log::info!("Cancelling ready removal for {} ({})", member.user.tag(), member.user.id); + } + _ = sleep(Duration::from_secs(5)) => { + play::remove(&mut member, &context).await.unwrap_or_else(|error| { + log::error!("Failed to remove ready role for {} ({}): {}", member.user.tag(), member.user.id, error); + }); + } + } + }); +} + +impl TypeMapKey for ScheduledReadyRemovals { + type Value = Self; +} diff --git a/src/main.rs b/src/main.rs index 3318728..9c30bcf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -42,6 +42,9 @@ async fn main() { let framework_options = FrameworkOptions { pre_command: |context| Box::pin(log_invocation(context)), on_error: |error| Box::pin(on_framework_error(error)), + event_handler: |context, event, framework, data| { + Box::pin(events::event_handler(context, event, framework, data)) + }, commands: vec![ @@ -87,7 +90,9 @@ async fn main() { /// Logs command invocations. async fn log_invocation(context: PoiseContext<'_>) { let author = &context.author().tag(); + let author_id = &context.author().id; let guild = &context.partial_guild().await.unwrap().name; + let guild_id = &context.guild_id().unwrap(); let parent_commands = context .parent_commands() .iter() @@ -102,7 +107,7 @@ async fn log_invocation(context: PoiseContext<'_>) { format!("{} {}", parent_commands, context.command().name) }; - log::info!("{author} ran [{command}] in {guild}"); + log::info!("{author} ({author_id}) ran [{command}] in {guild} ({guild_id})"); } /// Handles framework errors according to their severity. @@ -125,12 +130,14 @@ async fn on_framework_error(error: FrameworkError<'_, PoiseData, Error>) { /// Logs the error and notifies the affected guild. async fn log_error(message: &str, context: PoiseContext<'_>) { let user = &context.author().tag(); + let user_id = &context.author().id; let command = &context.command().name; let guild = &context.partial_guild().await.unwrap().name; + let guild_id = &context.guild_id().unwrap(); let trace_id = Uuid::new_v4(); - log::error!("{user} ran [{command}] in {guild} and got an error: {message} ({trace_id})",); + log::error!("{user} ({user_id}) ran [{command}] in {guild} ({guild_id}) and got an error: {message} ({trace_id})",); send_error_to_chat(message, trace_id, context).await; } diff --git a/src/util/lunaro_tracking.rs b/src/util/lunaro_tracking.rs index 89d68d7..356ce39 100644 --- a/src/util/lunaro_tracking.rs +++ b/src/util/lunaro_tracking.rs @@ -9,17 +9,17 @@ use crate::{errors::data::DataError, traits::config_file::ConfigFile, types::err use super::data; -static CONFIG: OnceLock> = OnceLock::new(); +static CONFIG: OnceLock> = OnceLock::new(); const CONFIG_FILE: &str = "lunaro_tracking.json"; /// Configures the Lunaro tracker's behaviour. #[derive(Serialize, Deserialize, Default, Debug)] -pub struct Config { +pub struct LunaroTrackingConfig { /// List of user IDs to ignore Lunaro updates from. pub blocklist: HashSet, } -impl Config { +impl LunaroTrackingConfig { /// Check if a user is on the tracking blocklist. pub fn is_blocked(&self, user_id: &UserId) -> bool { self.blocklist.contains(user_id) @@ -39,13 +39,13 @@ impl Config { } #[async_trait] -impl ConfigFile for Config { +impl ConfigFile for LunaroTrackingConfig { fn load() -> Result, Error> { match data::read_config(CONFIG_FILE) { Ok(config) => Ok(config), Err(error) => match error.downcast_ref::() { Some(DataError::MissingConfigFile(_)) => { - let config = Config::default(); + let config = LunaroTrackingConfig::default(); config.save()?; Ok(config.into()) } @@ -60,7 +60,7 @@ impl ConfigFile for Config { async fn instance<'a>() -> MutexGuard<'a, Self> { CONFIG - .get_or_init(|| Mutex::new(*Config::load().unwrap_or_default())) + .get_or_init(|| Mutex::new(*LunaroTrackingConfig::load().unwrap_or_default())) .lock() .await } @@ -68,7 +68,7 @@ impl ConfigFile for Config { /// Remove a user from the tracking blocklist, if present. pub async fn allow_for(user: &User) -> Result<(), Error> { - let mut config = Config::instance().await; + let mut config = LunaroTrackingConfig::instance().await; config.remove_from_blocklist(&user.id)?; @@ -83,7 +83,7 @@ pub async fn allow_for(user: &User) -> Result<(), Error> { /// Add a user to the tracking blocklist. pub async fn deny_for(user: &User) -> Result<(), Error> { - let mut config = Config::instance().await; + let mut config = LunaroTrackingConfig::instance().await; config.add_to_blocklist(&user.id)?; diff --git a/src/util/play.rs b/src/util/play.rs index 85d0851..17324bb 100644 --- a/src/util/play.rs +++ b/src/util/play.rs @@ -30,7 +30,11 @@ pub async fn add(member: &mut Member, context: &Context) -> Result<(), Error> { .add_role(context, Environment::instance().playing_role_id) .await?; - log::debug!("Added {} to playing role", member.user.tag()); + log::debug!( + "Added {} ({}) to playing role", + member.user.tag(), + member.user.id + ); Ok(()) } @@ -41,7 +45,11 @@ pub async fn remove(member: &mut Member, context: &Context) -> Result<(), Error> .remove_role(context, Environment::instance().playing_role_id) .await?; - log::debug!("Removed {} from playing role", member.user.tag()); + log::debug!( + "Removed {} ({}) from playing role", + member.user.tag(), + member.user.id + ); Ok(()) }