From 9e5eb619e162676edcd7b3916c6af001294804c6 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Tue, 12 Sep 2023 20:55:25 +0200 Subject: [PATCH] Make disconnection more resilient --- .../telethon/_impl/client/client/client.py | 4 ++-- .../src/telethon/_impl/client/client/net.py | 22 ++++++++++++++----- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/client/src/telethon/_impl/client/client/client.py b/client/src/telethon/_impl/client/client/client.py index 8939fb10..11b74d1c 100644 --- a/client/src/telethon/_impl/client/client/client.py +++ b/client/src/telethon/_impl/client/client/client.py @@ -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) diff --git a/client/src/telethon/_impl/client/client/net.py b/client/src/telethon/_impl/client/client/net.py index 9d087405..7cd05869 100644 --- a/client/src/telethon/_impl/client/client/net.py +++ b/client/src/telethon/_impl/client/client/net.py @@ -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)