jarvis/jarvis-bot
jarvis/jarvis-bot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add cache update method for autocomplete with locks

Browse source
This commit is contained in:
Zeva Rose 2023-09-14 13:32:38 -06:00
parent 9267dac97d
commit 81b7852ac3
1 changed files with 30 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
jarvis/cogs/extra/rolegiver.py
View file

@ -42,6 +42,7 @@ class RolegiverCog(Extension):
self.logger = logging.getLogger(__name__) self.logger = logging.getLogger(__name__)
self._group_cache: dict[int, dict[str, dict[str, int]]] = {} self._group_cache: dict[int, dict[str, dict[str, int]]] = {}
"""`{GUILD_ID: {GROUP_NAME: {ROLE_NAME: ROLE_ID}}}`""" """`{GUILD_ID: {GROUP_NAME: {ROLE_NAME: ROLE_ID}}}`"""
self._cache_update_locks: dict[int, asyncio.Lock] = {}
@listen() @listen()
async def on_ready(self) -> None: async def on_ready(self) -> None:
@ -95,12 +96,11 @@ class RolegiverCog(Extension):
await rolegiver.save() await rolegiver.save()
await ctx.send(f"Rolegiver group {group} created!") await ctx.send(f"Rolegiver group {group} created!")
if ctx.guild.id not in self._group_cache: if ctx.guild.id not in self._group_cache:
self._group_cache[ctx.guild.id] = {rolegiver.group: {}} self._group_cache[ctx.guild.id] = {rolegiver.group: {}}
elif rolegiver.group not in self._group_cache[ctx.guild.id]: elif rolegiver.group not in self._group_cache[ctx.guild.id]:
self._group_cache[ctx.guild.id][rolegiver.group] = {} self._group_cache[ctx.guild.id][rolegiver.group] = {}
@rg_group.subcommand( @rg_group.subcommand(
sub_cmd_name="delete", sub_cmd_description="DDelete a rolegiver group" sub_cmd_name="delete", sub_cmd_description="DDelete a rolegiver group"
@ -563,6 +563,7 @@ class RolegiverCog(Extension):
@_rolegiver_list.autocomplete("group") @_rolegiver_list.autocomplete("group")
async def _autocomplete_group(self, ctx: AutocompleteContext): async def _autocomplete_group(self, ctx: AutocompleteContext):
groups = list(self._group_cache.get(ctx.guild.id, {}).keys()) groups = list(self._group_cache.get(ctx.guild.id, {}).keys())
rolegivers = None
if not groups: if not groups:
rolegivers = await Rolegiver.find(Rolegiver.guild == ctx.guild.id).to_list() rolegivers = await Rolegiver.find(Rolegiver.guild == ctx.guild.id).to_list()
groups = [r.group for r in rolegivers] groups = [r.group for r in rolegivers]
@ -570,6 +571,33 @@ class RolegiverCog(Extension):
results = process.extract(ctx.input_text, groups, limit=5) results = process.extract(ctx.input_text, groups, limit=5)
choices = [{"name": r[0], "value": r[0]} for r in results] choices = [{"name": r[0], "value": r[0]} for r in results]
await ctx.send(choices) await ctx.send(choices)
if rolegivers:
await self._update_cache_from_autocomplete(ctx, rolegivers)
async def _update_cache_from_autocomplete(
self, ctx: AutocompleteContext, rolegivers: list[Rolegiver]
) -> None:
"""Update the cache with a lock to prevent multiple simultaneous updates"""
lock = self._cache_update_locks.get(ctx.guild.id)
if not lock:
lock = asyncio.Lock()
self._cache_update_locks[ctx.guild.id] = lock
if not lock.locked():
async with lock:
if ctx.guild.id not in self._group_cache:
self._group_cache[ctx.guild.id] = {}
for r in rolegivers:
roles = []
for role_id in r.roles:
role = await ctx.guild.fetch_role(role_id)
if not role:
r.roles.remove(role_id)
continue
roles.append(role)
await r.save()
self._group_cache[ctx.guild.id][r.group] = {
role.name: role.role.id for role in roles
}
def setup(bot: Client) -> None: def setup(bot: Client) -> None: