From 5695cf86f1b81af1daed1bbbbd3e56da179541d3 Mon Sep 17 00:00:00 2001
From: io mintz <io@mintz.cc>
Date: Tue, 2 Jun 2020 00:51:06 +0000
Subject: [PATCH] support multi-process sharding

---
 .gitignore    |  1 +
 bot.py        | 16 +++++++++++--
 cogs/stats.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 2 deletions(-)
 create mode 100644 cogs/stats.py

diff --git a/.gitignore b/.gitignore
index 4c778cd..573ebc2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ venv/
 
 /config.py
 /data/config.py
+/data/guild_counts.json
 
 *.png
 *.gif
diff --git a/bot.py b/bot.py
index e5f571e..90773a1 100755
--- a/bot.py
+++ b/bot.py
@@ -31,9 +31,9 @@ class Bot(Bot):
 	startup_extensions = (
 		'cogs.emote',
 		'cogs.meta',
+		'cogs.stats',
 		'bot_bin.debug',
 		'bot_bin.misc',
-		'bot_bin.stats',
 		'jishaku',
 	)
 
@@ -58,5 +58,17 @@ class Bot(Bot):
 			return super().activity
 		return super().activity or discord.Game(f'@{self.user.name} help')
 
+def main():
+	import sys
+
+	kwargs = dict(guild_subscriptions=False, request_offline_members=False)
+
+	if len(sys.argv) >= 3:
+		shard_count = int(sys.argv[1])
+		shard_ids = list(map(int, sys.argv[2].split(',')))
+		Bot(**kwargs, shard_count=shard_count, shard_ids=shard_ids).run()
+	else:
+		Bot(**kwargs).run()
+
 if __name__ == '__main__':
-	Bot().run()
+	main()
diff --git a/cogs/stats.py b/cogs/stats.py
new file mode 100644
index 0000000..c1f8e1b
--- /dev/null
+++ b/cogs/stats.py
@@ -0,0 +1,64 @@
+# © 2020 io mintz <io@mintz.cc>
+#
+# Emote Manager is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# Emote Manager is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with Emote Manager. If not, see <https://www.gnu.org/licenses/>.
+
+import collections
+import json
+
+from discord.ext import commands
+from bot_bin.stats import BotBinStats
+
+class Stats(BotBinStats):
+	path = 'data/guild_counts.json'
+
+	@commands.Cog.listener()
+	async def on_ready(self):
+		with open(self.path, 'w+') as f:
+			contents = f.read()
+			if contents:
+				guild_counts = json.loads(contents)
+			else:
+				guild_counts = {}
+
+			my_counts = collections.Counter()
+
+			for guild in self.bot.guilds:
+				my_counts[str(guild.shard_id)] += 1
+
+			guild_counts.update(my_counts)
+
+			f.seek(0)
+			json.dump(guild_counts, f)
+
+	@commands.Cog.listener()
+	async def on_guild_join(self, guild):
+		self.incr_guild_count(guild.shard_id, +1)
+
+	@commands.Cog.listener()
+	async def on_guild_remove(self, guild):
+		self.incr_guild_count(guild.shard_id, -1)
+
+	def incr_guild_count(self, shard_id, amount):
+		with open(self.path, 'w+') as f:
+			guild_counts = json.load(f)
+			guild_counts[str(shard_id)] += amount
+			f.seek(0)
+			json.dump(guild_counts, f)
+
+	async def guild_count(self):
+		with open(self.path) as f:
+			return sum(json.load(f).values())
+
+def setup(bot):
+	bot.add_cog(Stats(bot))