Refactor ratelimiting

hypergonial · Aug 31, 2023 · 2195a67 · 2195a67
1 parent 466e82c
commit 2195a67
Show file tree

Hide file tree

Showing 4 changed files with 193 additions and 97 deletions.
diff --git a/extensions/automod.py b/extensions/automod.py
@@ -16,7 +16,7 @@
 from models.events import AutoModMessageFlagEvent
 from models.plugin import SnedPlugin
 from utils import helpers
-from utils.ratelimiter import BucketType
+from utils.ratelimiter import MemberBucket
 
 INVITE_REGEX = re.compile(r"(?:https?://)?discord(?:app)?\.(?:com/invite|gg)/[a-zA-Z0-9]+/?")
 """Used to detect and handle Discord invites."""
@@ -25,12 +25,12 @@
 DISCORD_FORMATTING_REGEX = re.compile(r"<\S+>")
 """Remove Discord-specific formatting. Performance is key so some false-positives are acceptable here."""
 
-SPAM_RATELIMITER = utils.RateLimiter(10, 8, bucket=BucketType.MEMBER, wait=False)
-PUNISH_RATELIMITER = utils.RateLimiter(30, 1, bucket=BucketType.MEMBER, wait=False)
-ATTACH_SPAM_RATELIMITER = utils.RateLimiter(30, 2, bucket=BucketType.MEMBER, wait=False)
-LINK_SPAM_RATELIMITER = utils.RateLimiter(30, 2, bucket=BucketType.MEMBER, wait=False)
-ESCALATE_PREWARN_RATELIMITER = utils.RateLimiter(30, 1, bucket=BucketType.MEMBER, wait=False)
-ESCALATE_RATELIMITER = utils.RateLimiter(30, 1, bucket=BucketType.MEMBER, wait=False)
+SPAM_RATELIMITER = utils.RateLimiter(10, 8, bucket=MemberBucket, wait=False)
+PUNISH_RATELIMITER = utils.RateLimiter(30, 1, bucket=MemberBucket, wait=False)
+ATTACH_SPAM_RATELIMITER = utils.RateLimiter(30, 2, bucket=MemberBucket, wait=False)
+LINK_SPAM_RATELIMITER = utils.RateLimiter(30, 2, bucket=MemberBucket, wait=False)
+ESCALATE_PREWARN_RATELIMITER = utils.RateLimiter(30, 1, bucket=MemberBucket, wait=False)
+ESCALATE_RATELIMITER = utils.RateLimiter(30, 1, bucket=MemberBucket, wait=False)
 
 logger = logging.getLogger(__name__)
 

diff --git a/extensions/fun.py b/extensions/fun.py
@@ -23,8 +23,9 @@
 from models.context import SnedContext, SnedUserContext
 from models.plugin import SnedPlugin
 from models.views import AuthorOnlyNavigator, AuthorOnlyView
-from utils import BucketType, RateLimiter, helpers
+from utils import GlobalBucket, RateLimiter, helpers
 from utils.dictionaryapi import DictionaryClient, DictionaryEntry, DictionaryException, UrbanEntry
+from utils.ratelimiter import UserBucket
 from utils.rpn import InvalidExpressionError, Solver
 
 ANIMAL_EMOJI_MAPPING: dict[str, str] = {
@@ -37,7 +38,10 @@
     "racoon": "🦝",
 }
 
-animal_ratelimiter = RateLimiter(60, 45, BucketType.GLOBAL, wait=False)
+ANIMAL_RATELIMITER = RateLimiter(60, 45, GlobalBucket, wait=False)
+COMF_LIMITER = RateLimiter(60, 5, UserBucket, wait=False)
+VESZTETTEM_LIMITER = RateLimiter(1800, 1, GlobalBucket, wait=False)
+COMF_PROGRESS_BAR_WIDTH = 20
 
 logger = logging.getLogger(__name__)
 
@@ -903,8 +907,8 @@ async def on_dice_reroll(event: miru.ComponentInteractionCreateEvent) -> None:
 @lightbulb.command("animal", "Shows a random picture of the selected animal.", pass_options=True)
 @lightbulb.implements(lightbulb.SlashCommand)
 async def animal(ctx: SnedSlashContext, animal: str) -> None:
-    await animal_ratelimiter.acquire(ctx)
-    if animal_ratelimiter.is_rate_limited(ctx):
+    await ANIMAL_RATELIMITER.acquire(ctx)
+    if ANIMAL_RATELIMITER.is_rate_limited(ctx):
         await ctx.respond(
             embed=hikari.Embed(
                 title="❌ Ratelimited",
@@ -983,36 +987,29 @@ async def wiki(ctx: SnedSlashContext, query: str) -> None:
         await ctx.respond(embed=embed)
 
 
-vesztettem_limiter = RateLimiter(1800, 1, BucketType.GLOBAL, wait=False)
-
-
 @fun.listener(hikari.GuildMessageCreateEvent)
 async def lose_autoresponse(event: hikari.GuildMessageCreateEvent) -> None:
     if event.guild_id not in (Config().DEBUG_GUILDS or (1012448659029381190,)) or not event.is_human:
         return
 
     if event.content and "vesztettem" in event.content.lower():
-        await vesztettem_limiter.acquire(event.message)
+        await VESZTETTEM_LIMITER.acquire(event.message)
 
-        if vesztettem_limiter.is_rate_limited(event.message):
+        if VESZTETTEM_LIMITER.is_rate_limited(event.message):
             return
 
         await event.message.respond("Vesztettem")
 
 
-comf_ratelimiter = RateLimiter(60, 5, BucketType.USER, wait=False)
-COMF_PROGRESS_BAR_WIDTH = 20
-
-
 @fun.command
 @lightbulb.app_command_permissions(None, dm_enabled=False)
 @lightbulb.command("comf", "Shows your current and upcoming comfiness.")
 @lightbulb.implements(lightbulb.SlashCommand)
 async def comf(ctx: SnedSlashContext) -> None:
     assert ctx.member is not None
 
-    await comf_ratelimiter.acquire(ctx)
-    if comf_ratelimiter.is_rate_limited(ctx):
+    await COMF_LIMITER.acquire(ctx)
+    if COMF_LIMITER.is_rate_limited(ctx):
         await ctx.respond(
             embed=hikari.Embed(
                 title="❌ Ratelimited",

diff --git a/extensions/role_buttons.py b/extensions/role_buttons.py
@@ -11,7 +11,7 @@
 from models.plugin import SnedPlugin
 from models.rolebutton import RoleButton, RoleButtonMode
 from utils import helpers
-from utils.ratelimiter import BucketType, RateLimiter
+from utils.ratelimiter import MemberBucket, RateLimiter
 
 logger = logging.getLogger(__name__)
 
@@ -29,7 +29,7 @@
     "Remove": RoleButtonMode.REMOVE_ONLY,
 }
 
-role_button_ratelimiter = RateLimiter(2, 1, BucketType.MEMBER, wait=False)
+role_button_ratelimiter = RateLimiter(2, 1, MemberBucket, wait=False)
 
 
 class RoleButtonConfirmType(enum.Enum):