Croc/croccybot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Lazy member chunking.

Browse Source 
No longer try to fetch all members on startup.
Instead chunk on-demand.
This commit is contained in:
Conatum
2023-09-20 20:46:39 +03:00
parent 1a6e248e0e
commit 519fb976aa
15 changed files with 141 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
requirements.txt
View File

@@ -3,7 +3,7 @@ cachetools==4.2.2
configparser==5.0.2
discord.py
iso8601==0.1.16
psycopg
psycopg[pool]
pytz==2021.1
topggpy
psutil

2
src/analytics/snapshot.py
View File

@@ -22,7 +22,7 @@ async def shard_snapshot():
snap = ShardSnapshot(
guild_count=len(bot.guilds),
voice_count=sum(len(channel.members) for guild in bot.guilds for channel in guild.voice_channels),
member_count=sum(len(guild.members) for guild in bot.guilds),
member_count=sum(guild.member_count for guild in bot.guilds),
user_count=len(set(m.id for guild in bot.guilds for m in guild.members))
)
return snap

3
src/bot.py
View File

@@ -69,7 +69,8 @@ async def main():
shard_count=sharding.shard_count,
help_command=None,
proxy=conf.bot.get('proxy', None),
translator=translator
translator=translator,
chunk_guilds_at_startup=False,
) as lionbot:
ctx_bot.set(lionbot)
try:

28
src/meta/LionBot.py
View File

@@ -240,7 +240,7 @@ class LionBot(Bot):
details['guild'] = f"`{ctx.guild.id}` -- `{ctx.guild.name}`"
details['my_guild_perms'] = f"`{ctx.guild.me.guild_permissions.value}`"
if ctx.author:
ownerstr = ' (owner)' if ctx.author == ctx.guild.owner else ''
ownerstr = ' (owner)' if ctx.author.id == ctx.guild.owner_id else ''
details['author_guild_perms'] = f"`{ctx.author.guild_permissions.value}{ownerstr}`"
if ctx.channel.type is discord.enums.ChannelType.private:
details['channel'] = "`Direct Message`"
@@ -281,3 +281,29 @@ class LionBot(Bot):
def add_command(self, command):
if not hasattr(command, '_placeholder_group_'):
super().add_command(command)
def request_chunking_for(self, guild):
if not guild.chunked:
return asyncio.create_task(
self._connection.chunk_guild(guild, wait=False, cache=True),
name=f"Background chunkreq for {guild.id}"
)
async def on_interaction(self, interaction: discord.Interaction):
"""
Adds the interaction author to guild cache if appropriate.
This gets run a little bit late, so it is possible the interaction gets handled
without the author being in case.
"""
guild = interaction.guild
user = interaction.user
if guild is not None and user is not None and isinstance(user, discord.Member):
if not guild.get_member(user.id):
guild._add_member(user)
if guild is not None and not guild.chunked:
# Getting an interaction in the guild is a good enough reason to request chunking
logger.info(
f"Unchunked guild <gid: {guild.id}> requesting chunking after interaction."
)
self.request_chunking_for(guild)

2
src/meta/LionTree.py
View File

@@ -82,7 +82,7 @@ class LionTree(CommandTree):
details['guild'] = f"`{interaction.guild.id}` -- `{interaction.guild.name}`"
details['my_guild_perms'] = f"`{interaction.guild.me.guild_permissions.value}`"
if interaction.user:
ownerstr = ' (owner)' if interaction.user == interaction.guild.owner else ''
ownerstr = ' (owner)' if interaction.user.id == interaction.guild.owner_id else ''
details['user_guild_perms'] = f"`{interaction.user.guild_permissions.value}{ownerstr}`"
if interaction.channel.type is discord.enums.ChannelType.private:
details['channel'] = "`Direct Message`"

2
src/modules/economy/cog.py
View File

@@ -190,7 +190,7 @@ class Economy(LionCog):
# First fetch the members which currently exist
query = self.bot.core.data.Member.table.select_where(guildid=ctx.guild.id)
query.select('userid').with_no_adapter()
if 2 * len(targets) < len(ctx.guild.members):
if 2 * len(targets) < ctx.guild.member_count:
# More efficient to fetch the targets explicitly
query.where(userid=list(targetids))
existent_rows = await query

23
src/modules/member_admin/cog.py
View File

@@ -181,15 +181,17 @@ class MemberAdminCog(LionCog):
finally:
self._adding_roles.discard((member.guild.id, member.id))
@LionCog.listener('on_member_remove')
@LionCog.listener('on_raw_member_remove')
@log_wrap(action="Farewell")
async def admin_member_farewell(self, member: discord.Member):
async def admin_member_farewell(self, payload: discord.RawMemberRemoveEvent):
# Ignore members that just joined
if (member.guild.id, member.id) in self._adding_roles:
guildid = payload.guild_id
userid = payload.user.id
if (guildid, userid) in self._adding_roles:
return
# Set lion last_left, creating the lion_member if needed
lion = await self.bot.core.lions.fetch_member(member.guild.id, member.id)
lion = await self.bot.core.lions.fetch_member(guildid, userid)
await lion.data.update(last_left=utc_now())
# Save member roles
@@ -197,18 +199,21 @@ class MemberAdminCog(LionCog):
self.bot.db.conn = conn
async with conn.transaction():
await self.data.past_roles.delete_where(
guildid=member.guild.id,
userid=member.id
guildid=guildid,
userid=userid
)
# Insert current member roles
if member.roles:
print(type(payload.user))
if isinstance(payload.user, discord.Member) and payload.user.roles:
member = payload.user
await self.data.past_roles.insert_many(
('guildid', 'userid', 'roleid'),
*((member.guild.id, member.id, role.id) for role in member.roles)
*((guildid, userid, role.id) for role in member.roles)
)
logger.debug(
f"Stored persisting roles for member <uid:{member.id}> in <gid:{member.guild.id}>."
f"Stored persisting roles for member <uid:{userid}> in <gid:{guildid}>."
)
# TODO: Event log, and include info about unchunked members
@LionCog.listener('on_guild_join')
async def admin_init_guild(self, guild: discord.Guild):

4
src/modules/member_admin/settings.py
View File

@@ -173,7 +173,7 @@ class MemberAdminSettings(SettingGroup):
'{guild_name}': guild.name,
'{guild_icon}': guild.icon.url if guild.icon else member.default_avatar.url,
'{studying_count}': str(active),
'{member_count}': len(guild.members),
'{member_count}': guild.member_count,
}
recurse_map(
@@ -297,7 +297,7 @@ class MemberAdminSettings(SettingGroup):
'{guild_name}': guild.name,
'{guild_icon}': guild.icon.url if guild.icon else member.default_avatar.url,
'{studying_count}': str(active),
'{member_count}': str(len(guild.members)),
'{member_count}': str(guild.member_count),
'{last_time}': str(last_seen or member.joined_at.timestamp()),
}

10
src/modules/ranks/cog.py
View File

@@ -503,7 +503,15 @@ class RankCog(LionCog):
# Ensure guild is chunked
if not guild.chunked:
members = await guild.chunk()
try:
members = await asyncio.wait_for(guild.chunk(), timeout=60)
except asyncio.TimeoutError:
error = t(_p(
'rank_refresh|error:cannot_chunk|desc',
"Could not retrieve member list from Discord. Please try again later."
))
await ui.set_error(error)
return
else:
members = guild.members
ui.stage_members = True

13
src/modules/schedule/core/session.py
View File

@@ -253,6 +253,12 @@ class ScheduledSession:
overwrites = room.overwrites
for member in members:
mobj = guild.get_member(member.userid)
if not mobj and not guild.chunked:
self.bot.request_chunking_for(guild)
try:
mobj = await guild.fetch_member(member.userid)
except discord.HTTPException:
mobj = None
if mobj:
overwrites[mobj] = discord.PermissionOverwrite(connect=True, view_channel=True)
try:
@@ -297,6 +303,13 @@ class ScheduledSession:
}
for member in members:
mobj = guild.get_member(member.userid)
if not mobj and not guild.chunked:
self.bot.request_chunking_for(guild)
try:
mobj = await guild.fetch_member(member.userid)
except discord.HTTPException:
mobj = None
if mobj:
overwrites[mobj] = discord.PermissionOverwrite(connect=True, view_channel=True)
try:

2
src/modules/schedule/core/timeslot.py
View File

@@ -440,7 +440,7 @@ class TimeSlot:
)
def launch(self) -> asyncio.Task:
self.run_task = asyncio.create_task(self.run())
self.run_task = asyncio.create_task(self.run(), name=f"TimeSlot {self.slotid}")
return self.run_task
@log_wrap(action="TimeSlot Run")

7
src/modules/statistics/cog.py
View File

@@ -1,3 +1,4 @@
import asyncio
import logging
from typing import Optional
@@ -91,7 +92,11 @@ class StatsCog(LionCog):
timestamp=utc_now(),
)
await ctx.interaction.response.send_message(embed=waiting_embed)
await ctx.guild.chunk()
try:
await asyncio.wait_for(ctx.guild.chunk(), timeout=10)
pass
except asyncio.TimeoutError:
pass
else:
await ctx.interaction.response.defer(thinking=True)
ui = LeaderboardUI(self.bot, ctx.author, ctx.guild)

16
src/modules/statistics/ui/leaderboard.py
View File

@@ -75,6 +75,8 @@ class LeaderboardUI(StatsUI):
# (type, period) -> (pagen -> Optional[Future[Card]])
self.cache = {}
self.was_chunked: bool = guild.chunked
async def run(self, interaction: discord.Interaction):
self._original = interaction
@@ -136,6 +138,7 @@ class LeaderboardUI(StatsUI):
# Filter out members which are not in the server and unranked roles and bots
# Usually hits cache
self.was_chunked = self.guild.chunked
unranked_setting = await self.bot.get_cog('StatsCog').settings.UnrankedRoles.get(self.guild.id)
unranked_roleids = set(unranked_setting.data)
true_leaderboard = []
@@ -435,12 +438,19 @@ class LeaderboardUI(StatsUI):
Generate UI message arguments from stored data
"""
t = self.bot.translator.t
chunk_warning = t(_p(
'ui:leaderboard|chunk_warning',
"**Note:** Could not retrieve member list from Discord, so some members may be missing. "
"Try again in a minute!"
))
if self.card is not None:
period_start = self.period_starts[self.current_period]
header = t(_p(
'ui:leaderboard|since',
"Counting statistics since {timestamp}"
)).format(timestamp=discord.utils.format_dt(period_start))
if not self.was_chunked:
header = '\n'.join((header, chunk_warning))
args = MessageArgs(
embed=None,
content=header,
@@ -473,7 +483,11 @@ class LeaderboardUI(StatsUI):
)),
description=empty_description
)
args = MessageArgs(content=None, embed=embed, files=[])
args = MessageArgs(
content=chunk_warning if not self.was_chunked else None,
embed=embed,
files=[]
)
return args
async def refresh_components(self):

79
src/modules/sysadmin/guild_log.py
View File

@@ -1,3 +1,4 @@
import asyncio
import datetime
import discord
@@ -22,8 +23,8 @@ class GuildLog(LionCog):
embed.set_author(name="Left guild!")
# Add more specific information about the guild
embed.add_field(name="Owner", value="{0.name} (ID: {0.id})".format(guild.owner), inline=False)
embed.add_field(name="Members (cached)", value="{}".format(len(guild.members)), inline=False)
embed.add_field(name="Owner", value="<@{}>".format(guild.owner_id), inline=False)
embed.add_field(name="Members", value="{}".format(guild.member_count), inline=False)
embed.add_field(name="Now studying in", value="{} guilds".format(len(self.bot.guilds)), inline=False)
# Retrieve the guild log channel and log the event
@@ -35,39 +36,51 @@ class GuildLog(LionCog):
@LionCog.listener('on_guild_join')
@log_wrap(action="Log Guild Join")
async def log_join_guild(self, guild: discord.Guild):
owner = guild.owner
try:
await asyncio.wait_for(guild.chunk(), timeout=60)
except asyncio.TimeoutError:
pass
bots = 0
known = 0
unknown = 0
other_members = set(mem.id for mem in self.bot.get_all_members() if mem.guild != guild)
# TODO: Add info about when we last joined this guild etc once we have it.
for member in guild.members:
if member.bot:
bots += 1
elif member.id in other_members:
known += 1
else:
unknown += 1
if guild.chunked:
bots = 0
known = 0
unknown = 0
other_members = set(mem.id for mem in self.bot.get_all_members() if mem.guild != guild)
for member in guild.members:
if member.bot:
bots += 1
elif member.id in other_members:
known += 1
else:
unknown += 1
mem1 = "people I know" if known != 1 else "person I know"
mem2 = "new friends" if unknown != 1 else "new friend"
mem3 = "bots" if bots != 1 else "bot"
mem4 = "total members"
known = "`{}`".format(known)
unknown = "`{}`".format(unknown)
bots = "`{}`".format(bots)
total = "`{}`".format(guild.member_count)
mem_str = "{0:<5}\t{4},\n{1:<5}\t{5},\n{2:<5}\t{6}, and\n{3:<5}\t{7}.".format(
known,
unknown,
bots,
total,
mem1,
mem2,
mem3,
mem4
)
else:
mem_str = (
"`{count}` total members.\n"
"(Could not chunk guild within `60` seconds.)"
).format(count=guild.member_count)
mem1 = "people I know" if known != 1 else "person I know"
mem2 = "new friends" if unknown != 1 else "new friend"
mem3 = "bots" if bots != 1 else "bot"
mem4 = "total members"
known = "`{}`".format(known)
unknown = "`{}`".format(unknown)
bots = "`{}`".format(bots)
total = "`{}`".format(guild.member_count)
mem_str = "{0:<5}\t{4},\n{1:<5}\t{5},\n{2:<5}\t{6}, and\n{3:<5}\t{7}.".format(
known,
unknown,
bots,
total,
mem1,
mem2,
mem3,
mem4
)
created = "<t:{}>".format(int(guild.created_at.timestamp()))
embed = discord.Embed(
@@ -77,7 +90,7 @@ class GuildLog(LionCog):
)
embed.set_author(name="Joined guild!")
embed.add_field(name="Owner", value="{0} (ID: {0.id})".format(owner), inline=False)
embed.add_field(name="Owner", value="<@{}>".format(guild.owner_id), inline=False)
embed.add_field(name="Created at", value=created, inline=False)
embed.add_field(name="Members", value=mem_str, inline=False)
embed.add_field(name="Now studying in", value="{} guilds".format(len(self.bot.guilds)), inline=False)

4
src/wards.py
View File

@@ -32,7 +32,7 @@ async def high_management(bot: LionBot, member: discord.Member, guild: discord.G
async def low_management(bot: LionBot, member: discord.Member, guild: discord.Guild):
if not guild:
return True
if await high_management(bot, member):
if await high_management(bot, member, guild):
return True
return member.guild_permissions.manage_guild
@@ -196,7 +196,7 @@ async def equippable_role(bot: LionBot, target_role: discord.Role, actor: discor
"You need the `MANAGE_ROLES` permission before you can configure roles!"
)).format(role=target_role.mention)
)
elif actor.top_role <= target_role and not actor == guild.owner:
elif actor.top_role <= target_role and not actor.id == guild.owner_id:
raise UserInputError(
t(_p(
'ward:equippable_role|error:actor_top_role',