1
0
Fork 0
mirror of https://github.com/uhIgnacio/EmoteManager.git synced 2024-08-15 02:23:13 +00:00

stats: use proper IPC instead of a shared file

This commit is contained in:
io mintz 2020-06-02 03:37:18 +00:00
parent e8881788a8
commit 722c87c75b

View file

@ -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))