diff options
| -rw-r--r-- | discord/ext/commands/cooldowns.py | 12 | ||||
| -rw-r--r-- | discord/ext/commands/core.py | 2 |
2 files changed, 9 insertions, 5 deletions
diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index c7a200c9..873c832c 100644 --- a/discord/ext/commands/cooldowns.py +++ b/discord/ext/commands/cooldowns.py @@ -128,20 +128,20 @@ class CooldownMapping: elif bucket_type is BucketType.category: return (msg.channel.category or msg.channel).id - def _verify_cache_integrity(self): + def _verify_cache_integrity(self, current=None): # we want to delete all cache objects that haven't been used # in a cooldown window. e.g. if we have a command that has a # cooldown of 60s and it has not been used in 60s then that key should be deleted - current = time.time() + current = current or time.time() dead_keys = [k for k, v in self._cache.items() if current > v._last + v.per] for k in dead_keys: del self._cache[k] - def get_bucket(self, message): + def get_bucket(self, message, current=None): if self._cooldown.type is BucketType.default: return self._cooldown - self._verify_cache_integrity() + self._verify_cache_integrity(current) key = self._bucket_key(message) if key not in self._cache: bucket = self._cooldown.copy() @@ -150,3 +150,7 @@ class CooldownMapping: bucket = self._cache[key] return bucket + + def update_rate_limit(self, message, current=None): + bucket = self.get_bucket(message, current) + return bucket.update_rate_limit(current) diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index 551ccfd1..f5890202 100644 --- a/discord/ext/commands/core.py +++ b/discord/ext/commands/core.py @@ -664,8 +664,8 @@ class Command(_BaseCommand): def _prepare_cooldowns(self, ctx): if self._buckets.valid: - bucket = self._buckets.get_bucket(ctx.message) current = ctx.message.created_at.replace(tzinfo=datetime.timezone.utc).timestamp() + bucket = self._buckets.get_bucket(ctx.message, current) retry_after = bucket.update_rate_limit(current) if retry_after: raise CommandOnCooldown(bucket, retry_after) |