diff --git a/litecord/blueprints/guilds.py b/litecord/blueprints/guilds.py index 2f7b3b0..e2244bb 100644 --- a/litecord/blueprints/guilds.py +++ b/litecord/blueprints/guilds.py @@ -183,7 +183,7 @@ async def delete_guild(guild_id): # remove from the dispatcher so nobody # becomes the little memer that tries to fuck up with # everybody's gateway - app.dispatcher.remove_guild(guild_id) + await app.dispatcher.remove('guild', guild_id) return '', 204 diff --git a/litecord/blueprints/users.py b/litecord/blueprints/users.py index 02e83a4..117b9fd 100644 --- a/litecord/blueprints/users.py +++ b/litecord/blueprints/users.py @@ -106,7 +106,7 @@ async def leave_guild(guild_id: int): } ) - await app.dispatcher.unsub_guild(guild_id, user_id) + await app.dispatcher.unsub('guild', guild_id, user_id) await app.dispatcher.dispatch_guild('GUILD_MEMBER_REMOVE', { 'guild_id': str(guild_id), diff --git a/litecord/dispatcher.py b/litecord/dispatcher.py index 224cccf..a1e2c58 100644 --- a/litecord/dispatcher.py +++ b/litecord/dispatcher.py @@ -18,6 +18,8 @@ class EventDispatcher: 'guild': GuildDispatcher(self), 'member': MemberDispatcher(self), 'user': UserDispatcher(self), + + # TODO: channel, friends } async def action(self, backend_str: str, action: str, key, identifier): @@ -45,16 +47,35 @@ class EventDispatcher: """ backend = self.backends[backend_str] key = backend.KEY_TYPE(key) - return await backend._dispatch(key, *args, **kwargs) + return await backend.dispatch(key, *args, **kwargs) async def reset(self, backend_str: str, key: Any): """Reset the bucket in the given backend.""" backend = self.backends[backend_str] key = backend.KEY_TYPE(key) - return await backend._reset(key) + return await backend.reset(key) + + async def remove(self, backend_str: str, key: Any): + """Remove a key from the backend. This + might be a different operation than resetting.""" + backend = self.backends[backend_str] + key = backend.KEY_TYPE(key) + return await backend.remove(key) async def sub_many(self, backend_str: str, identifier: Any, keys: list): """Subscribe to many buckets inside a single backend at a time.""" for key in keys: await self.subscribe(backend_str, key, identifier) + + async def dispatch_guild(self, guild_id, event, data): + """Backwards compatibility.""" + return await self.dispatch('guild', guild_id, event, data) + + async def dispatch_user_guild(self, user_id, guild_id, event, data): + """Backwards compatibility.""" + return await self.dispatch('member', (guild_id, user_id), event, data) + + async def dispatch_user(self, user_id, event, data): + """Backwards compatibility.""" + return await self.dispatch('user', user_id, event, data) diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 339964b..cecf0e2 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -287,9 +287,11 @@ class GatewayWebsocket: """Subscribe to all available guilds""" guild_ids = await self._guild_ids() log.info('subscribing to {} guilds', len(guild_ids)) - self.ext.dispatcher.sub_many(self.state.user_id, guild_ids) + await self.ext.dispatcher.sub_many('guild', + self.state.user_id, guild_ids) async def update_status(self, status: dict): + """Update the status of the current websocket connection.""" if status is None: status = { 'afk': False, diff --git a/litecord/pubsub/dispatcher.py b/litecord/pubsub/dispatcher.py index 8c93d67..42d83e7 100644 --- a/litecord/pubsub/dispatcher.py +++ b/litecord/pubsub/dispatcher.py @@ -21,6 +21,12 @@ class Dispatcher: async def dispatch(self, _key, *_args, **_kwargs): raise NotImplementedError + async def reset(self, _key): + raise NotImplementedError + + async def remove(self, _key): + raise NotImplementedError + async def _dispatch_states(self, states: list, event: str, data) -> int: dispatched = 0 diff --git a/litecord/pubsub/guild.py b/litecord/pubsub/guild.py index c6f0ecf..e8d624f 100644 --- a/litecord/pubsub/guild.py +++ b/litecord/pubsub/guild.py @@ -26,6 +26,12 @@ class GuildDispatcher(Dispatcher): async def reset(self, guild_id: int): self.guild_buckets[guild_id] = set() + async def remove(self, guild_id: int): + try: + self.guild_buckets.pop(guild_id) + except KeyError: + pass + async def dispatch(self, guild_id: int, event_name: str, event_payload: Any): user_ids = self.guild_buckets[guild_id] @@ -41,7 +47,7 @@ class GuildDispatcher(Dispatcher): if not states: # user is actually disconnected, # so we should just unsub it - await self._unsub(guild_id, user_id) + await self.unsub(guild_id, user_id) continue dispatched += await self._dispatch_states( diff --git a/litecord/pubsub/member.py b/litecord/pubsub/member.py index 328f452..bdb61b2 100644 --- a/litecord/pubsub/member.py +++ b/litecord/pubsub/member.py @@ -2,14 +2,15 @@ from .dispatcher import Dispatcher class MemberDispatcher(Dispatcher): - KEY_TYPE = int - VAL_TYPE = int + KEY_TYPE = tuple - async def dispatch(self, guild_id: int, user_id: int, event, data): + async def dispatch(self, key, event, data): """Dispatch a single event to a member. This is shard-aware. """ + guild_id, user_id = key + # fetch shards states = self.sm.fetch_states(user_id, guild_id)