diff --git a/litecord/dispatcher.py b/litecord/dispatcher.py index 009ac5e..e7e392b 100644 --- a/litecord/dispatcher.py +++ b/litecord/dispatcher.py @@ -4,7 +4,8 @@ from typing import List, Any from logbook import Logger from .pubsub import GuildDispatcher, MemberDispatcher, \ - UserDispatcher, ChannelDispatcher, FriendDispatcher + UserDispatcher, ChannelDispatcher, FriendDispatcher, \ + LazyGuildDispatcher log = Logger(__name__) @@ -35,6 +36,7 @@ class EventDispatcher: 'channel': ChannelDispatcher(self), 'user': UserDispatcher(self), 'friend': FriendDispatcher(self), + 'lazy_guild': LazyGuildDispatcher(self), } async def action(self, backend_str: str, action: str, key, identifier): diff --git a/litecord/gateway/state_manager.py b/litecord/gateway/state_manager.py index 56d00bc..e393a79 100644 --- a/litecord/gateway/state_manager.py +++ b/litecord/gateway/state_manager.py @@ -22,12 +22,16 @@ class StateManager: # } self.states = defaultdict(dict) + #: raw mapping from session ids to GatewayState + self.states_raw = {} + def insert(self, state: GatewayState): """Insert a new state object.""" user_states = self.states[state.user_id] log.debug('inserting state: {!r}', state) user_states[state.session_id] = state + self.states_raw[state.session_id] = state def fetch(self, user_id: int, session_id: str) -> GatewayState: """Fetch a state object from the manager. @@ -40,11 +44,20 @@ class StateManager: """ return self.states[user_id][session_id] + def fetch_raw(self, session_id: str) -> GatewayState: + """Fetch a single state given the Session ID.""" + return self.states_raw[session_id] + def remove(self, state): """Remove a state from the registry""" if not state: return + try: + self.states_raw.pop(state.session_id) + except KeyError: + pass + try: log.debug('removing state: {!r}', state) self.states[state.user_id].pop(state.session_id) diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 9c6bd4d..7fce369 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -641,9 +641,11 @@ class GatewayWebsocket: This is the known structure of GUILD_MEMBER_LIST_UPDATE: + group_id = 'online' | 'offline' | role_id (string) + sync_item = { 'group': { - 'id': string, // 'online' | 'offline' | any role id + 'id': group_id, 'count': num } } | { @@ -678,7 +680,7 @@ class GatewayWebsocket: // separately from the online list? 'groups': [ { - 'id': string // 'online' | 'offline' | any role id + 'id': group_id 'count': num }, ... ] @@ -713,65 +715,16 @@ class GatewayWebsocket: if guild_id not in gids: return - member_ids = await self.storage.get_member_ids(guild_id) - log.debug('lazy: loading {} members', len(member_ids)) + # make shard query + lazy_guilds = self.ext.dispatcher.backends['lazy_guild'] - # the current implementation is rudimentary and only - # generates two groups: online and offline, using - # PresenceManager.guild_presences to fill list_data. + for chan_id, ranges in data['channels'].items(): + chan_id = int(chan_id) + member_list = await lazy_guilds.get_gml(chan_id) - # this also doesn't take account the channels in lazy_request. - - guild_presences = await self.presence.guild_presences(member_ids, - guild_id) - - online = [{'member': p} - for p in guild_presences - if p['status'] == 'online'] - offline = [{'member': p} - for p in guild_presences - if p['status'] == 'offline'] - - log.debug('lazy: {} presences, online={}, offline={}', - len(guild_presences), - len(online), - len(offline)) - - # construct items in the WORST WAY POSSIBLE. - items = [{ - 'group': { - 'id': 'online', - 'count': len(online), - } - }] + online + [{ - 'group': { - 'id': 'offline', - 'count': len(offline), - } - }] + offline - - await self.dispatch('GUILD_MEMBER_LIST_UPDATE', { - 'id': 'everyone', - 'guild_id': data['guild_id'], - 'groups': [ - { - 'id': 'online', - 'count': len(online), - }, - { - 'id': 'offline', - 'count': len(offline), - } - ], - - 'ops': [ - { - 'range': [0, 99], - 'op': 'SYNC', - 'items': items - } - ] - }) + await member_list.shard_query( + self.state.session_id, ranges + ) async def process_message(self, payload): """Process a single message coming in from the client.""" diff --git a/litecord/pubsub/__init__.py b/litecord/pubsub/__init__.py index 7320867..31388de 100644 --- a/litecord/pubsub/__init__.py +++ b/litecord/pubsub/__init__.py @@ -3,7 +3,8 @@ from .member import MemberDispatcher from .user import UserDispatcher from .channel import ChannelDispatcher from .friend import FriendDispatcher +from .lazy_guild import LazyGuildDispatcher __all__ = ['GuildDispatcher', 'MemberDispatcher', 'UserDispatcher', 'ChannelDispatcher', - 'FriendDispatcher'] + 'FriendDispatcher', 'LazyGuildDispatcher'] diff --git a/litecord/pubsub/lazy_guild.py b/litecord/pubsub/lazy_guild.py index ba8a451..ef7e0b2 100644 --- a/litecord/pubsub/lazy_guild.py +++ b/litecord/pubsub/lazy_guild.py @@ -1,5 +1,9 @@ +""" +Main code for Lazy Guild implementation in litecord. +""" +import pprint from collections import defaultdict -from typing import Any +from typing import Any, List, Dict from logbook import Logger @@ -8,36 +12,214 @@ from .dispatcher import Dispatcher log = Logger(__name__) -class GuildMemberList(): - def __init__(self, guild_id: int): +class GuildMemberList: + """This class stores the current member list information + for a guild (by channel). + + As channels can have different sets of roles that can + read them and so, different lists, this is more of a + "channel member list" than a guild member list. + + Attributes + ---------- + main_lg: LazyGuildDispatcher + Main instance of :class:`LazyGuildDispatcher`, + so that we're able to use things such as :class:`Storage`. + guild_id: int + The Guild ID this instance is referring to. + channel_id: int + The Channel ID this instance is referring to. + member_list: List + The actual member list information. + state: set + The set of session IDs that are subscribed to the guild. + + User IDs being used as the identifier in GuildMemberList + is a wrong assumption. It is true Discord rolled out + lazy guilds to all of the userbase, but users that are bots, + for example, can still rely on PRESENCE_UPDATEs. + """ + def __init__(self, guild_id: int, + channel_id: int, main_lg): + self.main_lg = main_lg self.guild_id = guild_id + self.channel_id = channel_id - # TODO: initialize list with actual member info - self._uninitialized = True - self.member_list = [] + # a really long chain of classes to get + # to the storage instance... + main = main_lg.main_dispatcher + self.storage = main.app.storage + self.presence = main.app.presence + self.state_man = main.app.state_manager - #: holds the state of subscribed users + self.member_list = None + self.items = None + + #: holds the state of subscribed shards + # to this channels' member list self.state = set() async def _init_check(self): """Check if the member list is initialized before messing with it.""" - if self._uninitialized: + if self.member_list is None: await self._init_member_list() + async def get_roles(self) -> List[Dict[str, Any]]: + """Get role information, but only: + - the ID + - the name + - the position + + of all HOISTED roles.""" + # TODO: write own query for this + # TODO: calculate channel overrides + roles = await self.storage.get_role_data(self.guild_id) + + return [{ + 'id': role['id'], + 'name': role['name'], + 'position': role['position'] + } for role in roles if role['hoist']] + async def _init_member_list(self): """Fill in :attr:`GuildMemberList.member_list` with information about the guilds' members.""" - pass + member_ids = await self.storage.get_member_ids(self.guild_id) - async def sub(self, user_id: int): - """Subscribe a user to the member list.""" + guild_presences = await self.presence.guild_presences( + member_ids, self.guild_id) + + guild_roles = await self.get_roles() + + # sort by position + guild_roles.sort(key=lambda role: role['position']) + roleids = [r['id'] for r in guild_roles] + + # groups are: + # - roles that are hoisted + # - "online" and "offline", with "online" + # being for people without any roles. + + groups = roleids + ['online', 'offline'] + + log.debug('{} presences, {} groups', + len(guild_presences), len(groups)) + + group_data = {group: [] for group in groups} + + print('group data', group_data) + + def _try_hier(role_id: str, roleids: list): + """Try to fetch a role's position in the hierarchy""" + try: + return roleids.index(role_id) + except ValueError: + # the given role isn't on a group + # so it doesn't count for our purposes. + return 0 + + for presence in guild_presences: + # simple group (online or offline) + # we'll decide on the best group for the presence later on + simple_group = ('offline' + if presence['status'] == 'offline' + else 'online') + + # get the best possible role + roles = sorted( + presence['roles'], + key=lambda role_id: _try_hier(role_id, roleids) + ) + + try: + best_role_id = roles[0] + except IndexError: + # no hoisted roles exist in the guild, assign + # the @everyone role + best_role_id = str(self.guild_id) + + print('best role', best_role_id, str(self.guild_id)) + print('simple group assign', simple_group) + + # if the best role is literally the @everyone role, + # this user has no hoisted roles + if best_role_id == str(self.guild_id): + # this user has no roles, put it on online/offline + group_data[simple_group].append(presence) + continue + + # this user has a best_role that isn't the + # @everyone role, so we'll put them in the respective group + group_data[best_role_id].append(presence) + + # go through each group and sort the resulting members by display name + + members = await self.storage.get_member_data(self.guild_id) + member_nicks = {member['user']['id']: member.get('nick') + for member in members} + + # now we'll sort each group by their display name + # (can be their current nickname OR their username + # if no nickname is set) + print('pre-sorted group data') + pprint.pprint(group_data) + + for _, group_list in group_data.items(): + def display_name(presence: dict) -> str: + uid = presence['user']['id'] + + uname = presence['user']['username'] + nick = member_nicks[uid] + + return nick or uname + + group_list.sort(key=display_name) + + pprint.pprint(group_data) + + self.member_list = { + 'groups': groups, + 'data': group_data + } + + def get_items(self) -> list: + """Generate the main items list,""" + if self.member_list is None: + return [] + + if self.items: + return self.items + + groups = self.member_list['groups'] + + res = [] + for group in groups: + members = self.member_list['data'][group] + + res.append({ + 'group': { + 'id': group, + 'count': len(members), + } + }) + + for member in members: + res.append({ + 'member': member + }) + + self.items = res + return res + + async def sub(self, session_id: str): + """Subscribe a shard to the member list.""" await self._init_check() - self.state.add(user_id) + self.state.add(session_id) - async def unsub(self, user_id: int): - """Unsubscribe a user from the member list""" - self.state.discard(user_id) + async def unsub(self, session_id: str): + """Unsubscribe a shard from the member list""" + self.state.discard(session_id) # once we reach 0 subscribers, # we drop the current member list we have (for memory) @@ -45,8 +227,70 @@ class GuildMemberList(): # uninitialized) for a future subscriber. if not self.state: - self.member_list = [] - self._uninitialized = True + self.member_list = None + + async def shard_query(self, session_id: str, ranges: list): + """Send a GUILD_MEMBER_LIST_UPDATE event + for a shard that is querying about the member list. + + Paramteters + ----------- + session_id: str + The Session ID querying information. + channel_id: int + The Channel ID that we want information on. + ranges: List[List[int]] + ranges of the list that we want. + """ + + await self._init_check() + + # make sure this is a sane state + state = self.state_man.fetch_raw(session_id) + if not state: + await self.unsub(session_id) + return + + # since this is a sane state AND + # trying to query, we automatically + # subscribe the state to this list + await self.sub(session_id) + + reply = { + 'guild_id': str(self.guild_id), + + # TODO: everyone for channels without overrides + # channel_id for channels WITH overrides. + + 'id': 'everyone', + # 'id': str(self.channel_id), + + 'groups': [ + { + 'count': len(self.member_list['data'][group]), + 'id': group + } for group in self.member_list['groups'] + ], + + 'ops': [], + } + + for start, end in ranges: + itemcount = end - start + + # ignore incorrect ranges + if itemcount < 0: + continue + + items = self.get_items() + + reply['ops'].append({ + 'op': 'SYNC', + 'range': [start, end], + 'items': items[start:end], + }) + + await state.ws.dispatch('GUILD_MEMBER_LIST_UPDATE', reply) async def dispatch(self, event: str, data: Any): """The dispatch() method here, instead of being @@ -61,7 +305,7 @@ class GuildMemberList(): calls to :meth:`GuildMemberList.dispatch` """ - if self._uninitialized: + if self.member_list is None: # if the list is currently uninitialized, # no subscribers actually happened, so # we can safely drop the incoming event. @@ -70,25 +314,36 @@ class GuildMemberList(): class LazyGuildDispatcher(Dispatcher): """Main class holding the member lists for lazy guilds.""" + # channel ids KEY_TYPE = int - VAL_TYPE = int + + # the session ids subscribing to channels + VAL_TYPE = str def __init__(self, main): super().__init__(main) - self.state = defaultdict(GuildMemberList) - async def sub(self, guild_id, user_id): - await self.state[guild_id].sub(user_id) + self.storage = main.app.storage - async def unsub(self, guild_id, user_id): - await self.state[guild_id].unsub(user_id) + # {chan_id: gml, ...} + self.state = {} - async def dispatch(self, guild_id: int, event: str, data): - """Dispatch an event to the member list. + async def get_gml(self, channel_id: int): + try: + return self.state[channel_id] + except KeyError: + guild_id = await self.storage.guild_from_channel( + channel_id + ) - GuildMemberList will make sure of converting it to - GUILD_MEMBER_LIST_UPDATE events. - """ - member_list = self.state[guild_id] - await member_list.dispatch(event, data) + gml = GuildMemberList(guild_id, channel_id, self) + self.state[channel_id] = gml + return gml + async def sub(self, chan_id, session_id): + gml = await self.get_gml(chan_id) + await gml.sub(session_id) + + async def unsub(self, chan_id, session_id): + gml = await self.get_gml(chan_id) + await gml.unsub(session_id) diff --git a/litecord/storage.py b/litecord/storage.py index a7e37eb..7337e54 100644 --- a/litecord/storage.py +++ b/litecord/storage.py @@ -440,6 +440,7 @@ class Storage: permissions, managed, mentionable FROM roles WHERE guild_id = $1 + ORDER BY position ASC """, guild_id) return list(map(dict, roledata)) @@ -966,7 +967,6 @@ class Storage: """, user_id) for row in settings: - print(dict(row)) gid = int(row['guild_id']) drow = dict(row)