From 332bf050b33e883543361862566b5125ab21b259 Mon Sep 17 00:00:00 2001 From: apple050620312 Date: Thu, 10 Sep 2026 14:00:10 +0800 Subject: [PATCH] perf: bound hot-path work and distribute shards --- README.md | 4 ++ cogs/link_fix.py | 110 ++++++++++++++++++++++++++++++------- cogs/setup.py | 7 ++- database/models/Event.py | 33 ++++++++++- docker-compose.example.yml | 10 ++++ main.py | 3 +- src/hot_path_cache.py | 76 +++++++++++++++++++++++++ src/runtime.py | 34 ++++++++++++ src/settings.py | 2 + src/sharding.py | 34 ++++++++++++ src/websites.py | 24 +++++--- 11 files changed, 306 insertions(+), 31 deletions(-) create mode 100644 src/hot_path_cache.py create mode 100644 src/runtime.py create mode 100644 src/sharding.py diff --git a/README.md b/README.md index 6a9d1ae..d01a92f 100644 --- a/README.md +++ b/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 | diff --git a/cogs/link_fix.py b/cogs/link_fix.py index d5bf6de..77f2d97 100644 --- a/cogs/link_fix.py +++ b/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,14 +254,21 @@ 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(( @@ -253,9 +276,12 @@ async def get_or_create_webhook(channel: GuildMessageableChannel) -> discore.Web 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() diff --git a/cogs/setup.py b/cogs/setup.py index 8471c57..637c777 100644 --- a/cogs/setup.py +++ b/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 diff --git a/database/models/Event.py b/database/models/Event.py index 39e52dc..21cefb5 100644 --- a/database/models/Event.py +++ b/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,6 +50,18 @@ 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""" @@ -50,8 +69,16 @@ async def _flush_loop(cls) -> None: 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()) diff --git a/docker-compose.example.yml b/docker-compose.example.yml index 70cf47b..e420094 100644 --- a/docker-compose.example.yml +++ b/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 diff --git a/main.py b/main.py index c5dab18..0a56491 100644 --- a/main.py +++ b/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() diff --git a/src/hot_path_cache.py b/src/hot_path_cache.py new file mode 100644 index 0000000..555bd0a --- /dev/null +++ b/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) diff --git a/src/runtime.py b/src/runtime.py new file mode 100644 index 0000000..8f0c638 --- /dev/null +++ b/src/runtime.py @@ -0,0 +1,34 @@ +"""Bounded helpers for work that must not block the Discord event loop.""" + +from __future__ import annotations + +import asyncio +import os +from collections.abc import Callable +from typing import TypeVar + + +T = TypeVar('T') + + +class RuntimeBusyError(RuntimeError): + """Raised when bounded blocking work cannot start before its deadline.""" + + +_database_concurrency = max(1, int(os.getenv('DATABASE_WORKER_CONCURRENCY', '8'))) +_database_wait_timeout = max(0.01, float(os.getenv('DATABASE_WORKER_WAIT_TIMEOUT', '0.5'))) +_database_semaphore = asyncio.Semaphore(_database_concurrency) + + +async def run_database(func: Callable[..., T], *args, **kwargs) -> T: + """Run synchronous ORM work off-loop with bounded concurrency and wait time.""" + + try: + await asyncio.wait_for(_database_semaphore.acquire(), timeout=_database_wait_timeout) + except asyncio.TimeoutError as exc: + raise RuntimeBusyError('database worker capacity exhausted') from exc + + try: + return await asyncio.to_thread(func, *args, **kwargs) + finally: + _database_semaphore.release() diff --git a/src/settings.py b/src/settings.py index 0305f54..2087010 100644 --- a/src/settings.py +++ b/src/settings.py @@ -15,6 +15,7 @@ from database.models.CustomWebsite import CustomWebsite from src.utils import * +from src.hot_path_cache import invalidate_guild __all__ = ('SettingsView',) @@ -1859,6 +1860,7 @@ async def refresh(self, interaction: discore.Interaction) -> None: Send or refresh the built view (if already sent) with the current settings :param interaction: The interaction to respond to """ + invalidate_guild(self.ctx.guild.id) await self.build() # Discord API sometimes returns incorrect error code, in this case 404 Unknown interaction when interaction diff --git a/src/sharding.py b/src/sharding.py new file mode 100644 index 0000000..30a9a19 --- /dev/null +++ b/src/sharding.py @@ -0,0 +1,34 @@ +"""Shard assignment parsing shared by the entry point and tests.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping + + +def shard_options(environ: Mapping[str, str] | None = None) -> dict: + """Build validated discord.py shard options from environment variables.""" + + environ = os.environ if environ is None else environ + raw_count = environ.get('SHARD_COUNT') + raw_ids = environ.get('SHARD_IDS') + if not raw_count and not raw_ids: + return {} + if not raw_count: + raise RuntimeError('SHARD_COUNT must be set when SHARD_IDS is configured') + + shard_count = int(raw_count) + if shard_count < 1: + raise ValueError('SHARD_COUNT must be greater than zero') + + options: dict = {'shard_count': shard_count} + if raw_ids: + shard_ids = [int(value.strip()) for value in raw_ids.split(',') if value.strip()] + if not shard_ids: + raise ValueError('SHARD_IDS must contain at least one shard ID') + if len(shard_ids) != len(set(shard_ids)): + raise ValueError('SHARD_IDS must not contain duplicates') + if any(shard_id < 0 or shard_id >= shard_count for shard_id in shard_ids): + raise ValueError('Every SHARD_IDS value must be within SHARD_COUNT') + options['shard_ids'] = shard_ids + return options diff --git a/src/websites.py b/src/websites.py index f710ee6..719f3d6 100644 --- a/src/websites.py +++ b/src/websites.py @@ -2,16 +2,20 @@ Allows fixing links from various websites. """ import logging +import asyncio +import os import re from typing import Type, Iterable, Callable from database.models.Event import * from database.models.Guild import * from src import utils +import aiohttp __all__ = ('WebsiteLink', 'websites') _logger = logging.getLogger(__name__) +_embedez_concurrency = asyncio.Semaphore(max(1, int(os.getenv('EMBEDEZ_CONCURRENCY', '50')))) def call_if_valid(func: Callable) -> Callable: @@ -353,17 +357,23 @@ async def get_fixed_url(self) -> tuple[str | None, str | None]: subdomain = self.route_fix_subdomain() + subdomain prepared_url = self.get_patched_url(self.match['domain'], subdomain, self.route_fix_post_path_segments()) try: - async with utils.session.get("https://embedez.com/api/v1/providers/combined", params={'q': prepared_url}) as response: - if response.status != 200: - _logger.warning("EmbedEZ request error for link: %s (status code: %d, body: %s)", prepared_url, response.status, await response.text()) - await Event.buff_cr({'name': 'embedez_fixer_error', 'data': {'link': prepared_url, 'status_code': response.status, 'response_body': await response.text()}}) - return None, None - search_hash = (await response.json())['data']['key'] - return f"https://embedez.com/embed/{search_hash}", self.fixer_name + async with _embedez_concurrency: + async with utils.session.get("https://embedez.com/api/v1/providers/combined", params={'q': prepared_url}) as response: + if response.status != 200: + response_body = await response.text() + _logger.warning("EmbedEZ request error for link: %s (status code: %d, body: %s)", prepared_url, response.status, response_body) + await Event.buff_cr({'name': 'embedez_fixer_error', 'data': {'link': prepared_url, 'status_code': response.status, 'response_body': response_body}}) + return None, None + search_hash = (await response.json())['data']['key'] + return f"https://embedez.com/embed/{search_hash}", self.fixer_name except asyncio.TimeoutError: _logger.warning("EmbedEZ request timeout for link: %s", prepared_url) await Event.buff_cr({'name': 'embedez_fixer_timeout', 'data': {'link': prepared_url}}) return None, None + except (aiohttp.ClientError, KeyError, TypeError, ValueError) as error: + _logger.warning("EmbedEZ request failed: %s", type(error).__name__) + await Event.buff_cr({'name': 'embedez_fixer_error', 'data': {'error': type(error).__name__}}) + return None, None class TwitterLink(GenericWebsiteLink):