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

support multi-process sharding

This commit is contained in:
io mintz 2020-06-02 00:51:06 +00:00
parent d16bb2c6f1
commit 5695cf86f1
3 changed files with 79 additions and 2 deletions

1
.gitignore vendored
View file

@ -9,6 +9,7 @@ venv/
/config.py /config.py
/data/config.py /data/config.py
/data/guild_counts.json
*.png *.png
*.gif *.gif

16
bot.py
View file

@ -31,9 +31,9 @@ class Bot(Bot):
startup_extensions = ( startup_extensions = (
'cogs.emote', 'cogs.emote',
'cogs.meta', 'cogs.meta',
'cogs.stats',
'bot_bin.debug', 'bot_bin.debug',
'bot_bin.misc', 'bot_bin.misc',
'bot_bin.stats',
'jishaku', 'jishaku',
) )
@ -58,5 +58,17 @@ class Bot(Bot):
return super().activity return super().activity
return super().activity or discord.Game(f'@{self.user.name} help') 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__': if __name__ == '__main__':
Bot().run() main()

64
cogs/stats.py Normal file
View file

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