From 99f52a4143b82a66633c5b28e736a2011a2d8f3c Mon Sep 17 00:00:00 2001 From: AEnterprise Date: Mon, 1 Jun 2020 15:56:36 +0200 Subject: [PATCH] clustering --- GearBot/Bot/GearBot.py | 2 ++ GearBot/Bot/TheRealGearBot.py | 14 +++++++------- GearBot/Cogs/PromMonitoring.py | 2 +- GearBot/GearBot.py | 26 +++++++++++++++++++++----- GearBot/Util/Emoji.py | 3 ++- GearBot/Util/GearbotLogging.py | 2 +- clusterloader.sh | 19 +++++++++++++++++++ 7 files changed, 53 insertions(+), 15 deletions(-) create mode 100644 clusterloader.sh diff --git a/GearBot/Bot/GearBot.py b/GearBot/Bot/GearBot.py index 86120afe..99978ac5 100644 --- a/GearBot/Bot/GearBot.py +++ b/GearBot/Bot/GearBot.py @@ -27,10 +27,12 @@ class GearBot(AutoShardedBot): version = "" dash_guild_users = set() dash_guild_watchers = dict() + cluster = 0 def __init__(self, *args, loop=None, **kwargs): super().__init__(*args, loop=loop, **kwargs) self.metrics = PromMonitors(self) + self.cluster = kwargs.get("cluster", 0) def dispatch(self, event_name, *args, **kwargs): if "socket" not in event_name not in ["message_edit"]: diff --git a/GearBot/Bot/TheRealGearBot.py b/GearBot/Bot/TheRealGearBot.py index d4b224f9..59e21249 100644 --- a/GearBot/Bot/TheRealGearBot.py +++ b/GearBot/Bot/TheRealGearBot.py @@ -35,11 +35,11 @@ async def initialize(bot, startup=False): bot.locked = True try: #database - GearbotLogging.info("Connecting to the database.") + GearbotLogging.info(f"Cluster {bot.cluster} connecting to the database.") await DatabaseConnector.init() - GearbotLogging.info("Database connection established.") + GearbotLogging.info(f"Cluster {bot.cluster} database connection established.") - Emoji.initialize(bot) + await Emoji.initialize(bot) Utils.initialize(bot) InfractionUtils.initialize(bot) bot.data = { @@ -52,8 +52,8 @@ async def initialize(bot, startup=False): if startup: c = await Utils.get_commit() bot.version = c - GearbotLogging.info(f"GearBot spinning up version {c}") - await GearbotLogging.bot_log(f"{Emoji.get_chat_emoji('ALTER')} GearBot spinning up version {c}") + GearbotLogging.info(f"GearBot cluster {bot.cluster} spinning up version {c}") + await GearbotLogging.bot_log(f"{Emoji.get_chat_emoji('ALTER')} GearBot cluster {bot.cluster} spinning up version {c}") if bot.redis_pool is None: try: @@ -66,8 +66,8 @@ async def initialize(bot, startup=False): GearbotLogging.error("==============Failed to connect to redis==============") await GearbotLogging.bot_log(f"{Emoji.get_chat_emoji('NO')} Failed to connect to redis, caching unavailable") else: - GearbotLogging.info("Redis connection established") - await GearbotLogging.bot_log(f"{Emoji.get_chat_emoji('YES')} Redis connection established, let's go full speed!") + GearbotLogging.info("Cluster {bot.cluster} redis connection established") + await GearbotLogging.bot_log(f"{Emoji.get_chat_emoji('YES')} Cluster {bot.cluster} redis connection established, let's go full speed!") if bot.aiosession is None: bot.aiosession = aiohttp.ClientSession() diff --git a/GearBot/Cogs/PromMonitoring.py b/GearBot/Cogs/PromMonitoring.py index 07fce301..2a9546cc 100644 --- a/GearBot/Cogs/PromMonitoring.py +++ b/GearBot/Cogs/PromMonitoring.py @@ -45,7 +45,7 @@ async def create_site(self): runner = web.AppRunner(metrics_app) await self.bot.loop.create_task(runner.setup()) - site = web.TCPSite(runner) + site = web.TCPSite(runner, 'localhost', 8090 + self.bot.cluster) await site.start() self.metric_server = site diff --git a/GearBot/GearBot.py b/GearBot/GearBot.py index b19eac24..f904ed50 100644 --- a/GearBot/GearBot.py +++ b/GearBot/GearBot.py @@ -14,14 +14,13 @@ def prefix_callable(bot, message): return TheRealGearBot.prefix_callable(bot, message) -gearbot = GearBot(command_prefix=prefix_callable, case_insensitive=True, max_messages = 100) #100 is the min for some reason - - - if __name__ == '__main__': parser = ArgumentParser() parser.add_argument("--token", help="Specify your Discord token") + parser.add_argument("--total_shards", help="Total shard count") + parser.add_argument("--num_shards", help="Amount of shards to start in this cluster") + parser.add_argument("--offset", help="Shard offset") GearbotLogging.init_logger() @@ -34,10 +33,27 @@ def prefix_callable(bot, message): token = Configuration.get_master_var("LOGIN_TOKEN") else: token = input("Please enter your Discord token: ") + + args = { + "command_prefix": prefix_callable, + "case_insensitive": True, + "max_messages": None, + } + if clargs.total_shards: + total_shards = int(clargs.total_shards) + offset = int(clargs.offset) + num_shards = int(clargs.num_shards) + args.update({ + "shard_count": total_shards, + "cluster": offset, + "shard_ids": [*range(offset * num_shards, (offset * num_shards) + num_shards)] + }) + + gearbot = GearBot(**args) + gearbot.remove_command("help") GearbotLogging.info("Ready to go, spinning up the gears") gearbot.run(token) GearbotLogging.info("GearBot shutting down, cleaning up") gearbot.database_connection.close() GearbotLogging.info("Cleanup complete") - diff --git a/GearBot/Util/Emoji.py b/GearBot/Util/Emoji.py index 07814656..c57a7f67 100644 --- a/GearBot/Util/Emoji.py +++ b/GearBot/Util/Emoji.py @@ -73,7 +73,8 @@ } -def initialize(bot): +async def initialize(bot): + emoji_guild = await bot.fetch_guild(Configuration.get_master_var("EMOJI_GUILD")) for name, eid in Configuration.get_master_var("EMOJI", {}).items(): emojis[name] = utils.get(bot.emojis, id=eid) diff --git a/GearBot/Util/GearbotLogging.py b/GearBot/Util/GearbotLogging.py index 9af3d811..b53a50f1 100644 --- a/GearBot/Util/GearbotLogging.py +++ b/GearBot/Util/GearbotLogging.py @@ -293,7 +293,7 @@ def init_logger(): async def initialize(bot: commands.Bot, channelID): global BOT_LOG_CHANNEL, BOT, STARTUP_ERRORS BOT = bot - BOT_LOG_CHANNEL = bot.get_channel(int(channelID)) + BOT_LOG_CHANNEL = await bot.fetch_channel(int(channelID)) if BOT_LOG_CHANNEL is None: LOGGER.error( "==========================Logging channel is misconfigured, aborting startup!==========================") diff --git a/clusterloader.sh b/clusterloader.sh new file mode 100644 index 00000000..18867264 --- /dev/null +++ b/clusterloader.sh @@ -0,0 +1,19 @@ +#! /bin/bash +if [ -s upgradeRequest ]; then + git pull origin + python3 -m pip install -U -r requirements.txt --user + rm -rf upgradeRequest +fi +SHARDS=2 +CLUSTERS=2 +COUNT=0 +TOTAL_SHARDS=$(($SHARDS * $CLUSTERS)) +LAST=$(($SHARDS-1)) +while [[ $COUNT < $LAST ]]; do + OFFSET=$((SHARDS*$COUNT)) + echo "Starting GearBot cluster $COUNT with $SHARDS shards (offset $OFFSET)" + $(python3 GearBot/GearBot.py --total_shards $TOTAL_SHARDS --num_shards $SHARDS --offset $OFFSET &) + sleep $((5*$SHARDS)) + COUNT=$(($COUNT+1)) +done +$(python3 GearBot/GearBot.py --total_shards $TOTAL_SHARDS --num_shards $SHARDS --offset $OFFSET) \ No newline at end of file