From 722c87c75b0a619ede47549dd9a0161c7ffe0fcf Mon Sep 17 00:00:00 2001
From: io mintz <io@mintz.cc>
Date: Tue, 2 Jun 2020 03:37:18 +0000
Subject: [PATCH] stats: use proper IPC instead of a shared file

---
 cogs/stats.py | 58 +++++++++++++++++++++++++--------------------------
 1 file changed, 29 insertions(+), 29 deletions(-)

diff --git a/cogs/stats.py b/cogs/stats.py
index c1f8e1b..fc1fecd 100644
--- a/cogs/stats.py
+++ b/cogs/stats.py
@@ -13,52 +13,52 @@
 # 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 multiprocessing.shared_memory import ShareableList
 
 from discord.ext import commands
 from bot_bin.stats import BotBinStats
 
 class Stats(BotBinStats):
-	path = 'data/guild_counts.json'
+	def __init__(self, bot):
+		super().__init__(bot)
+		seq = [0] * self.bot.shard_count if self.is_opener() else None
+		# Use our user ID as part of the shm name to allow running multiple instances of the bot on the same machine.
+		# on_ready() hasn't run yet, so we don't have self.bot.user.
+		token = self.bot.config['tokens']['discord']
+		user_id = token.partition('.')[0]
+		self.shlist = ShareableList(seq, name=f'emote-manager-{user_id}')
+
+	def is_opener(self):
+		"""return whether this is the process that should open the shared memory"""
+		return 0 in self.bot.shard_ids
+
+	def is_reporter(self):
+		"""return whether we should report stats to the bot lists"""
+		return self.bot.shard_count - 1 in self.bot.shard_ids
+
+	def cog_unload(self):
+		self.shlist.shm.close()
+		if self.is_opener():
+			self.shlist.shm.unlink()
 
 	@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 = {}
+		for guild in self.bot.guilds:
+			self.shlist[guild.shard_id] += 1
 
-			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)
+		if self.is_reporter():
+			await self.send()
 
 	@commands.Cog.listener()
 	async def on_guild_join(self, guild):
-		self.incr_guild_count(guild.shard_id, +1)
+		self.shlist[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)
+		self.shlist[guild.shard_id] -= 1
 
 	async def guild_count(self):
-		with open(self.path) as f:
-			return sum(json.load(f).values())
+		return sum(self.shlist)
 
 def setup(bot):
 	bot.add_cog(Stats(bot))