refactor(session): replace recursion with loop and add backoff

HitaloM · HitaloM · commit fd5beaad8865 · 2024-05-29T22:45:16.000-03:00
This refactor replaces recursion with a loop in the session invoke logic. Additionally, a backoff mechanism has been introduced to prevent frequent restarts from crashing the bot.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -11,7 +11,7 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.5
+    rev: v0.4.6
     hooks:
       - id: ruff-format
       - id: ruff
diff --git a/hydrogram/session/session.py b/hydrogram/session/session.py
@@ -22,6 +22,7 @@
 import contextlib
 import logging
 import os
+from datetime import datetime, timedelta
 from hashlib import sha1
 from io import BytesIO
 from typing import ClassVar
@@ -61,6 +62,7 @@ class Session:
     ACKS_THRESHOLD = 10
     PING_INTERVAL = 5
     STORED_MSG_IDS_MAX_SIZE = 1000 * 2
+    RECONNECT_THRESHOLD = timedelta(seconds=10)
 
     TRANSPORT_ERRORS: ClassVar = {
         404: "auth key not found",
@@ -108,6 +110,8 @@ def __init__(
 
         self.loop = asyncio.get_event_loop()
 
+        self.last_reconnect_attempt = None
+
     async def start(self):
         while True:
             self.connection = Connection(
@@ -189,6 +193,15 @@ async def stop(self):
         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()
 
@@ -406,7 +419,7 @@ async def invoke(
 
         query_name = ".".join(inner_query.QUALNAME.split(".")[1:])
 
-        while True:
+        while retries > 0:
             try:
                 return await self.send(query, timeout=timeout)
             except FloodWait as e:
@@ -424,16 +437,17 @@ async def invoke(
 
                 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")
diff --git a/hydrogram/types/user_and_chats/chat.py b/hydrogram/types/user_and_chats/chat.py
@@ -276,37 +276,28 @@ def _parse_chat_chat(client, chat: raw.types.Chat) -> Chat:
     @staticmethod
     def _parse_channel_chat(client, channel: raw.types.Channel) -> Chat:
         peer_id = utils.get_channel_id(channel.id)
-        restriction_reason = getattr(channel, "restriction_reason", [])
-        usernames = getattr(channel, "usernames", [])
 
         return Chat(
             id=peer_id,
-            type=enums.ChatType.SUPERGROUP
-            if getattr(channel, "megagroup", None)
-            else enums.ChatType.CHANNEL,
-            is_verified=getattr(channel, "verified", None),
-            is_restricted=getattr(channel, "restricted", None),
-            is_creator=getattr(channel, "creator", None),
-            is_scam=getattr(channel, "scam", None),
-            is_fake=getattr(channel, "fake", None),
-            is_forum=getattr(channel, "forum", None),
+            type=enums.ChatType.SUPERGROUP if channel.megagroup else enums.ChatType.CHANNEL,
+            is_verified=channel.verified,
+            is_restricted=channel.restricted,
+            is_creator=channel.creator,
+            is_scam=channel.scam,
+            is_fake=channel.fake,
+            is_forum=channel.forum,
             title=channel.title,
-            username=getattr(channel, "username", None),
-            photo=types.ChatPhoto._parse(
-                client,
-                getattr(channel, "photo", None),
-                peer_id,
-                getattr(channel, "access_hash", 0),
-            ),
-            restrictions=types.List([types.Restriction._parse(r) for r in restriction_reason])
+            username=channel.username,
+            photo=types.ChatPhoto._parse(client, channel.photo, peer_id, channel.access_hash),
+            restrictions=types.List([
+                types.Restriction._parse(r) for r in channel.restriction_reason
+            ])
             or None,
-            permissions=types.ChatPermissions._parse(
-                getattr(channel, "default_banned_rights", None)
-            ),
-            members_count=getattr(channel, "participants_count", None),
-            dc_id=getattr(getattr(channel, "photo", None), "dc_id", None),
-            has_protected_content=getattr(channel, "noforwards", None),
-            usernames=types.List([types.Username._parse(r) for r in usernames]) or None,
+            permissions=types.ChatPermissions._parse(channel.default_banned_rights),
+            members_count=channel.participants_count,
+            dc_id=getattr(channel.photo, "dc_id", None),
+            has_protected_content=channel.noforwards,
+            usernames=types.List([types.Username._parse(r) for r in channel.usernames]) or None,
             client=client,
         )
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -51,7 +51,7 @@ build-backend = "hatchling.build"
 [tool.rye]
 managed = true
 dev-dependencies = [
-    "ruff>=0.4.5",
+    "ruff>=0.4.6",
     "pytest>=7.4.3",
     "pytest-asyncio>=0.23.2",
     "pytest-cov>=4.1.0",
