Skip to content

Commit

Permalink
put cooldown impl inside utils (#65)
Browse files Browse the repository at this point in the history
  • Loading branch information
circuitsacul authored Sep 6, 2022
1 parent b387b4b commit ccf9694
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 110 deletions.
111 changes: 1 addition & 110 deletions src/client/cooldowns.rs
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -24,105 +17,3 @@ impl Cooldowns {
}
}
}

pub struct FixedMapping<K> {
mapping: FlexibleMapping<K>,
period: Duration,
capacity: u32,
}

impl<K> FixedMapping<K>
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<K> {
windows: RwLock<HashMap<K, JumpingWindow>>,
}

impl<K> FlexibleMapping<K>
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
}
}
}
109 changes: 109 additions & 0 deletions src/utils/cooldowns.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
use std::{
collections::HashMap,
hash::Hash,
time::{Duration, Instant},
};

use tokio::sync::RwLock;

pub struct FixedMapping<K> {
mapping: FlexibleMapping<K>,
period: Duration,
capacity: u32,
}

impl<K> FixedMapping<K>
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<K> {
windows: RwLock<HashMap<K, JumpingWindow>>,
}

impl<K> FlexibleMapping<K>
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
}
}
}
1 change: 1 addition & 0 deletions src/utils/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
pub mod async_dash;
pub mod cooldowns;
pub mod dashset_lock;
pub mod dm;
pub mod embed;
Expand Down

0 comments on commit ccf9694

Please sign in to comment.