From 1c05d7a88072f4a134217525a92590b86aa0ff96 Mon Sep 17 00:00:00 2001 From: Conatum Date: Wed, 22 Dec 2021 13:07:20 +0200 Subject: [PATCH] sharding (blacklists): Blacklist shard support. Moved the `user_blacklist` and `guild_blacklist` to a client TTL cache. --- bot/LionModule.py | 4 +- bot/core/blacklists.py | 44 ++++++++----------- bot/modules/economy/cointop_cmd.py | 2 +- bot/modules/reminders/reminder.py | 2 +- bot/modules/study/stats_cmd.py | 2 +- bot/modules/study/top_cmd.py | 2 +- bot/modules/study/tracking/session_tracker.py | 2 +- bot/modules/study/tracking/time_tracker.py | 2 +- bot/modules/sysadmin/blacklist.py | 33 +++++++++----- bot/modules/workout/tracker.py | 2 +- 10 files changed, 48 insertions(+), 47 deletions(-) diff --git a/bot/LionModule.py b/bot/LionModule.py index 6e89cf5e..9c2a4671 100644 --- a/bot/LionModule.py +++ b/bot/LionModule.py @@ -82,7 +82,7 @@ class LionModule(Module): raise SafeCancellation(details="Module '{}' is not ready.".format(self.name)) # Check global user blacklist - if ctx.author.id in ctx.client.objects['blacklisted_users']: + if ctx.author.id in ctx.client.user_blacklist(): raise SafeCancellation(details='User is blacklisted.') if ctx.guild: @@ -91,7 +91,7 @@ class LionModule(Module): raise SafeCancellation(details='Command channel is no longer reachable.') # Check global guild blacklist - if ctx.guild.id in ctx.client.objects['blacklisted_guilds']: + if ctx.guild.id in ctx.client.guild_blacklist(): raise SafeCancellation(details='Guild is blacklisted.') # Check guild's own member blacklist diff --git a/bot/core/blacklists.py b/bot/core/blacklists.py index 942bd012..1ca5bd9c 100644 --- a/bot/core/blacklists.py +++ b/bot/core/blacklists.py @@ -1,9 +1,8 @@ """ Guild, user, and member blacklists. - -NOTE: The pre-loading methods are not shard-optimised. """ from collections import defaultdict +import cachetools.func from data import tables from meta import client @@ -11,32 +10,22 @@ from meta import client from .module import module -@module.init_task -def load_guild_blacklist(client): +@cachetools.func.ttl_cache(ttl=300) +def guild_blacklist(): """ - Load the blacklisted guilds. + Get the guild blacklist """ rows = tables.global_guild_blacklist.select_where() - client.objects['blacklisted_guilds'] = set(row['guildid'] for row in rows) - if rows: - client.log( - "Loaded {} blacklisted guilds.".format(len(rows)), - context="GUILD_BLACKLIST" - ) + return set(row['guildid'] for row in rows) -@module.init_task -def load_user_blacklist(client): +@cachetools.func.ttl_cache(ttl=300) +def user_blacklist(): """ - Load the blacklisted users. + Get the global user blacklist. """ rows = tables.global_user_blacklist.select_where() - client.objects['blacklisted_users'] = set(row['userid'] for row in rows) - if rows: - client.log( - "Loaded {} globally blacklisted users.".format(len(rows)), - context="USER_BLACKLIST" - ) + return set(row['userid'] for row in rows) @module.init_task @@ -62,18 +51,20 @@ def load_ignored_members(client): ) +@module.init_task +def attach_client_blacklists(client): + client.guild_blacklist = guild_blacklist + client.user_blacklist = user_blacklist + + @module.launch_task async def leave_blacklisted_guilds(client): """ Launch task to leave any blacklisted guilds we are in. - Assumes that the blacklisted guild list has been initialised. """ - # Cache to avoic repeated lookups - blacklisted = client.objects['blacklisted_guilds'] - to_leave = [ guild for guild in client.guilds - if guild.id in blacklisted + if guild.id in guild_blacklist() ] for guild in to_leave: @@ -92,7 +83,8 @@ async def check_guild_blacklist(client, guild): Guild join event handler to check whether the guild is blacklisted. If so, leaves the guild. """ - if guild.id in client.objects['blacklisted_guilds']: + # First refresh the blacklist cache + if guild.id in guild_blacklist(): await guild.leave() client.log( "Automatically left blacklisted guild '{}' (gid:{}) upon join.".format(guild.name, guild.id), diff --git a/bot/modules/economy/cointop_cmd.py b/bot/modules/economy/cointop_cmd.py index 9d1b9b2d..81bdbad9 100644 --- a/bot/modules/economy/cointop_cmd.py +++ b/bot/modules/economy/cointop_cmd.py @@ -43,7 +43,7 @@ async def cmd_topcoin(ctx): # Fetch the leaderboard exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) - exclude.update(ctx.client.objects['blacklisted_users']) + exclude.update(ctx.client.user_blacklist()) exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) args = { diff --git a/bot/modules/reminders/reminder.py b/bot/modules/reminders/reminder.py index d3e4f764..73870341 100644 --- a/bot/modules/reminders/reminder.py +++ b/bot/modules/reminders/reminder.py @@ -134,7 +134,7 @@ class Reminder: """ Execute the reminder. """ - if self.data.userid in client.objects['blacklisted_users']: + if self.data.userid in client.user_blacklist(): self.delete(self.reminderid) return diff --git a/bot/modules/study/stats_cmd.py b/bot/modules/study/stats_cmd.py index 29e412e6..88bc8be5 100644 --- a/bot/modules/study/stats_cmd.py +++ b/bot/modules/study/stats_cmd.py @@ -59,7 +59,7 @@ async def cmd_stats(ctx): # Leaderboard ranks exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) - exclude.update(ctx.client.objects['blacklisted_users']) + exclude.update(ctx.client.user_blacklist()) exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) if target.id in exclude: time_rank = None diff --git a/bot/modules/study/top_cmd.py b/bot/modules/study/top_cmd.py index cb4008f7..79564c1f 100644 --- a/bot/modules/study/top_cmd.py +++ b/bot/modules/study/top_cmd.py @@ -40,7 +40,7 @@ async def cmd_top(ctx): # Fetch the leaderboard exclude = set(m.id for m in ctx.guild_settings.unranked_roles.members) - exclude.update(ctx.client.objects['blacklisted_users']) + exclude.update(ctx.client.user_blacklist()) exclude.update(ctx.client.objects['ignored_members'][ctx.guild.id]) args = { diff --git a/bot/modules/study/tracking/session_tracker.py b/bot/modules/study/tracking/session_tracker.py index 96262f04..9cbe53be 100644 --- a/bot/modules/study/tracking/session_tracker.py +++ b/bot/modules/study/tracking/session_tracker.py @@ -298,7 +298,7 @@ async def session_voice_tracker(client, member, before, after): pending.cancel() if after.channel: - blacklist = client.objects['blacklisted_users'] + blacklist = client.user_blacklist() guild_blacklist = client.objects['ignored_members'][guild.id] untracked = untracked_channels.get(guild.id).data start_session = ( diff --git a/bot/modules/study/tracking/time_tracker.py b/bot/modules/study/tracking/time_tracker.py index 1cb35fa0..46f88ec7 100644 --- a/bot/modules/study/tracking/time_tracker.py +++ b/bot/modules/study/tracking/time_tracker.py @@ -47,7 +47,7 @@ def _scan(guild): members = itertools.chain(*channel_members) # TODO filter out blacklisted users - blacklist = client.objects['blacklisted_users'] + blacklist = client.user_blacklist() guild_blacklist = client.objects['ignored_members'][guild.id] for member in members: diff --git a/bot/modules/sysadmin/blacklist.py b/bot/modules/sysadmin/blacklist.py index 12a2ed9b..90202407 100644 --- a/bot/modules/sysadmin/blacklist.py +++ b/bot/modules/sysadmin/blacklist.py @@ -7,6 +7,8 @@ import discord from cmdClient.checks import is_owner from cmdClient.lib import ResponseTimedOut +from meta.sharding import sharded + from .module import module @@ -26,14 +28,14 @@ async def cmd_guildblacklist(ctx, flags): Description: View, add, or remove guilds from the blacklist. """ - blacklist = ctx.client.objects['blacklisted_guilds'] + blacklist = ctx.client.guild_blacklist() if ctx.args: # guildid parsing items = [item.strip() for item in ctx.args.split(',')] if any(not item.isdigit() for item in items): return await ctx.error_reply( - "Please provide guilds as comma seprated guild ids." + "Please provide guilds as comma separated guild ids." ) guildids = set(int(item) for item in items) @@ -80,9 +82,18 @@ async def cmd_guildblacklist(ctx, flags): insert_keys=('guildid', 'ownerid', 'reason') ) - # Check if we are in any of these guilds - to_leave = (ctx.client.get_guild(guildid) for guildid in to_add) - to_leave = [guild for guild in to_leave if guild is not None] + # Leave freshly blacklisted guilds, accounting for shards + to_leave = [] + for guildid in to_add: + guild = ctx.client.get_guild(guildid) + if not guild and sharded: + try: + guild = await ctx.client.fetch_guild(guildid) + except discord.HTTPException: + pass + if guild: + to_leave.append(guild) + for guild in to_leave: await guild.leave() @@ -102,9 +113,8 @@ async def cmd_guildblacklist(ctx, flags): ) # Refresh the cached blacklist after modification - ctx.client.objects['blacklisted_guilds'] = set( - row['guildid'] for row in ctx.client.data.global_guild_blacklist.select_where() - ) + ctx.client.guild_blacklist.cache_clear() + ctx.client.guild_blacklist() else: # Display the current blacklist # First fetch the full blacklist data @@ -183,7 +193,7 @@ async def cmd_userblacklist(ctx, flags): Description: View, add, or remove users from the blacklist. """ - blacklist = ctx.client.objects['blacklisted_users'] + blacklist = ctx.client.user_blacklist() if ctx.args: # userid parsing @@ -245,9 +255,8 @@ async def cmd_userblacklist(ctx, flags): ) # Refresh the cached blacklist after modification - ctx.client.objects['blacklisted_users'] = set( - row['userid'] for row in ctx.client.data.global_user_blacklist.select_where() - ) + ctx.client.user_blacklist.cache_clear() + ctx.client.user_blacklist() else: # Display the current blacklist # First fetch the full blacklist data diff --git a/bot/modules/workout/tracker.py b/bot/modules/workout/tracker.py index 90eea397..79dc9378 100644 --- a/bot/modules/workout/tracker.py +++ b/bot/modules/workout/tracker.py @@ -170,7 +170,7 @@ async def workout_voice_tracker(client, member, before, after): if member.bot: return - if member.id in client.objects['blacklisted_users']: + if member.id in client.user_blacklist(): return if member.id in client.objects['ignored_members'][member.guild.id]: return