diff --git a/Cargo.toml b/Cargo.toml index 6666edc8..40d0a59b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,3 +31,4 @@ humantime = "2.1.0" reqwest = "0.11.12" thiserror = "1.0.37" sentry = "0.29.0" +floodgate = "0.4.0" diff --git a/src/client/cooldowns.rs b/src/client/cooldowns.rs index 5e456166..945c8422 100644 --- a/src/client/cooldowns.rs +++ b/src/client/cooldowns.rs @@ -1,20 +1,24 @@ +use std::sync::Arc; + +use floodgate::FixedMapping; use twilight_model::id::{marker::ChannelMarker, Id}; -use crate::{constants, utils::cooldowns::FixedMapping}; +use crate::constants; pub struct Cooldowns { // restricts per-channel - pub autostar_send: FixedMapping>, + pub autostar_send: Arc>>, } impl Cooldowns { pub fn new() -> Self { - Self { - autostar_send: FixedMapping::new( - constants::AUTOSTAR_COOLDOWN.0, - constants::AUTOSTAR_COOLDOWN.1, - ), - } + let autostar_send = Arc::new(FixedMapping::new( + constants::AUTOSTAR_COOLDOWN.0, + constants::AUTOSTAR_COOLDOWN.1, + )); + FixedMapping::start(autostar_send.clone(), None); + + Self { autostar_send } } } diff --git a/src/constants.rs b/src/constants.rs index b5351147..ef8706c6 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -11,7 +11,7 @@ pub const MAX_MESSAGES: u32 = 10_000; pub const MAX_NAMES: u32 = 100; // Cooldowns -pub const AUTOSTAR_COOLDOWN: (u32, Duration) = (5, Duration::from_secs(20)); +pub const AUTOSTAR_COOLDOWN: (u64, Duration) = (5, Duration::from_secs(20)); // Common Validation pub const MAX_NAME_LENGTH: u32 = 32; diff --git a/src/core/autostar.rs b/src/core/autostar.rs index aa43930c..b67ffa51 100644 --- a/src/core/autostar.rs +++ b/src/core/autostar.rs @@ -28,9 +28,8 @@ pub async fn handle(bot: &StarboardBot, event: &MessageCreate) -> StarboardResul if bot .cooldowns .autostar_send - .trigger(event.channel_id) - .await - .is_none() + .trigger(&event.channel_id) + .is_some() { return Ok(()); } diff --git a/src/utils/cooldowns.rs b/src/utils/cooldowns.rs deleted file mode 100644 index d6dea2c2..00000000 --- a/src/utils/cooldowns.rs +++ /dev/null @@ -1,118 +0,0 @@ -use std::{ - collections::HashMap, - hash::Hash, - time::{Duration, Instant}, -}; - -use tokio::sync::RwLock; - -pub struct FixedMapping { - mapping: FlexibleMapping, - period: Duration, - capacity: u32, -} - -impl FixedMapping -where - K: Hash + Eq + Clone, -{ - pub fn new(capacity: u32, period: Duration) -> Self { - Self { - mapping: FlexibleMapping::new(), - period, - capacity, - } - } - - pub async fn trigger(&self, key: K) -> Option<()> { - self.mapping.trigger(key, self.capacity, self.period).await - } -} - -pub struct FlexibleMapping { - windows: RwLock>, -} - -impl FlexibleMapping -where - K: Hash + Eq + Clone, -{ - pub fn new() -> Self { - Self { - windows: RwLock::new(HashMap::new()), - } - } - - pub async fn trigger(&self, key: K, capacity: u32, period: Duration) -> Option<()> { - let mut windows = self.windows.write().await; - - let mut to_remove = Vec::new(); - for window in windows.iter_mut() { - if window.1.refresh() { - to_remove.push(window.0.clone()); - } - } - for key in to_remove { - windows.remove(&key); - } - - let window = match windows.get_mut(&key) { - None => { - windows.insert(key.clone(), JumpingWindow::new(capacity, period)); - windows.get_mut(&key).unwrap() - } - Some(window) => window, - }; - - window.trigger() - } -} - -impl Default for FlexibleMapping -where - K: Hash + Eq + Clone, -{ - fn default() -> Self { - Self::new() - } -} - -pub struct JumpingWindow { - last_reset: Instant, - period: Duration, - capacity: u32, - tokens: u32, -} - -impl JumpingWindow { - pub fn new(capacity: u32, period: Duration) -> Self { - Self { - last_reset: Instant::now(), - period, - capacity, - tokens: capacity, - } - } - - pub fn trigger(&mut self) -> Option<()> { - self.refresh(); - - if self.tokens == 0 { - return None; - } - - self.tokens -= 1; - - Some(()) - } - - fn refresh(&mut self) -> bool { - if self.last_reset.elapsed() >= self.period { - self.tokens = self.capacity; - self.last_reset = Instant::now(); - true - } else { - false - } - } -} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 5809e922..020f00e2 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,4 @@ pub mod async_dash; -pub mod cooldowns; pub mod dashset_lock; pub mod dm; pub mod embed;