Skip to content

Commit

Permalink
Add a fixed-window cooldown algorithm implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
tandemdude committed Dec 9, 2022
1 parent 68a9b2b commit e3e653b
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 1 deletion.
1 change: 1 addition & 0 deletions lightbulb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
"ExtensionMissingUnload",
"ExtensionNotFound",
"ExtensionNotLoaded",
"FixedWindowCooldownAlgorithm",
"GlobalBucket",
"GuildBucket",
"GuildCategoryConverter",
Expand Down
33 changes: 32 additions & 1 deletion lightbulb/cooldown_algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,13 @@
# along with Lightbulb. If not, see <https://www.gnu.org/licenses/>.
from __future__ import annotations

__all__ = ["CooldownStatus", "CooldownAlgorithm", "BangBangCooldownAlgorithm", "SlidingWindowCooldownAlgorithm"]
__all__ = [
"CooldownStatus",
"CooldownAlgorithm",
"BangBangCooldownAlgorithm",
"FixedWindowCooldownAlgorithm",
"SlidingWindowCooldownAlgorithm",
]

import abc
import enum
Expand Down Expand Up @@ -77,6 +83,31 @@ def evaluate(self, bucket: buckets.Bucket) -> CooldownStatus:
return CooldownStatus.INACTIVE


class FixedWindowCooldownAlgorithm(CooldownAlgorithm):
"""
Cooldown algorithm that allows ``n`` command invocations within a fixed window time period. I.e. 5 invocations
within a 30-second time period, allowing additional invocations once the entire period has elapsed.
"""

__slots__ = ("_commands_run", "_expires")

def __init__(self) -> None:
self._commands_run: int = 0
self._expires: t.Optional[float] = None

def evaluate(self, bucket: buckets.Bucket) -> CooldownStatus:
self._expires = self._expires or (time.perf_counter() + bucket.length)
if self._expires > time.perf_counter():
self._commands_run = 0

self._commands_run += 1
if self._commands_run > bucket.usages:
raise errors.CommandIsOnCooldown(
"This command is on cooldown", retry_after=(self._expires - time.perf_counter())
)
return CooldownStatus.INACTIVE


class SlidingWindowCooldownAlgorithm(CooldownAlgorithm):
"""
Cooldown algorithm that allows ``n`` command invocations within a sliding window time period. I.e. 5 invocations
Expand Down

0 comments on commit e3e653b

Please sign in to comment.