[P1] Scale shard processing and bound hot-path concurrency

README.md

@@ -116,6 +116,10 @@ To use it, uncomment the proper lines in your docker-compose:
 
 Then, simply run `docker compose up -d`.
 
+For large installations, run multiple bot instances with the same `SHARD_COUNT` and assign each instance a
+non-overlapping comma-separated `SHARD_IDS` set. Use the shard count recommended by Discord's Gateway Bot endpoint.
+Each shard ID must be owned by exactly one running instance.
+
 #### Available environment variables
 
 | Environment variable | Comment                                                                         |

cogs/link_fix.py

@@ -7,6 +7,8 @@
 import discord_markdown_ast_parser as dmap
 from discord_markdown_ast_parser.parser import NodeType
 import logging
+import os
+import time
 
 from database.models.Member import *
 from database.models.Role import Role
@@ -15,12 +17,26 @@
 from database.models.Event import *
 from src.websites import *
 from src.utils import *
+from src.hot_path_cache import MISSING, filter_cache, guild_cache, webhook_cache
+from src.runtime import RuntimeBusyError, run_database
 
 import discore
 
 __all__ = ('LinkFix',)
 
 _logger = logging.getLogger(__name__)
+_fix_concurrency = asyncio.Semaphore(max(1, int(os.getenv('FIX_CONCURRENCY', '200'))))
+_last_capacity_warning = 0.0
+
+
+def _warn_capacity(message: str) -> None:
+    """Avoid turning overload protection into a log storm."""
+
+    global _last_capacity_warning
+    now = time.monotonic()
+    if now - _last_capacity_warning >= 30:
+        _last_capacity_warning = now
+        _logger.warning(message)
 
 
 def get_website(guild: Guild, url: str, spoiler: bool = False) -> WebsiteLink | None:
@@ -238,24 +254,34 @@ async def get_or_create_webhook(channel: GuildMessageableChannel) -> discore.Web
     if webhook_channel is None:
         return None
 
+    cached = webhook_cache.get(webhook_channel.id)
+    if cached is not MISSING:
+        return cached
+
     if not hasattr(webhook_channel, 'webhooks') or not hasattr(webhook_channel, 'create_webhook'):
+        webhook_cache.set(webhook_channel.id, None, ttl=300)
         return None
 
     if not webhook_channel.permissions_for(channel.guild.me).manage_webhooks:
+        webhook_cache.set(webhook_channel.id, None, ttl=300)
         return None
 
     success, webhooks = await safe_send_coro(webhook_channel.webhooks(), forbidden=True)
     if not success:
+        webhook_cache.set(webhook_channel.id, None, ttl=60)
         return None
     bot = discore.Bot.get()
     webhook = next((
         w for w in webhooks
         if getattr(w.user, 'id', None) == bot.user.id
     ), None)
     if webhook is not None:
+        webhook_cache.set(webhook_channel.id, webhook)
         return webhook
     success, webhook = await safe_send_coro(webhook_channel.create_webhook(name=bot.user.display_name), forbidden=True)
-    return webhook if success else None
+    webhook = webhook if success else None
+    webhook_cache.set(webhook_channel.id, webhook, ttl=None if webhook else 60)
+    return webhook
 
 
 async def webhook_send(
@@ -340,36 +366,80 @@ async def on_message(self, message: discore.Message) -> None:
         entrypoint_context.set(f"event on_message {{message={message!r}}}")
 
         if (
-                message.author == message.guild.me
+                not message.guild
+                or message.author == message.guild.me
                 or not message.content
                 or not message.channel
-                or not message.guild
                 or message.is_system()
         ):
             return
 
+        if 'http://' not in message.content.lower() and 'https://' not in message.content.lower():
+            return
+
         urls = get_embeddable_urls(dmap.parse(message.content))
 
         if not urls:
             return
 
-        guild = Guild.find_or_create(message.guild)
-        links = filter_fixable_links(urls, guild)
-
-        if not links:
+        def resolve_context() -> tuple[Guild | None, list[WebsiteLink]]:
+            guild = guild_cache.get(message.guild.id)
+            if guild is MISSING:
+                guild = Guild.find_or_create(message.guild)
+                guild_cache.set(message.guild.id, guild)
+
+            links = filter_fixable_links(urls, guild)
+            if not links:
+                return None, []
+
+            if any(
+                re.search(rf"\b{re.escape(keyword)}\b", message.content)
+                for keyword in guild.keywords
+            ) != guild.keywords_use_allow_list:
+                return None, []
+
+            role_ids = tuple(sorted(role.id for role in message.author.roles)) \
+                if isinstance(message.author, discore.Member) else ()
+            filter_key = (
+                guild.id,
+                message.channel.id,
+                message.author.id,
+                role_ids,
+                bool(message.webhook_id),
+            )
+            allowed = filter_cache.get(filter_key)
+            if allowed is MISSING:
+                allowed = (
+                    TextChannel.find_get_enabled(message.channel, guild)
+                    and (
+                        not isinstance(message.author, discore.Member)
+                        or (
+                            Member.find_get_enabled(message.author, guild)
+                            and (any if guild.roles_use_any_rule else all)(
+                                Role.finds_get_enabled(message.author.roles, guild)
+                            )
+                        )
+                    )
+                    and (message.webhook_id is None or bool(guild.webhooks))
+                )
+                filter_cache.set(filter_key, allowed)
+            return (guild, links) if allowed else (None, [])
+
+        try:
+            guild, links = await run_database(resolve_context)
+        except RuntimeBusyError:
+            _warn_capacity('Skipping link fixes because database worker capacity is exhausted')
             return
-        if any(
-                re.search(rf"\b{re.escape(k)}\b", message.content) for k in guild.keywords
-        ) != guild.keywords_use_allow_list:
-            return
-        if not TextChannel.find_get_enabled(message.channel, guild):
-            return
-        if isinstance(message.author, discore.Member) and (
-            not Member.find_get_enabled(message.author, guild)
-            or not (any if (guild and guild.roles_use_any_rule) else all)(Role.finds_get_enabled(message.author.roles, guild))
-        ):
-            return
-        if message.webhook_id is not None and not bool(guild.webhooks):
+
+        if not guild or not links:
             return
 
-        await fix_embeds(message, guild, links)
+        try:
+            await asyncio.wait_for(_fix_concurrency.acquire(), timeout=0.5)
+        except asyncio.TimeoutError:
+            _warn_capacity('Skipping link fixes because fix concurrency is exhausted')
+            return
+        try:
+            await fix_embeds(message, guild, links)
+        finally:
+            _fix_concurrency.release()

cogs/setup.py

@@ -3,6 +3,7 @@
 
 from src import utils
 from database.models.Event import *
+from src.runtime import RuntimeBusyError, run_database
 
 import discore
 
@@ -58,7 +59,11 @@ def format_count(n: int) -> str:
                 return f"{n / 1_000:.1f}".rstrip('0').rstrip('.') + 'k'
             return str(n)
 
-        fixed_links_nb = len(Event.since('fixed_link', days=1))
+        try:
+            fixed_links_nb = await run_database(Event.count_since, 'fixed_link', days=1)
+        except RuntimeBusyError:
+            _logger.warning('[ACTIVITY] Skipping update because database workers are busy')
+            return
         if fixed_links_nb == 0:
             return
 

database/models/Event.py

@@ -4,11 +4,18 @@
 from typing import Self
 import datetime as dt
 import json
+import logging
+import os
 
 from masoniteorm.models import Model
 
 import discore
 
+
+_logger = logging.getLogger(__name__)
+_buffer_limit = max(1000, int(os.getenv('ANALYTICS_BUFFER_LIMIT', '100000')))
+_flush_batch_size = max(100, int(os.getenv('ANALYTICS_FLUSH_BATCH_SIZE', '5000')))
+
 class Event(Model):
     """Event Model"""
 
@@ -43,15 +50,35 @@ def since(cls, event_name: str | None = None, days: int = 0, hours: int = 0, min
 
         return query.get()
 
+    @classmethod
+    def count_since(cls, event_name: str | None = None, **delta_kwargs) -> int:
+        """Count matching events in the database without materialising model objects."""
+
+        if not discore.config.analytic:
+            return 0
+        query = cls.where('name', event_name) if event_name else cls
+        delta = dt.timedelta(**delta_kwargs)
+        if delta:
+            query = query.where('created_at', '>=', dt.datetime.now() - delta)
+        return query.count()
+
     @classmethod
     async def _flush_loop(cls) -> None:
         """Flush the buffer every 5 seconds"""
         while True:
             await asyncio.sleep(5)
             async with cls._lock:
                 if cls._buffer:
-                    cls.bulk_create(cls._buffer)
-                    cls._buffer.clear()
+                    events = cls._buffer[:_flush_batch_size]
+                    del cls._buffer[:_flush_batch_size]
+                else:
+                    continue
+            try:
+                await asyncio.to_thread(cls.bulk_create, events)
+            except Exception:
+                _logger.exception('Failed to flush analytics events')
+                async with cls._lock:
+                    cls._buffer = (events + cls._buffer)[-_buffer_limit:]
 
     @classmethod
     async def buff_cr(cls, *events: dict) -> None:
@@ -74,6 +101,8 @@ async def buff_cr(cls, *events: dict) -> None:
             events[i] = event
         async with cls._lock:
             cls._buffer.extend(events)
+            if len(cls._buffer) > _buffer_limit:
+                del cls._buffer[:-_buffer_limit]
 
         if cls._flush_task is None or cls._flush_task.done():
             cls._flush_task = asyncio.create_task(cls._flush_loop())

docker-compose.example.yml

@@ -19,6 +19,16 @@ services:
       - DISCORD_TOKEN=your_discord_bot_token
       # Optional: Server ID that you want to use for controlling your bot's instance.
       # - DEV_GUILD=your_discord_dev_guild_id
+      # Optional: split Discord shards between multiple bot instances.
+      # SHARD_COUNT must be the same for every instance; SHARD_IDS is comma-separated.
+      # - SHARD_COUNT=50
+      # - SHARD_IDS=0,1,2,3,4
+      # Optional bounded-concurrency tuning.
+      # - DATABASE_WORKER_CONCURRENCY=8
+      # - FIX_CONCURRENCY=200
+      # - EMBEDEZ_CONCURRENCY=50
+      # - ANALYTICS_BUFFER_LIMIT=100000
+      # - ANALYTICS_FLUSH_BATCH_SIZE=5000
     # uncomment and create file if you want to override any default settings
     # volumes:
     #    - ./override.config.yml:/usr/local/app/override.config.yml:ro

main.py

@@ -2,12 +2,13 @@
 import os
 import discore
 
+from src.sharding import shard_options
 from src.utils import I18nTranslator
 
 os.environ['DB_CONFIG_PATH'] = 'database/config.py'
 
 intents = discore.Intents(guild_messages=True, message_content=True, guilds=True)
 
-bot = discore.Bot(help_command=None, intents=intents)
+bot = discore.Bot(help_command=None, intents=intents, **shard_options())
 asyncio.run(bot.tree.set_translator(I18nTranslator()))
 bot.run()

src/hot_path_cache.py

@@ -0,0 +1,76 @@
+"""Small process-local caches for frequently read guild configuration."""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+from collections import OrderedDict
+from typing import Generic, TypeVar
+
+
+K = TypeVar('K')
+V = TypeVar('V')
+MISSING = object()
+
+
+class TTLCache(Generic[K, V]):
+    """A thread-safe, size-bounded TTL cache."""
+
+    def __init__(self, max_size: int, ttl: float):
+        self.max_size = max(1, max_size)
+        self.ttl = max(0.1, ttl)
+        self._items: OrderedDict[K, tuple[float, V]] = OrderedDict()
+        self._lock = threading.Lock()
+
+    def get(self, key: K, default=MISSING):
+        now = time.monotonic()
+        with self._lock:
+            item = self._items.get(key)
+            if item is None:
+                return default
+            expires_at, value = item
+            if expires_at <= now:
+                del self._items[key]
+                return default
+            self._items.move_to_end(key)
+            return value
+
+    def set(self, key: K, value: V, ttl: float | None = None) -> None:
+        expires_at = time.monotonic() + (self.ttl if ttl is None else max(0.1, ttl))
+        with self._lock:
+            self._items[key] = (expires_at, value)
+            self._items.move_to_end(key)
+            while len(self._items) > self.max_size:
+                self._items.popitem(last=False)
+
+    def pop(self, key: K) -> None:
+        with self._lock:
+            self._items.pop(key, None)
+
+    def pop_guild(self, guild_id: int) -> None:
+        with self._lock:
+            for key in tuple(self._items):
+                if isinstance(key, tuple) and key and key[0] == guild_id:
+                    del self._items[key]
+
+
+guild_cache: TTLCache[int, object] = TTLCache(
+    max_size=int(os.getenv('GUILD_CACHE_SIZE', '75000')),
+    ttl=float(os.getenv('GUILD_CACHE_TTL', '300')),
+)
+filter_cache: TTLCache[tuple, bool] = TTLCache(
+    max_size=int(os.getenv('FILTER_CACHE_SIZE', '200000')),
+    ttl=float(os.getenv('FILTER_CACHE_TTL', '60')),
+)
+webhook_cache: TTLCache[int, object | None] = TTLCache(
+    max_size=int(os.getenv('WEBHOOK_CACHE_SIZE', '100000')),
+    ttl=float(os.getenv('WEBHOOK_CACHE_TTL', '3600')),
+)
+
+
+def invalidate_guild(guild_id: int) -> None:
+    """Invalidate locally cached settings after a settings interaction."""
+
+    guild_cache.pop(guild_id)
+    filter_cache.pop_guild(guild_id)