velvox
/
litecord
mirror of https://gitlab.com/litecord/litecord.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Plug state resume into deletion unschedule

This commit is contained in:
Luna 2020-04-04 22:50:00 -03:00
parent 720417d14d
commit 24ff453c6b
2 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
litecord/gateway/state_manager.py
View File

@ -245,6 +245,8 @@ class StateManager:
async def _future_cleanup(self, state: GatewayState): async def _future_cleanup(self, state: GatewayState):
await asyncio.sleep(30) await asyncio.sleep(30)
self.remove(state) self.remove(state)
state.ws.state = None
state.ws = None
async def schedule_deletion(self, state: GatewayState): async def schedule_deletion(self, state: GatewayState):
task = app.loop.create_task(self._future_cleanup(state)) task = app.loop.create_task(self._future_cleanup(state))
@ -253,6 +255,7 @@ class StateManager:
async def unschedule_deletion(self, state: GatewayState): async def unschedule_deletion(self, state: GatewayState):
try: try:
task = self.tasks.pop(state.session_id) task = self.tasks.pop(state.session_id)
task.cancel()
except KeyError: except KeyError:
pass return
task.cancel()

8
litecord/gateway/websocket.py
View File

@ -824,6 +824,7 @@ class GatewayWebsocket:
return await self.invalidate_session(False) return await self.invalidate_session(False)
# relink this connection # relink this connection
await self.app.state_manager.unschedule_deletion(state)
self.state = state self.state = state
state.ws = self state.ws = self
@ -1085,9 +1086,10 @@ class GatewayWebsocket:
task.cancel() task.cancel()
if self.state: if self.state:
self.app.state_manager.remove(self.state) self.app.state_manager.schedule_deletion(self.state)
self.state.ws = None # self.app.state_manager.remove(self.state)
self.state = None # self.state.ws = None
# self.state = None
async def _check_conns(self, user_id): async def _check_conns(self, user_id):
"""Check if there are any existing connections. """Check if there are any existing connections.