From 24ff453c6ba3b9cffa8de76da8e91d0aef2770b2 Mon Sep 17 00:00:00 2001 From: Luna Date: Sat, 4 Apr 2020 22:50:00 -0300 Subject: [PATCH] Plug state resume into deletion unschedule --- litecord/gateway/state_manager.py | 7 +++++-- litecord/gateway/websocket.py | 8 +++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/litecord/gateway/state_manager.py b/litecord/gateway/state_manager.py index ce600fb..eb2e0fb 100644 --- a/litecord/gateway/state_manager.py +++ b/litecord/gateway/state_manager.py @@ -245,6 +245,8 @@ class StateManager: async def _future_cleanup(self, state: GatewayState): await asyncio.sleep(30) self.remove(state) + state.ws.state = None + state.ws = None async def schedule_deletion(self, state: GatewayState): task = app.loop.create_task(self._future_cleanup(state)) @@ -253,6 +255,7 @@ class StateManager: async def unschedule_deletion(self, state: GatewayState): try: task = self.tasks.pop(state.session_id) - task.cancel() except KeyError: - pass + return + + task.cancel() diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index a98752f..2a47fdc 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -824,6 +824,7 @@ class GatewayWebsocket: return await self.invalidate_session(False) # relink this connection + await self.app.state_manager.unschedule_deletion(state) self.state = state state.ws = self @@ -1085,9 +1086,10 @@ class GatewayWebsocket: task.cancel() if self.state: - self.app.state_manager.remove(self.state) - self.state.ws = None - self.state = None + self.app.state_manager.schedule_deletion(self.state) + # self.app.state_manager.remove(self.state) + # self.state.ws = None + # self.state = None async def _check_conns(self, user_id): """Check if there are any existing connections.