diff --git a/src/client/cooldowns.rs b/src/client/cooldowns.rs index 57feeee0..f0214155 100644 --- a/src/client/cooldowns.rs +++ b/src/client/cooldowns.rs @@ -1,13 +1,6 @@ -use std::{ - collections::HashMap, - hash::Hash, - time::{Duration, Instant}, -}; - -use tokio::sync::RwLock; use twilight_model::id::{marker::ChannelMarker, Id}; -use crate::constants; +use crate::{constants, utils::cooldowns::FixedMapping}; pub struct Cooldowns { // restricts per-channel @@ -24,105 +17,3 @@ impl Cooldowns { } } } - -pub struct FixedMapping { - mapping: FlexibleMapping, - period: Duration, - capacity: u32, -} - -impl FixedMapping -where - K: Hash + Eq + Clone, -{ - 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, -{ - 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() - } -} - -pub struct JumpingWindow { - last_reset: Instant, - period: Duration, - capacity: u32, - tokens: u32, -} - -impl JumpingWindow { - 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/cooldowns.rs b/src/utils/cooldowns.rs new file mode 100644 index 00000000..18fd3657 --- /dev/null +++ b/src/utils/cooldowns.rs @@ -0,0 +1,109 @@ +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() + } +} + +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 71b6d90c..20ce39c3 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,4 +1,5 @@ pub mod async_dash; +pub mod cooldowns; pub mod dashset_lock; pub mod dm; pub mod embed;