rootspring / cache.rs

use super::core::{Limit, LimitTypes}; use governor::{DefaultKeyedRateLimiter, Quota}; use moka::future::Cache; use serenity::all::{GuildId, UserId}; use std::collections::HashMap; use std::num::NonZeroU32; use std::sync::{Arc, LazyLock}; // Hashmap of limit types to a hashmap of limit ids to its ratelimiter pub type RatelimiterMap<RlKey> = HashMap<LimitTypes, HashMap<String, DefaultKeyedRateLimiter<RlKey>>>; #[derive(Debug)] pub struct GuildLimitsCache { pub global: RatelimiterMap<()>, pub per_user: RatelimiterMap<UserId>, } impl GuildLimitsCache { /// Attempts to limit a user, returning a tuple of whether the user is allowed to continue, the time at which the bucket will be replenished, limit id that was hit pub async fn limit( &self, user_id: UserId, limit_type: LimitTypes, ) -> (bool, Option<governor::clock::QuantaInstant>, Option<String>) { if let Some(limits) = self.per_user.get(&limit_type) { for (limit_id, lim) in limits.iter() { match lim.check_key(&user_id) { Ok(()) => continue, // TODO: Return the time at which the bucket will be replenished Err(wait) => { return ( false, Some(wait.earliest_possible()), Some(limit_id.clone()), ) } } } } (true, None, None) } } static GUILD_LIMITS: LazyLock<Cache<GuildId, Arc<GuildLimitsCache>>> = LazyLock::new(|| Cache::builder().support_invalidation_closures().build()); pub async fn get_limits( data: &silverpelt::data::Data, guild_id: GuildId, ) -> Result<Arc<GuildLimitsCache>, silverpelt::Error> { if let Some(limits) = GUILD_LIMITS.get(&guild_id).await { Ok(limits.clone()) } else { let mut limits = GuildLimitsCache { global: HashMap::new(), per_user: HashMap::new(), }; // Init limits map here let limits_db = Limit::guild(&data.pool, guild_id).await?; for limit in limits_db { let limit_per = NonZeroU32::new(limit.limit_per as u32).ok_or("Invalid limit_per")?; let quota = Quota::with_period(std::time::Duration::from_secs(limit.limit_time as u64)) .ok_or("Failed to create quota")? .allow_burst(limit_per); let lim = DefaultKeyedRateLimiter::keyed(quota); // TODO: Support global limits limits .per_user .entry(limit.limit_type) .or_default() .insert(limit.limit_id.clone(), lim); } let limits = Arc::new(limits); GUILD_LIMITS.insert(guild_id, limits.clone()).await; Ok(limits) } }