From c570c71fd843287f56d292e4206edab92809237d Mon Sep 17 00:00:00 2001 From: Itsydv Date: Thu, 14 Aug 2025 20:26:59 +0530 Subject: [PATCH 1/2] fixed and modified restart logic for TimeoutError or ConnectionError --- pyrogram/session/session.py | 44 ++++++++++++++++++++----------------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index 538522f9a..47ca78863 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -105,7 +105,7 @@ def __init__( self.recv_task = None self.is_started = asyncio.Event() - self.restart_event = asyncio.Event() + self._restart_lock = asyncio.Lock() async def start(self): while True: @@ -192,7 +192,8 @@ async def stop(self): self.ping_task_event.clear() - await self.connection.close() + if self.connection: + await self.connection.close() if self.recv_task: try: @@ -205,10 +206,20 @@ async def stop(self): log.info("Session stopped") async def restart(self): - self.restart_event.set() - await self.stop() - await self.start() - self.restart_event.clear() + async with self._restart_lock: + if self.is_started.is_set(): + log.info("A restart was requested, stopping session...") + try: + await self.stop() + except Exception as e: + log.error("An error occurred during session stop: %s", e) + + log.info("Restarting session...") + try: + await self.start() + log.info("Session restarted successfully.") + except Exception as e: + log.error("An error occurred during session start: %s", e) async def handle_packet(self, packet): try: @@ -333,6 +344,10 @@ async def recv_worker(self): packet = await asyncio.wait_for(self.connection.recv(), timeout=1) except asyncio.TimeoutError: continue + except ConnectionError: + if self.is_started.is_set(): + self.client.loop.create_task(self.restart()) + break if packet is None or len(packet) == 4: if packet: @@ -446,21 +461,10 @@ async def invoke( raise e from None (log.warning if retries < 2 else log.info)( - '[%s] Retrying "%s" due to: %s', + '[%s] Retrying "%s" due to: %s. Performing a restart.', Session.MAX_RETRIES - retries + 1, query_name, str(e) or repr(e) ) - # restart was never being called after Exception block - if not self.restart_event.is_set(): - self.client.loop.create_task(self.restart()) - else: - # multiple Exceptions can be raised in a row, so we need to wait for the restart to finish - try: - await asyncio.wait_for(self.restart_event.wait(), self.WAIT_TIMEOUT) - except asyncio.TimeoutError: - pass - - await asyncio.sleep(0.5) - - return await self.invoke(query, retries - 1, timeout) + await self.restart() + retries -= 1 From 63e9a9a5f60a5c8b205c4c836ef95a523bba9a6c Mon Sep 17 00:00:00 2001 From: Itsydv Date: Fri, 15 Aug 2025 22:55:33 +0530 Subject: [PATCH 2/2] fixed and modified restart logic v2 --- pyrogram/session/session.py | 826 +++++++++++++++++------------------- 1 file changed, 400 insertions(+), 426 deletions(-) diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index 47ca78863..10116c638 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -29,15 +29,15 @@ from pyrogram.connection import Connection from pyrogram.crypto import mtproto from pyrogram.errors import ( - AuthKeyDuplicated, - BadMsgNotification, - FloodPremiumWait, - FloodWait, - InternalServerError, - RPCError, - SecurityCheckMismatch, - ServiceUnavailable, - Unauthorized, + AuthKeyDuplicated, + BadMsgNotification, + FloodPremiumWait, + FloodWait, + InternalServerError, + RPCError, + SecurityCheckMismatch, + ServiceUnavailable, + Unauthorized, ) from pyrogram.raw.all import layer from pyrogram.raw.core import FutureSalts, Int, MsgContainer, TLObject @@ -48,423 +48,397 @@ class Result: - def __init__(self): - self.value = None - self.event = asyncio.Event() + def __init__(self): + self.value = None + self.event = asyncio.Event() class Session: - START_TIMEOUT = 2 - WAIT_TIMEOUT = 15 - SLEEP_THRESHOLD = 10 - MAX_RETRIES = 10 - ACKS_THRESHOLD = 10 - PING_INTERVAL = 5 - STORED_MSG_IDS_MAX_SIZE = 1000 * 2 - - TRANSPORT_ERRORS = { - 404: "auth key not found", - 429: "transport flood", - 444: "invalid DC" - } - - def __init__( - self, - client: "pyrogram.Client", - dc_id: int, - auth_key: bytes, - test_mode: bool, - is_media: bool = False, - is_cdn: bool = False - ): - self.client = client - self.dc_id = dc_id - self.auth_key = auth_key - self.test_mode = test_mode - self.is_media = is_media - self.is_cdn = is_cdn - - self.connection: Optional[Connection] = None - - self.auth_key_id = sha1(auth_key).digest()[-8:] - - self.session_id = os.urandom(8) - self.msg_factory = MsgFactory() - - self.salt = 0 - - self.pending_acks = set() - - self.results = {} - - self.stored_msg_ids = [] - - self.ping_task = None - self.ping_task_event = asyncio.Event() - - self.recv_task = None - - self.is_started = asyncio.Event() - self._restart_lock = asyncio.Lock() - - async def start(self): - while True: - self.connection = self.client.connection_factory( - dc_id=self.dc_id, - test_mode=self.test_mode, - ipv6=self.client.ipv6, - proxy=self.client.proxy, - media=self.is_media, - protocol_factory=self.client.protocol_factory, - loop=self.client.loop - ) - - try: - await self.connection.connect() - - self.recv_task = self.client.loop.create_task(self.recv_worker()) - - await self.send(raw.functions.Ping(ping_id=0), timeout=self.START_TIMEOUT) - - if not self.is_cdn: - await self.send( - raw.functions.InvokeWithLayer( - layer=layer, - query=raw.functions.InitConnection( - api_id=await self.client.storage.api_id(), - app_version=self.client.app_version, - device_model=self.client.device_model, - system_version=self.client.system_version, - system_lang_code=self.client.system_lang_code, - lang_pack=self.client.lang_pack, - lang_code=self.client.lang_code, - query=raw.functions.help.GetConfig(), - params=self.client.init_connection_params, - ) - ), - timeout=self.START_TIMEOUT - ) - - self.ping_task = self.client.loop.create_task(self.ping_worker()) - - log.info("Session initialized: Pyrogram v%s (Layer %s)", pyrogram.__version__, layer) - log.info("Device: %s - %s", self.client.device_model, self.client.app_version) - log.info("System: %s (%s)", self.client.system_version, self.client.lang_code) - except (AuthKeyDuplicated, Unauthorized) as e: - await self.stop() - raise e - except ConnectionError as e: - await self.stop() - # raise e - except (OSError, RPCError): - await self.stop() - except Exception as e: - await self.stop() - raise e - else: - break - - self.is_started.set() - - log.info("Session started") - - if callable(self.client.connect_handler): - try: - await self.client.connect_handler(self.client, self) - except Exception as e: - log.exception(e) - - async def stop(self): - if callable(self.client.disconnect_handler): - try: - await self.client.disconnect_handler(self.client, self) - except Exception as e: - log.exception(e) - - self.is_started.clear() - - self.stored_msg_ids.clear() - - self.ping_task_event.set() - - if self.ping_task is not None: - await self.ping_task - - self.ping_task_event.clear() - - if self.connection: - await self.connection.close() - - if self.recv_task: - try: - await self.recv_task - except asyncio.CancelledError: - pass - - self.recv_task = None - - log.info("Session stopped") - - async def restart(self): - async with self._restart_lock: - if self.is_started.is_set(): - log.info("A restart was requested, stopping session...") - try: - await self.stop() - except Exception as e: - log.error("An error occurred during session stop: %s", e) - - log.info("Restarting session...") - try: - await self.start() - log.info("Session restarted successfully.") - except Exception as e: - log.error("An error occurred during session start: %s", e) - - async def handle_packet(self, packet): - try: - data = await self.client.loop.run_in_executor( - pyrogram.crypto_executor, - mtproto.unpack, - BytesIO(packet), - self.session_id, - self.auth_key, - self.auth_key_id - ) - except ValueError as e: - log.debug(e) - self.client.loop.create_task(self.restart()) - return - - messages = ( - data.body.messages - if isinstance(data.body, MsgContainer) - else [data] - ) - - log.debug("Received: %s", data) - - for msg in messages: - if msg.seq_no % 2 != 0: - if msg.msg_id in self.pending_acks: - continue - else: - self.pending_acks.add(msg.msg_id) - - try: - if len(self.stored_msg_ids) > Session.STORED_MSG_IDS_MAX_SIZE: - del self.stored_msg_ids[:Session.STORED_MSG_IDS_MAX_SIZE // 2] - - if self.stored_msg_ids: - if msg.msg_id < self.stored_msg_ids[0]: - raise SecurityCheckMismatch("The msg_id is lower than all the stored values") - - if msg.msg_id in self.stored_msg_ids: - raise SecurityCheckMismatch("The msg_id is equal to any of the stored values") - - time_diff = (msg.msg_id - MsgId()) / 2 ** 32 - - if time_diff > 30: - raise SecurityCheckMismatch("The msg_id belongs to over 30 seconds in the future. " - "Most likely the client time has to be synchronized.") - - if time_diff < -300: - raise SecurityCheckMismatch("The msg_id belongs to over 300 seconds in the past. " - "Most likely the client time has to be synchronized.") - except SecurityCheckMismatch as e: - log.info("Discarding packet: %s", e) - await self.connection.close() - return - else: - bisect.insort(self.stored_msg_ids, msg.msg_id) - - if isinstance(msg.body, (raw.types.MsgDetailedInfo, raw.types.MsgNewDetailedInfo)): - self.pending_acks.add(msg.body.answer_msg_id) - continue - - if isinstance(msg.body, raw.types.NewSessionCreated): - continue - - msg_id = None - - if isinstance(msg.body, (raw.types.BadMsgNotification, raw.types.BadServerSalt)): - msg_id = msg.body.bad_msg_id - elif isinstance(msg.body, (FutureSalts, raw.types.RpcResult)): - msg_id = msg.body.req_msg_id - elif isinstance(msg.body, raw.types.Pong): - msg_id = msg.body.msg_id - else: - if self.client is not None: - self.client.loop.create_task(self.client.handle_updates(msg.body)) - - if msg_id in self.results: - self.results[msg_id].value = getattr(msg.body, "result", msg.body) - self.results[msg_id].event.set() - - if len(self.pending_acks) >= self.ACKS_THRESHOLD: - log.debug("Sending %s acks", len(self.pending_acks)) - - try: - await self.send(raw.types.MsgsAck(msg_ids=list(self.pending_acks)), False) - except OSError: - pass - else: - self.pending_acks.clear() - - async def ping_worker(self): - log.info("PingTask started") - - while True: - try: - await asyncio.wait_for(self.ping_task_event.wait(), self.PING_INTERVAL) - except asyncio.TimeoutError: - pass - else: - break - - try: - await self.send( - raw.functions.PingDelayDisconnect( - ping_id=0, disconnect_delay=self.WAIT_TIMEOUT + 10 - ), False - ) - except OSError: - self.client.loop.create_task(self.restart()) - break - except RPCError: - pass - - log.info("PingTask stopped") - - async def recv_worker(self): - log.info("NetworkTask started") - - while True: - try: - packet = await asyncio.wait_for(self.connection.recv(), timeout=1) - except asyncio.TimeoutError: - continue - except ConnectionError: - if self.is_started.is_set(): - self.client.loop.create_task(self.restart()) - break - - if packet is None or len(packet) == 4: - if packet: - error_code = -Int.read(BytesIO(packet)) - - if error_code == 404: - raise Unauthorized( - "Auth key not found in the system. You must delete your session file " - "and log in again with your phone number or bot token." - ) - - log.warning( - "Server sent transport error: %s (%s)", - error_code, Session.TRANSPORT_ERRORS.get(error_code, "unknown error") - ) - - if self.is_started.is_set(): - self.client.loop.create_task(self.restart()) - - break - - self.client.loop.create_task(self.handle_packet(packet)) - - log.info("NetworkTask stopped") - - async def send(self, data: TLObject, wait_response: bool = True, timeout: float = WAIT_TIMEOUT): - message = self.msg_factory(data) - msg_id = message.msg_id - - if wait_response: - self.results[msg_id] = Result() - - log.debug("Sent: %s", message) - - payload = await self.client.loop.run_in_executor( - pyrogram.crypto_executor, - mtproto.pack, - message, - self.salt, - self.session_id, - self.auth_key, - self.auth_key_id - ) - - try: - await self.connection.send(payload) - except OSError as e: - self.results.pop(msg_id, None) - raise e - - if wait_response: - try: - await asyncio.wait_for(self.results[msg_id].event.wait(), timeout) - except asyncio.TimeoutError: - pass - - result = self.results.pop(msg_id).value - - if result is None: - raise TimeoutError("Request timed out") - - if isinstance(result, raw.types.RpcError): - if isinstance(data, (raw.functions.InvokeWithoutUpdates, raw.functions.InvokeWithTakeout)): - data = data.query - - RPCError.raise_it(result, type(data)) - - if isinstance(result, raw.types.BadMsgNotification): - log.warning("%s: %s", BadMsgNotification.__name__, BadMsgNotification(result.error_code)) - - if isinstance(result, raw.types.BadServerSalt): - self.salt = result.new_server_salt - return await self.send(data, wait_response, timeout) - - return result - - async def invoke( - self, - query: TLObject, - retries: int = MAX_RETRIES, - timeout: float = WAIT_TIMEOUT, - sleep_threshold: float = SLEEP_THRESHOLD - ): - try: - await asyncio.wait_for(self.is_started.wait(), self.WAIT_TIMEOUT) - except asyncio.TimeoutError: - pass - - if isinstance(query, (raw.functions.InvokeWithoutUpdates, raw.functions.InvokeWithTakeout)): - inner_query = query.query - else: - inner_query = query - - query_name = ".".join(inner_query.QUALNAME.split(".")[1:]) - - while True: - try: - return await self.send(query, timeout=timeout) - except (FloodWait, FloodPremiumWait) as e: - amount = e.value - - if amount > sleep_threshold >= 0: - raise - - log.warning('[%s] Waiting for %s seconds before continuing (required by "%s")', - self.client.name, amount, query_name) - - await asyncio.sleep(amount) - except (OSError, InternalServerError, ServiceUnavailable) as e: - if retries == 0: - raise e from None - - (log.warning if retries < 2 else log.info)( - '[%s] Retrying "%s" due to: %s. Performing a restart.', - Session.MAX_RETRIES - retries + 1, - query_name, str(e) or repr(e) - ) - - await self.restart() - retries -= 1 + START_TIMEOUT = 2 + WAIT_TIMEOUT = 15 + SLEEP_THRESHOLD = 10 + MAX_RETRIES = 10 + ACKS_THRESHOLD = 10 + PING_INTERVAL = 5 + STORED_MSG_IDS_MAX_SIZE = 1000 * 2 + + TRANSPORT_ERRORS = { + 404: "auth key not found", + 429: "transport flood", + 444: "invalid DC" + } + + def __init__( + self, + client: "pyrogram.Client", + dc_id: int, + auth_key: bytes, + test_mode: bool, + is_media: bool = False, + is_cdn: bool = False + ): + self.client = client + self.dc_id = dc_id + self.auth_key = auth_key + self.test_mode = test_mode + self.is_media = is_media + self.is_cdn = is_cdn + + self.connection: Optional[Connection] = None + + self.auth_key_id = sha1(auth_key).digest()[-8:] + + self.session_id = os.urandom(8) + self.msg_factory = MsgFactory() + + self.salt = 0 + + self.pending_acks = set() + + self.results = {} + + self.stored_msg_ids = [] + + self.ping_task = None + self.ping_task_event = asyncio.Event() + + self.recv_task = None + + self.is_started = asyncio.Event() + self.restart_lock = asyncio.Lock() + + async def start(self): + while True: + self.connection = self.client.connection_factory( + dc_id=self.dc_id, + test_mode=self.test_mode, + ipv6=self.client.ipv6, + proxy=self.client.proxy, + media=self.is_media, + protocol_factory=self.client.protocol_factory, + loop=self.client.loop + ) + + try: + await self.connection.connect() + + self.recv_task = self.client.loop.create_task(self.recv_worker()) + + await self.send(raw.functions.Ping(ping_id=0), timeout=self.START_TIMEOUT) + + if not self.is_cdn: + await self.send( + raw.functions.InvokeWithLayer( + layer=layer, + query=raw.functions.InitConnection( + api_id=await self.client.storage.api_id(), + app_version=self.client.app_version, + device_model=self.client.device_model, + system_version=self.client.system_version, + system_lang_code=self.client.system_lang_code, + lang_pack=self.client.lang_pack, + lang_code=self.client.lang_code, + query=raw.functions.help.GetConfig(), + params=self.client.init_connection_params, + ) + ), + timeout=self.START_TIMEOUT + ) + + self.ping_task = self.client.loop.create_task(self.ping_worker()) + + log.info("Session initialized: Pyrogram v%s (Layer %s)", pyrogram.__version__, layer) + log.info("Device: %s - %s", self.client.device_model, self.client.app_version) + log.info("System: %s (%s)", self.client.system_version, self.client.lang_code) + except (AuthKeyDuplicated, Unauthorized) as e: + await self.stop() + raise e + except ConnectionError as e: + await self.stop() + except (OSError, RPCError): + await self.stop() + except Exception as e: + await self.stop() + raise e + else: + break + + self.is_started.set() + log.info("Session started") + + if callable(self.client.connect_handler): + try: + await self.client.connect_handler(self.client, self) + except Exception as e: + log.exception(e) + + async def stop(self): + if callable(self.client.disconnect_handler): + try: + await self.client.disconnect_handler(self.client, self) + except Exception as e: + log.exception(e) + + self.is_started.clear() + self.stored_msg_ids.clear() + self.ping_task_event.set() + + if self.ping_task is not None: + await self.ping_task + + self.ping_task_event.clear() + + if self.connection: + await self.connection.close() + + if self.recv_task: + try: + await self.recv_task + except asyncio.CancelledError: + pass + self.recv_task = None + + log.info("Session stopped") + + async def restart(self): + async with self.restart_lock: + if self.is_started.is_set(): + await self.stop() + await self.start() + + async def handle_packet(self, packet): + try: + data = await self.client.loop.run_in_executor( + pyrogram.crypto_executor, + mtproto.unpack, + BytesIO(packet), + self.session_id, + self.auth_key, + self.auth_key_id + ) + except ValueError as e: + log.debug(e) + await self.restart() + return + + messages = ( + data.body.messages + if isinstance(data.body, MsgContainer) + else [data] + ) + + log.debug("Received: %s", data) + + for msg in messages: + if msg.seq_no % 2 != 0: + if msg.msg_id in self.pending_acks: + continue + else: + self.pending_acks.add(msg.msg_id) + try: + if len(self.stored_msg_ids) > Session.STORED_MSG_IDS_MAX_SIZE: + del self.stored_msg_ids[:Session.STORED_MSG_IDS_MAX_SIZE // 2] + if self.stored_msg_ids: + if msg.msg_id < self.stored_msg_ids[0]: + raise SecurityCheckMismatch("The msg_id is lower than all the stored values") + if msg.msg_id in self.stored_msg_ids: + raise SecurityCheckMismatch("The msg_id is equal to any of the stored values") + + time_diff = (msg.msg_id - MsgId()) / 2 ** 32 + if time_diff > 30: + raise SecurityCheckMismatch("The msg_id belongs to over 30 seconds in the future. " + "Most likely the client time has to be synchronized.") + if time_diff < -300: + raise SecurityCheckMismatch("The msg_id belongs to over 300 seconds in the past. " + "Most likely the client time has to be synchronized.") + except SecurityCheckMismatch as e: + log.info("Discarding packet: %s", e) + await self.connection.close() + return + else: + bisect.insort(self.stored_msg_ids, msg.msg_id) + + if isinstance(msg.body, (raw.types.MsgDetailedInfo, raw.types.MsgNewDetailedInfo)): + self.pending_acks.add(msg.body.answer_msg_id) + continue + + if isinstance(msg.body, raw.types.NewSessionCreated): + continue + + msg_id = None + + if isinstance(msg.body, (raw.types.BadMsgNotification, raw.types.BadServerSalt)): + msg_id = msg.body.bad_msg_id + elif isinstance(msg.body, (FutureSalts, raw.types.RpcResult)): + msg_id = msg.body.req_msg_id + elif isinstance(msg.body, raw.types.Pong): + msg_id = msg.body.msg_id + else: + if self.client is not None: + self.client.loop.create_task(self.client.handle_updates(msg.body)) + + if msg_id in self.results: + self.results[msg_id].value = getattr(msg.body, "result", msg.body) + self.results[msg_id].event.set() + + if len(self.pending_acks) >= self.ACKS_THRESHOLD: + log.debug("Sending %s acks", len(self.pending_acks)) + try: + await self.send(raw.types.MsgsAck(msg_ids=list(self.pending_acks)), False) + except OSError: + pass + else: + self.pending_acks.clear() + + async def ping_worker(self): + log.info("PingTask started") + + while True: + try: + await asyncio.wait_for(self.ping_task_event.wait(), self.PING_INTERVAL) + except asyncio.TimeoutError: + pass + else: + break + try: + await self.send( + raw.functions.PingDelayDisconnect( + ping_id=0, disconnect_delay=self.WAIT_TIMEOUT + 10 + ), False + ) + except OSError: + log.warning("Ping failed, restarting session...") + await self.restart() + continue + + except RPCError: + pass + + log.info("PingTask stopped") + + async def recv_worker(self): + log.info("NetworkTask started") + + while True: + try: + packet = await asyncio.wait_for(self.connection.recv(), timeout=1) + except asyncio.TimeoutError: + continue + except ConnectionError: + break + + if packet is None or len(packet) == 4: + if packet: + error_code = -Int.read(BytesIO(packet)) + + if error_code == 404: + raise Unauthorized( + "Auth key not found in the system. You must delete your session file " + "and log in again with your phone number or bot token." + ) + + log.warning( + "Server sent transport error: %s (%s)", + error_code, Session.TRANSPORT_ERRORS.get(error_code, "unknown error") + ) + + self.client.loop.create_task(self.handle_packet(packet)) + + log.info("NetworkTask stopped") + + async def send(self, data: TLObject, wait_response: bool = True, timeout: float = WAIT_TIMEOUT): + message = self.msg_factory(data) + msg_id = message.msg_id + + if wait_response: + self.results[msg_id] = Result() + + log.debug("Sent: %s", message) + + payload = await self.client.loop.run_in_executor( + pyrogram.crypto_executor, + mtproto.pack, + message, + self.salt, + self.session_id, + self.auth_key, + self.auth_key_id + ) + + try: + await self.connection.send(payload) + except OSError as e: + self.results.pop(msg_id, None) + raise e + + if wait_response: + try: + await asyncio.wait_for(self.results[msg_id].event.wait(), timeout) + except asyncio.TimeoutError: + pass + + result = self.results.pop(msg_id).value + + if result is None: + raise TimeoutError("Request timed out") + + if isinstance(result, raw.types.RpcError): + if isinstance(data, (raw.functions.InvokeWithoutUpdates, raw.functions.InvokeWithTakeout)): + data = data.query + + RPCError.raise_it(result, type(data)) + + if isinstance(result, raw.types.BadMsgNotification): + log.warning("%s: %s", BadMsgNotification.__name__, BadMsgNotification(result.error_code)) + + if isinstance(result, raw.types.BadServerSalt): + self.salt = result.new_server_salt + return await self.send(data, wait_response, timeout) + + return result + + async def invoke( + self, + query: TLObject, + retries: int = MAX_RETRIES, + timeout: float = WAIT_TIMEOUT, + sleep_threshold: float = SLEEP_THRESHOLD + ): + try: + await asyncio.wait_for(self.is_started.wait(), self.WAIT_TIMEOUT) + except asyncio.TimeoutError: + pass + + if isinstance(query, (raw.functions.InvokeWithoutUpdates, raw.functions.InvokeWithTakeout)): + inner_query = query.query + else: + inner_query = query + + query_name = ".".join(inner_query.QUALNAME.split(".")[1:]) + + for attempt in range(1, retries + 1): + try: + return await self.send(query, timeout=timeout) + except (FloodWait, FloodPremiumWait) as e: + amount = e.value + + if amount > sleep_threshold >= 0: + raise + + log.warning('[%s] Waiting for %s seconds before continuing (required by "%s")', + self.client.name, amount, query_name) + + await asyncio.sleep(amount) + except (OSError, InternalServerError, ServiceUnavailable) as e: + if attempt >= retries: + raise e from None + + (log.warning if attempt >= retries - 1 else log.info)( + '[%s] Retrying "%s" due to: %s', + attempt, + query_name, str(e) or repr(e) + ) + + await self.restart() + + raise TimeoutError(f'Failed to invoke "{query_name}" after {retries} retries')