diff --git a/.env.docker.example b/.env.docker.example index 6dcece8..203a56a 100644 --- a/.env.docker.example +++ b/.env.docker.example @@ -7,6 +7,12 @@ brand_name=Axiom OWNER_IDS=YOUR_DISCORD_USER_ID_HERE DASHBOARD_ADMIN_IDS=YOUR_DISCORD_USER_ID_HERE +# Sharding — leave unset for auto (recommended). Multi-process: same TOKEN, shared bot-db volume. +# SHARD_COUNT= +# SHARD_IDS=0,1 +# SHARD_PRIMARY= +# Enable API only on the primary shard process. + # Shared secret — must match between bot API and dashboard DASHBOARD_API_KEY=generate_a_long_random_secret_here diff --git a/bot/.env.example b/bot/.env.example index 438ab65..e79f1d1 100644 --- a/bot/.env.example +++ b/bot/.env.example @@ -9,6 +9,16 @@ NEXT_PUBLIC_BRAND_NAME='Axiom' # BOT_INSTANCE_ID=prod-a # API_PORT=8001 +# ── Discord sharding (optional) ─────────────────────────────────────────────── +# Default (all unset): AutoShardedBot asks Discord for recommended shard count (one process). +# Multi-process cluster: same TOKEN + shared DB volume; set SHARD_COUNT + SHARD_IDS per process. +# API / slash sync / emoji sync only on the primary (shard 0 in SHARD_IDS, or SHARD_PRIMARY=true). +# SHARD_COUNT= +# SHARD_IDS=0,1 +# SHARD_PRIMARY= +# Only enable API on the primary process: +# API_ENABLED=true + # ── Owner / Staff IDs (REQUIRED — comma-separated Discord user IDs) ─────────── # Set your own Discord user ID(s). No foreign IDs are used as fallback. OWNER_IDS=YOUR_DISCORD_USER_ID_HERE diff --git a/bot/Axiom.py b/bot/Axiom.py index 99fc103..b72c866 100644 --- a/bot/Axiom.py +++ b/bot/Axiom.py @@ -96,6 +96,17 @@ async def on_ready(): print(f"Logged in as: {client.user}") print(f"Connected to: {len(client.guilds)} guilds") print(f"Connected to: {len(client.users)} users") + shard_count = client.shard_count or 1 + shard_ids = list(client.shards.keys()) if client.shards else [0] + print(f"Shards: {shard_count} (this process: {shard_ids})") + print(f"Shard primary: {IS_SHARD_PRIMARY}") + if client.latencies: + for sid, lat in client.latencies: + print(f" Shard {sid}: {round(lat * 1000)}ms") + + if not IS_SHARD_PRIMARY: + print("\033[33m◈ Non-primary shard process — skipping API-side sync tasks\033[0m") + return # Sync application emojis on startup await run_sync(TOKEN) @@ -112,17 +123,38 @@ async def on_ready(): client.loop.create_task(update_stats()) +@client.event +async def on_shard_connect(shard_id: int): + print(f"\033[32m◈ Shard {shard_id}: connected\033[0m") + + +@client.event +async def on_shard_ready(shard_id: int): + print(f"\033[32m◈ Shard {shard_id}: ready\033[0m") + + +@client.event +async def on_shard_disconnect(shard_id: int): + print(f"\033[33m◈ Shard {shard_id}: disconnected\033[0m") + + +@client.event +async def on_shard_resumed(shard_id: int): + print(f"\033[32m◈ Shard {shard_id}: resumed\033[0m") + + @client.event async def on_guild_join(guild: discord.Guild): # Only guild-mode needs per-server registration; global mode must not # copy commands into the guild or Discord shows duplicates (/afk twice). - try: - from utils.slash_sync import _sync_mode, sync_guild - if _sync_mode() == "guild": - synced = await sync_guild(client, guild, copy_global=True) - print(f"◈ Slash sync for new guild {guild.id}: {len(synced)} command(s)") - except Exception as e: - print(f"Slash sync on guild join failed: {e}") + if IS_SHARD_PRIMARY: + try: + from utils.slash_sync import _sync_mode, sync_guild + if _sync_mode() == "guild": + synced = await sync_guild(client, guild, copy_global=True) + print(f"◈ Slash sync for new guild {guild.id}: {len(synced)} command(s)") + except Exception as e: + print(f"Slash sync on guild join failed: {e}") if not LOG_CHANNEL_ID: return @@ -328,6 +360,9 @@ def run_api(): uvicorn.run(fastapi_app, host=API_HOST, port=API_PORT, log_level="warning") def keep_alive(): + if not IS_SHARD_PRIMARY: + print(f"\033[33m◈ API Server: Skipped (non-primary shard process)\033[0m") + return if not API_ENABLED: print(f"\033[33m◈ API Server: Disabled via API_ENABLED=false\033[0m") return @@ -341,7 +376,10 @@ keep_alive() # --- Cloudflare Tunnel (HTTPS for API) — only when explicitly enabled --- from utils.tunnel import start_tunnel -start_tunnel() +if IS_SHARD_PRIMARY: + start_tunnel() +else: + print("\033[33m◈ Cloudflare Tunnel: Skipped (non-primary shard process)\033[0m") # --- Main Bot Execution --- async def main(): diff --git a/bot/api/routes/admin.py b/bot/api/routes/admin.py index 2c72f81..02527dd 100644 --- a/bot/api/routes/admin.py +++ b/bot/api/routes/admin.py @@ -87,7 +87,16 @@ async def get_admin_stats(bot: "Axiom" = Depends(get_bot)): AdminNodeStatus( name="Auth Sockets", status="Healthy", - load=f"Shard: {bot.shard_count} | Latency: {round(bot.latency * 1000)}ms", + load=( + f"Shards: {bot.shard_count or 1} | " + f"Latency: {round(bot.latency * 1000)}ms" + + ( + " | " + + ", ".join(f"s{sid}:{round(lat * 1000)}ms" for sid, lat in bot.latencies) + if bot.latencies and len(bot.latencies) > 1 + else "" + ) + ), icon="Lock" ) ] diff --git a/bot/api/routes/bot.py b/bot/api/routes/bot.py index 6a1dfb6..71f96f8 100644 --- a/bot/api/routes/bot.py +++ b/bot/api/routes/bot.py @@ -28,13 +28,20 @@ async def get_status(bot: "Axiom" = Depends(get_bot)): """ Returns the live status of the bot. """ + shard_ids = list(bot.shards.keys()) if bot.shards else ([0] if bot.shard_count else []) + shard_latencies = [ + {"id": sid, "latency_ms": round(lat * 1000, 2)} + for sid, lat in (bot.latencies or []) + ] return BotStatus( user=str(bot.user), id=str(bot.user.id) if bot.user else None, latency=bot.latency * 1000, guild_count=len(bot.guilds), user_count=sum(g.member_count or 0 for g in bot.guilds), - shards=bot.shard_count + shards=bot.shard_count, + shard_ids=shard_ids, + shard_latencies=shard_latencies, ) @router.get("/info", response_model=BotInfo, summary="Get bot info", description="Returns general information about the bot including command count and user reach.") diff --git a/bot/api/schemas.py b/bot/api/schemas.py index 1a4f3b2..5c50f81 100644 --- a/bot/api/schemas.py +++ b/bot/api/schemas.py @@ -31,6 +31,8 @@ class BotStatus(BaseModel): guild_count: int user_count: int shards: Optional[int] + shard_ids: Optional[List[int]] = None + shard_latencies: Optional[List[Dict[str, Any]]] = None # --- Guild Schemas --- diff --git a/bot/cogs/commands/extra.py b/bot/cogs/commands/extra.py index 5157d85..859bfda 100644 --- a/bot/cogs/commands/extra.py +++ b/bot/cogs/commands/extra.py @@ -905,6 +905,11 @@ class Extra(commands.Cog): f"**API (Roundtrip):** `{api_latency}ms`\n" f"**Database:** `{db_latency}`" ) + if self.bot.latencies and len(self.bot.latencies) > 1: + shard_lines = "\n".join( + f"**Shard {sid}:** `{round(lat * 1000)}ms`" for sid, lat in self.bot.latencies + ) + latency_text += f"\n\n**Shards ({self.bot.shard_count or len(self.bot.latencies)}):**\n{shard_lines}" await msg.edit(view=CV2("System Latency Report", latency_text)) diff --git a/bot/core/axiom.py b/bot/core/axiom.py index a499972..6f1db23 100644 --- a/bot/core/axiom.py +++ b/bot/core/axiom.py @@ -19,7 +19,7 @@ import aiohttp import typing from typing import List import aiosqlite -from utils.config import OWNER_IDS, BotName +from utils.config import OWNER_IDS, BotName, get_shard_kwargs from utils import getConfig, updateConfig from utils.Tools import setup_db from .Context import Context @@ -39,6 +39,11 @@ class Axiom(commands.AutoShardedBot): intents = discord.Intents.all() intents.presences = True intents.members = True + shard_kwargs = get_shard_kwargs() + # Allow callers to override shard settings via kwargs + for key in ("shard_count", "shard_ids"): + if key in kwargs: + shard_kwargs[key] = kwargs.pop(key) super().__init__(command_prefix=self.get_prefix, case_insensitive=True, intents=intents, @@ -48,7 +53,7 @@ class Axiom(commands.AutoShardedBot): owner_ids=OWNER_IDS, allowed_mentions=discord.AllowedMentions( everyone=False, replied_user=False, roles=False), - shard_count=1) + **shard_kwargs) self.status_index = 0 self.status_list = [] self._slash_synced = False diff --git a/bot/utils/config.py b/bot/utils/config.py index e82d6fd..26d7c87 100644 --- a/bot/utils/config.py +++ b/bot/utils/config.py @@ -80,6 +80,48 @@ OWNER_IDS_STR: list[str] = [str(i) for i in OWNER_IDS] # Discord application / OAuth client id (same as dashboard DISCORD_CLIENT_ID) DISCORD_CLIENT_ID = os.getenv("DISCORD_CLIENT_ID", "").strip() +# ── Discord sharding ────────────────────────────────────────────────────────── +# SHARD_COUNT empty = AutoShardedBot asks Discord for recommended count (in-process). +# SHARD_IDS (e.g. "0,1") + SHARD_COUNT = this process only runs those shards (cluster). +# Primary process runs API / slash sync / emoji sync / stats (shard 0 or override). + +def _parse_optional_int(env_key: str) -> int | None: + raw = os.getenv(env_key, "").strip() + if not raw or not raw.isdigit(): + return None + return int(raw) + + +SHARD_COUNT: int | None = _parse_optional_int("SHARD_COUNT") +SHARD_IDS: list[int] = _parse_ids("SHARD_IDS") + + +def get_shard_kwargs() -> dict: + """Kwargs for AutoShardedBot.__init__ (omit keys for full auto-sharding).""" + kwargs: dict = {} + if SHARD_COUNT is not None: + kwargs["shard_count"] = SHARD_COUNT + if SHARD_IDS: + if SHARD_COUNT is None: + raise ValueError("SHARD_IDS requires SHARD_COUNT to be set") + kwargs["shard_ids"] = SHARD_IDS + kwargs["shard_count"] = SHARD_COUNT + return kwargs + + +def _resolve_is_shard_primary() -> bool: + raw = os.getenv("SHARD_PRIMARY", "").strip().lower() + if raw in ("1", "true", "yes", "on"): + return True + if raw in ("0", "false", "no", "off"): + return False + if not SHARD_IDS: + return True + return 0 in SHARD_IDS + + +IS_SHARD_PRIMARY: bool = _resolve_is_shard_primary() + # Aliases kept for backwards compatibility with files that import these names BOT_OWNER_IDS = OWNER_IDS BOT_OWNER_IDS_STR = OWNER_IDS_STR