diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py index c84a3cb4..e60d48c6 100644 --- a/pyrogram/session/session.py +++ b/pyrogram/session/session.py @@ -21,6 +21,7 @@ import asyncio import bisect import logging import os +from datetime import datetime, timedelta from hashlib import sha1 from io import BytesIO from typing import Optional @@ -56,6 +57,7 @@ class Session: ACKS_THRESHOLD = 10 PING_INTERVAL = 5 STORED_MSG_IDS_MAX_SIZE = 500 + RECONNECT_THRESHOLD = timedelta(seconds=10) TRANSPORT_ERRORS = { 404: "auth key not found", @@ -103,6 +105,8 @@ class Session: self.loop = asyncio.get_event_loop() + self.last_reconnect_attempt = None + async def start(self): while True: self.connection = self.client.connection_factory( @@ -186,6 +190,15 @@ class Session: log.info("Session stopped") async def restart(self): + now = datetime.now() + if ( + self.last_reconnect_attempt + and now - self.last_reconnect_attempt < self.RECONNECT_THRESHOLD + ): + log.info("Reconnecting too frequently, sleeping for a while") + await asyncio.sleep(5) + + self.last_reconnect_attempt = now await self.stop() await self.start() @@ -415,7 +428,7 @@ class Session: query_name = ".".join(inner_query.QUALNAME.split(".")[1:]) - while True: + while retries > 0: try: return await self.send(query, timeout=timeout) except (FloodWait, FloodPremiumWait) as e: @@ -429,15 +442,16 @@ class Session: await asyncio.sleep(amount) except (OSError, InternalServerError, ServiceUnavailable) as e: + retries -= 1 if retries == 0: - raise e from None + raise e (log.warning if retries < 2 else log.info)( '[%s] Retrying "%s" due to: %s', - Session.MAX_RETRIES - retries + 1, + Session.MAX_RETRIES - retries, query_name, str(e) or repr(e) ) await asyncio.sleep(0.5) - return await self.invoke(query, retries - 1, timeout) + raise TimeoutError("Exceeded maximum number of retries")