Gracefully handle user-invoked disconnects

This commit is contained in:
Lonami Exo 2023-10-19 19:28:00 +02:00
parent 854d3745b3
commit 864d5cd444

View File

@ -207,6 +207,9 @@ async def disconnect(self: Client) -> None:
return return
assert self._dispatcher assert self._dispatcher
sender = self._sender
self._sender = None # treated as disconnected
self._dispatcher.cancel() self._dispatcher.cancel()
try: try:
await self._dispatcher await self._dispatcher
@ -220,11 +223,9 @@ async def disconnect(self: Client) -> None:
self._dispatcher = None self._dispatcher = None
try: try:
await self._sender.disconnect() await sender.disconnect()
except Exception: except Exception:
self._logger.exception("unhandled exception during disconnect; this is a bug") self._logger.exception("unhandled exception during disconnect; this is a bug")
finally:
self._sender = None
self._session.state = self._message_box.session_state() self._session.state = self._message_box.session_state()
await self._storage.save(self._session) await self._storage.save(self._session)
@ -272,7 +273,15 @@ async def step_sender(client: Client, sender: Sender, lock: asyncio.Lock) -> Non
pass pass
else: else:
async with lock: async with lock:
try:
updates = await sender.step() updates = await sender.step()
except ConnectionError:
if client.connected:
raise
else:
# disconnect was called, so the socket returning 0 bytes is expected
return
process_socket_updates(client, updates) process_socket_updates(client, updates)