Refactored ingestion to semaphore-bounded tasks, centralized error handling, and consolidated metrics.
CI / Formatting (push) Successful in 5s
CI / Linting (push) Successful in 5s
CI / Tests (push) Successful in 22s
CI / Type Checking (push) Successful in 11s
CI / Spelling (push) Successful in 5s

This commit is contained in:
2026-03-22 22:45:31 -04:00
parent c1f967ece9
commit 093426ac61
7 changed files with 180 additions and 135 deletions
-21
View File
@@ -80,26 +80,6 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
' variable or "crabstero.db").' ' variable or "crabstero.db").'
), ),
) )
workers_env = os.environ.get("INGESTION_WORKERS", "4")
try:
workers_default = int(workers_env)
except ValueError:
logger.warning(
"Invalid INGESTION_WORKERS value '%s', defaulting to 4.",
workers_env,
)
workers_default = 4
parser.add_argument(
"--ingestion-workers",
type=int,
default=workers_default,
help=(
"Number of concurrent ingestion workers"
" (default: INGESTION_WORKERS environment"
" variable or 4)."
),
)
parser.add_argument( parser.add_argument(
"--ingest-only", "--ingest-only",
action="store_true", action="store_true",
@@ -157,7 +137,6 @@ def main() -> None:
bot = Crabstero( bot = Crabstero(
token=args.token, token=args.token,
database_path=args.database_path, database_path=args.database_path,
ingestion_workers=args.ingestion_workers,
ingest_only=args.ingest_only, ingest_only=args.ingest_only,
metrics_address=args.listen_metrics, metrics_address=args.listen_metrics,
) )
+144 -37
View File
@@ -15,7 +15,7 @@
"""The simple nonversation Discord bot. """The simple nonversation Discord bot.
Provides the Crabstero subclass that owns the full bot lifecycle: database connection, Provides the Crabstero subclass that owns the full bot lifecycle: database connection,
cog loading, ingestion worker pool, and graceful shutdown. cog loading, background ingestion, and graceful shutdown.
""" """
import asyncio import asyncio
@@ -24,6 +24,7 @@ from typing import override
import discord import discord
from discord import app_commands from discord import app_commands
from discord.app_commands import AppCommandError, CommandInvokeError
from discord.ext import commands from discord.ext import commands
from crabstero import __version__ as crabstero_version from crabstero import __version__ as crabstero_version
@@ -31,11 +32,11 @@ from crabstero.cache import IngestCache
from crabstero.database import Database from crabstero.database import Database
from crabstero.listeners import interaction, message, server_events from crabstero.listeners import interaction, message, server_events
from crabstero.metrics import ( from crabstero.metrics import (
CHANNEL_INGESTION_ERRORS,
DISCORD_EVENTS, DISCORD_EVENTS,
DISCORD_LATENCY, DISCORD_LATENCY,
ERRORS,
GUILD_COUNT, GUILD_COUNT,
INGESTION_BACKLOG, INGESTION_ACTIVE,
MetricsServer, MetricsServer,
) )
from crabstero.tasks.ingestion import ingest_channel from crabstero.tasks.ingestion import ingest_channel
@@ -44,27 +45,27 @@ logger = logging.getLogger(__name__)
type IngestableChannel = discord.TextChannel | discord.VoiceChannel type IngestableChannel = discord.TextChannel | discord.VoiceChannel
_MAX_CONCURRENT_INGESTIONS = 4
class Crabstero(commands.Bot): class Crabstero(commands.Bot):
"""Central bot subclass that owns all lifecycle state. """Central bot subclass that owns all lifecycle state.
The database is opened in setup_hook and closed in close(). Background The database is opened in setup_hook and closed in close(). Background
ingestion is handled by a bounded queue and a fixed worker pool. ingestion is handled by semaphore-bounded dynamic tasks.
""" """
def __init__( def __init__(
self, self,
token: str, token: str,
database_path: str, database_path: str,
ingestion_workers: int = 4,
ingest_only: bool = False, ingest_only: bool = False,
metrics_address: tuple[str, int] | None = None, metrics_address: tuple[str, int] | None = None,
) -> None: ) -> None:
"""Configure intents, store configuration, and prepare ingestion queue state. """Configure intents, store configuration, and prepare ingestion state.
:param token: The Discord bot token. :param token: The Discord bot token.
:param database_path: The file path to the SQLite database. :param database_path: The file path to the SQLite database.
:param ingestion_workers: The number of concurrent ingestion workers.
:param ingest_only: When True, the bot only ingests data and never responds. :param ingest_only: When True, the bot only ingests data and never responds.
:param metrics_address: Optional (host, port) for the Prometheus metrics server. :param metrics_address: Optional (host, port) for the Prometheus metrics server.
""" """
@@ -81,10 +82,9 @@ class Crabstero(commands.Bot):
self._token = token self._token = token
self._database_path = database_path self._database_path = database_path
self._ingestion_worker_count = ingestion_workers
self._ingest_only = ingest_only self._ingest_only = ingest_only
self._ingestion_queue: asyncio.Queue[IngestableChannel] = asyncio.Queue() self._ingestion_semaphore = asyncio.Semaphore(_MAX_CONCURRENT_INGESTIONS)
self._ingestion_workers: list[asyncio.Task[None]] = [] self._ingestion_tasks: dict[int, asyncio.Task[None]] = {}
self._db: Database | None = None self._db: Database | None = None
self.ingest_cache = IngestCache() self.ingest_cache = IngestCache()
@@ -93,7 +93,7 @@ class Crabstero(commands.Bot):
DISCORD_LATENCY.set_function(lambda: self.latency) DISCORD_LATENCY.set_function(lambda: self.latency)
GUILD_COUNT.set_function(lambda: len(self.guilds)) GUILD_COUNT.set_function(lambda: len(self.guilds))
INGESTION_BACKLOG.set_function(self._ingestion_queue.qsize) INGESTION_ACTIVE.set_function(lambda: len(self._ingestion_tasks))
repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero" repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero"
self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})" self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})"
@@ -115,10 +115,9 @@ class Crabstero(commands.Bot):
@override @override
async def setup_hook(self) -> None: async def setup_hook(self) -> None:
"""Open the database, start ingestion workers, and load all cogs.""" """Open the database, start the ingest cache, and load all cogs."""
self._db = await Database.connect(self._database_path) self._db = await Database.connect(self._database_path)
self.ingest_cache.start() self.ingest_cache.start()
self._start_ingestion_workers()
if self._metrics_address is not None: if self._metrics_address is not None:
server = MetricsServer(*self._metrics_address) server = MetricsServer(*self._metrics_address)
@@ -131,6 +130,47 @@ class Crabstero(commands.Bot):
await message.setup(self) await message.setup(self)
await server_events.setup(self) await server_events.setup(self)
@self.tree.error
async def on_app_command_error(
interaction: discord.Interaction, error: AppCommandError
) -> None:
original = (
error.original if isinstance(error, CommandInvokeError) else error
)
command_name = (
interaction.command.name if interaction.command else "unknown"
)
ERRORS.labels(source="command").inc()
logger.error(
"Unhandled exception in app command '%s'.",
command_name,
exc_info=original,
)
if self._metrics_server is not None:
reply = (
"An error occurred while processing this command."
" The developer has been notified,"
" please try again later."
)
else:
reply = (
"An error occurred while processing this command."
" Please try again later."
)
try:
if interaction.response.is_done():
await interaction.followup.send(reply, ephemeral=True)
else:
await interaction.response.send_message(reply, ephemeral=True)
except discord.HTTPException:
logger.debug(
"Failed to send error response for command '%s'.",
command_name,
)
if not self._ingest_only: if not self._ingest_only:
# Only sync slash commands if the registered commands # Only sync slash commands if the registered commands
# differ from local definitions. # differ from local definitions.
@@ -160,6 +200,17 @@ class Crabstero(commands.Bot):
DISCORD_EVENTS.labels(event=event).inc() DISCORD_EVENTS.labels(event=event).inc()
super().dispatch(event, *args, **kwargs) super().dispatch(event, *args, **kwargs)
@override
async def on_error(
self, event_method: str, /, *args: object, **kwargs: object
) -> None:
"""Increment the global error counter for event listener exceptions.
:param event_method: The name of the event that raised the exception.
"""
ERRORS.labels(source=event_method).inc()
logger.error("Unhandled exception in %s.", event_method, exc_info=True)
async def on_ready(self) -> None: async def on_ready(self) -> None:
"""Log that the bot has started successfully.""" """Log that the bot has started successfully."""
logger.info("Crabstero started!") logger.info("Crabstero started!")
@@ -171,14 +222,15 @@ class Crabstero(commands.Bot):
@override @override
async def close(self) -> None: async def close(self) -> None:
"""Cancel ingestion workers, close the database, and then the bot connection.""" """Cancel ingestion tasks, close the database, and then the bot connection."""
if self.is_closed(): if self.is_closed():
return return
logger.info("Shutting down Crabstero...") logger.info("Shutting down Crabstero...")
for worker in self._ingestion_workers: tasks = list(self._ingestion_tasks.values())
worker.cancel() self._ingestion_tasks.clear()
await asyncio.gather(*self._ingestion_workers, return_exceptions=True) for task in tasks:
self._ingestion_workers.clear() task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
await self.ingest_cache.stop() await self.ingest_cache.stop()
if self._metrics_server is not None: if self._metrics_server is not None:
await self._metrics_server.stop() await self._metrics_server.stop()
@@ -187,28 +239,83 @@ class Crabstero(commands.Bot):
await super().close() await super().close()
def queue_channel_for_ingestion(self, channel: IngestableChannel) -> None: def queue_channel_for_ingestion(self, channel: IngestableChannel) -> None:
"""Enqueue a single channel for background message history ingestion. """Create a background task to ingest a single channel.
:param channel: The channel to enqueue. Duplicate requests for a channel that is already in-flight are ignored.
Concurrency is bounded by the ingestion semaphore.
:param channel: The channel to ingest.
""" """
self._ingestion_queue.put_nowait(channel) if channel.id in self._ingestion_tasks:
return
def _start_ingestion_workers(self) -> None:
"""Spawn the fixed pool of ingestion worker tasks."""
for i in range(self._ingestion_worker_count):
task = asyncio.create_task( task = asyncio.create_task(
self._ingestion_worker(), name=f"ingestion-worker-{i}" self._ingest_one(channel), name=f"ingest-{channel.id}"
) )
self._ingestion_workers.append(task) self._ingestion_tasks[channel.id] = task
task.add_done_callback(lambda t: self._on_ingestion_done(channel.id, t))
async def _ingestion_worker(self) -> None: async def _ingest_one(self, channel: IngestableChannel) -> None:
"""Loop forever pulling channels from the ingestion queue and ingesting them.""" """Acquire the semaphore and ingest one channel."""
while True: async with self._ingestion_semaphore:
channel = await self._ingestion_queue.get()
try:
await ingest_channel(channel, self.db) await ingest_channel(channel, self.db)
except Exception:
CHANNEL_INGESTION_ERRORS.inc() def _on_ingestion_done(self, channel_id: int, task: asyncio.Task[None]) -> None:
logger.exception("Ingestion failed for channel %s", channel.id) """Clean up a finished ingestion task and log any errors."""
finally: self._ingestion_tasks.pop(channel_id, None)
self._ingestion_queue.task_done() if task.cancelled():
return
exc = task.exception()
if exc is not None:
ERRORS.labels(source="ingestion").inc()
logger.error("Ingestion task failed.", exc_info=exc)
class TrackedView(discord.ui.View):
"""Base View that increments the global error counter on failures."""
@override
async def on_error(
self,
interaction: discord.Interaction,
error: Exception,
item: discord.ui.Item[TrackedView],
/,
) -> None:
"""Increment the error counter and log the exception.
:param interaction: The interaction that led to the failure.
:param error: The exception that was raised.
:param item: The item that failed the dispatch.
"""
ERRORS.labels(source="view").inc()
logger.error(
"Unhandled exception in view %r for item %r.",
self,
item,
exc_info=error,
)
class TrackedModal(discord.ui.Modal):
"""Base Modal that increments the global error counter on failures."""
@override
async def on_error(
self,
interaction: discord.Interaction,
error: Exception,
item: discord.ui.Item[TrackedModal] | None = None,
/,
) -> None:
"""Increment the error counter and log the exception.
:param interaction: The interaction that led to the failure.
:param error: The exception that was raised.
:param item: Unused. Present for BaseView signature compatibility.
"""
ERRORS.labels(source="modal").inc()
logger.error(
"Unhandled exception in modal %r.",
self,
exc_info=error,
)
+2 -19
View File
@@ -17,7 +17,6 @@
Provides the /pingme slash command as a Cog with an app command. Provides the /pingme slash command as a Cog with an app command.
""" """
import logging
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
import discord import discord
@@ -30,8 +29,6 @@ from crabstero.flags import EntityType, Flag
if TYPE_CHECKING: if TYPE_CHECKING:
from crabstero.bot import Crabstero from crabstero.bot import Crabstero
logger = logging.getLogger(__name__)
class InteractionCog(commands.Cog): class InteractionCog(commands.Cog):
"""Cog for handling slash command interactions.""" """Cog for handling slash command interactions."""
@@ -55,15 +52,13 @@ class InteractionCog(commands.Cog):
:param interaction: The interaction event. :param interaction: The interaction event.
""" """
# Wrap in try/except to always send a response even if the database fails.
try:
if await flags.is_flag_set( if await flags.is_flag_set(
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
): ):
await flags.clear_flag( await flags.clear_flag(
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
) )
metrics.PINGME_OPTED_OUT.inc() metrics.PINGME.labels(outcome="opted_out").inc()
await interaction.response.send_message( await interaction.response.send_message(
"I will no longer ping you for messages which" "I will no longer ping you for messages which"
" mention you. If you decide to opt back in," " mention you. If you decide to opt back in,"
@@ -74,7 +69,7 @@ class InteractionCog(commands.Cog):
await flags.set_flag( await flags.set_flag(
self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS self.bot.db, interaction.user, EntityType.USER, Flag.ALLOW_PINGS
) )
metrics.PINGME_OPTED_IN.inc() metrics.PINGME.labels(outcome="opted_in").inc()
await interaction.response.send_message( await interaction.response.send_message(
"I will now ping you for messages which mention" "I will now ping you for messages which mention"
" you. If you change your mind, run" " you. If you change your mind, run"
@@ -82,18 +77,6 @@ class InteractionCog(commands.Cog):
ephemeral=True, ephemeral=True,
) )
except Exception:
metrics.PINGME_ERRORS.inc()
logger.exception(
"An exception occurred while attempting to execute"
" slash command responder for command 'pingme'."
)
await interaction.response.send_message(
"An error occurred while executing your command."
" Please try again later.",
ephemeral=True,
)
async def setup(bot: Crabstero) -> None: async def setup(bot: Crabstero) -> None:
"""Add the InteractionCog to the bot. """Add the InteractionCog to the bot.
+11 -17
View File
@@ -61,21 +61,15 @@ CHANNELS_INGESTED = Counter(
"crabstero_channels_ingested_total", "crabstero_channels_ingested_total",
"Channels fully ingested successfully", "Channels fully ingested successfully",
) )
CHANNEL_INGESTION_ERRORS = Counter( PINGME = Counter(
"crabstero_channel_ingestion_errors_total", "crabstero_pingme_total",
"Channel ingestions that failed", "/pingme command outcomes",
["outcome"],
) )
PINGME_OPTED_IN = Counter( ERRORS = Counter(
"crabstero_pingme_opted_in_total", "crabstero_errors_total",
"/pingme opt-ins", "Unhandled exceptions by source",
) ["source"],
PINGME_OPTED_OUT = Counter(
"crabstero_pingme_opted_out_total",
"/pingme opt-outs",
)
PINGME_ERRORS = Counter(
"crabstero_pingme_errors_total",
"/pingme exceptions",
) )
DISCORD_EVENTS = Counter( DISCORD_EVENTS = Counter(
"crabstero_discord_events_total", "crabstero_discord_events_total",
@@ -91,9 +85,9 @@ GUILD_COUNT = Gauge(
"crabstero_guild_count", "crabstero_guild_count",
"Guilds the bot is currently in", "Guilds the bot is currently in",
) )
INGESTION_BACKLOG = Gauge( INGESTION_ACTIVE = Gauge(
"crabstero_ingestion_backlog_channels", "crabstero_ingestion_active_channels",
"Channels currently queued for ingestion", "Channels pending or in-progress for ingestion",
) )
MESSAGE_INGESTION_DURATION = Histogram( MESSAGE_INGESTION_DURATION = Histogram(
+1 -1
View File
@@ -14,7 +14,7 @@
"""Bulk-ingests the message history of channels. """Bulk-ingests the message history of channels.
Channels are enqueued for processing by the bot's fixed ingestion worker pool. Channels are processed by the bot's semaphore-bounded dynamic tasks.
""" """
import logging import logging
-11
View File
@@ -46,17 +46,6 @@ class TestParseArgs:
args = _parse_args(["--token", "test", "--database-path", "/custom.db"]) args = _parse_args(["--token", "test", "--database-path", "/custom.db"])
assert args.database_path == "/custom.db" assert args.database_path == "/custom.db"
def test_ingestion_workers_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
"""Omitting --ingestion-workers defaults to 4."""
monkeypatch.delenv("INGESTION_WORKERS", raising=False)
args = _parse_args(["--token", "test"])
assert args.ingestion_workers == 4
def test_ingestion_workers_override(self) -> None:
"""--ingestion-workers overrides the default."""
args = _parse_args(["--token", "test", "--ingestion-workers", "8"])
assert args.ingestion_workers == 8
def test_ingest_only_flag(self) -> None: def test_ingest_only_flag(self) -> None:
"""--ingest-only sets ingest_only to True.""" """--ingest-only sets ingest_only to True."""
args = _parse_args(["--token", "test", "--ingest-only"]) args = _parse_args(["--token", "test", "--ingest-only"])
+3 -10
View File
@@ -41,18 +41,11 @@ class TestMetricObjects:
pytest.param("crabstero_embeds_generated_total", id="embeds-generated"), pytest.param("crabstero_embeds_generated_total", id="embeds-generated"),
pytest.param("crabstero_messages_processed_total", id="messages-processed"), pytest.param("crabstero_messages_processed_total", id="messages-processed"),
pytest.param("crabstero_channels_ingested_total", id="channels-ingested"), pytest.param("crabstero_channels_ingested_total", id="channels-ingested"),
pytest.param( pytest.param("crabstero_pingme_total", id="pingme"),
"crabstero_channel_ingestion_errors_total", pytest.param("crabstero_errors_total", id="errors"),
id="channel-ingestion-errors",
),
pytest.param("crabstero_pingme_opted_in_total", id="pingme-opted-in"),
pytest.param("crabstero_pingme_opted_out_total", id="pingme-opted-out"),
pytest.param("crabstero_pingme_errors_total", id="pingme-errors"),
pytest.param("crabstero_discord_latency_seconds", id="discord-latency"), pytest.param("crabstero_discord_latency_seconds", id="discord-latency"),
pytest.param("crabstero_guild_count", id="guild-count"), pytest.param("crabstero_guild_count", id="guild-count"),
pytest.param( pytest.param("crabstero_ingestion_active_channels", id="ingestion-active"),
"crabstero_ingestion_backlog_channels", id="ingestion-backlog"
),
pytest.param( pytest.param(
"crabstero_message_ingestion_duration_seconds", "crabstero_message_ingestion_duration_seconds",
id="message-ingestion-duration", id="message-ingestion-duration",