Make disconnection more resilient

This commit is contained in:
Lonami Exo 2023-09-12 20:55:25 +02:00
parent f2b8a91fa9
commit 9e5eb619e1
2 changed files with 18 additions and 8 deletions

View File

@ -485,7 +485,7 @@ class Client:
return await invoke_request(self, self._sender, self._sender_lock, request)
async def __aenter__(self) -> Self:
await self.connect()
await connect(self)
return self
async def __aexit__(
@ -495,4 +495,4 @@ class Client:
tb: Optional[TracebackType],
) -> None:
exc_type, exc, tb
await self.disconnect()
await disconnect(self)

View File

@ -141,14 +141,24 @@ async def connect(self: Client) -> None:
async def disconnect(self: Client) -> None:
if not self._sender:
return
assert self._dispatcher
self._dispatcher.cancel()
await self._dispatcher
self._dispatcher = None
await self._sender.disconnect()
self._sender = None
self._dispatcher.cancel()
try:
await self._dispatcher
except asyncio.CancelledError:
pass
except Exception:
pass # TODO log
finally:
self._dispatcher = None
try:
await self._sender.disconnect()
except Exception:
pass # TODO log
finally:
self._sender = None
self._config.session.state = self._message_box.session_state()
await self._storage.save(self._config.session)