fix: handle connection closure and retry logic in session management

HitaloM · HitaloM · commit 08a756a00933 · 2025-04-17T13:49:26.000-03:00
diff --git a/hydrogram/connection/transport/tcp/tcp.py b/hydrogram/connection/transport/tcp/tcp.py
@@ -56,6 +56,13 @@ def __init__(self, ipv6: bool, proxy: Proxy) -> None:
 
         self.lock = asyncio.Lock()
         self.loop = asyncio.get_running_loop()
+        self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        return (
+            self._closed or self.writer is None or self.writer.is_closing() or self.reader is None
+        )
 
     async def _connect_via_proxy(self, destination: tuple[str, int]) -> None:
         scheme = self.proxy.get("scheme")
@@ -108,45 +115,56 @@ async def _connect(self, destination: tuple[str, int]) -> None:
     async def connect(self, address: tuple[str, int]) -> None:
         try:
             await asyncio.wait_for(self._connect(address), TCP.TIMEOUT)
+            self._closed = False
         except (
             asyncio.TimeoutError
         ):  # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11
+            self._closed = True
             raise TimeoutError("Connection timed out")
 
     async def close(self) -> None:
         if self.writer is None:
+            self._closed = True
             return
 
         try:
             self.writer.close()
             await asyncio.wait_for(self.writer.wait_closed(), TCP.TIMEOUT)
         except Exception as e:
             log.info("Close exception: %s %s", type(e).__name__, e)
+        finally:
+            self._closed = True
 
     async def send(self, data: bytes) -> None:
-        if self.writer is None:
-            return
+        if self.writer is None or self._closed:
+            raise OSError("Connection is closed")
 
         async with self.lock:
             try:
                 self.writer.write(data)
                 await self.writer.drain()
             except Exception as e:
                 log.info("Send exception: %s %s", type(e).__name__, e)
+                self._closed = True
                 raise OSError(e) from e
 
     async def recv(self, length: int = 0) -> bytes | None:
+        if self._closed or self.reader is None:
+            return None
+
         data = b""
 
         while len(data) < length:
             try:
                 chunk = await asyncio.wait_for(self.reader.read(length - len(data)), TCP.TIMEOUT)
             except (OSError, asyncio.TimeoutError):
+                self._closed = True
                 return None
             else:
                 if chunk:
                     data += chunk
                 else:
+                    self._closed = True
                     return None
 
         return data
diff --git a/hydrogram/session/session.py b/hydrogram/session/session.py
@@ -424,6 +424,19 @@ async def invoke(
 
         while retries > 0:
             try:
+                if (
+                    self.connection is None
+                    or self.connection.protocol is None
+                    or getattr(self.connection.protocol, "closed", True)
+                ):
+                    log.warning(
+                        "[%s] Connection is closed or not established. Attempting to reconnect...",
+                        self.client.name,
+                    )
+                    await self.restart()
+                    await asyncio.sleep(1)
+                    continue
+
                 return await self.send(query, timeout=timeout)
             except FloodWait as e:
                 amount = e.value
@@ -451,6 +464,16 @@ async def invoke(
                     str(e) or repr(e),
                 )
 
+                if isinstance(e, OSError) and retries > 1:
+                    try:
+                        await self.restart()
+                    except Exception as restart_error:
+                        log.warning(
+                            "[%s] Failed to restart session: %s",
+                            self.client.name,
+                            str(restart_error) or repr(restart_error),
+                        )
+
                 await asyncio.sleep(0.5)
 
         raise TimeoutError("Exceeded maximum number of retries")
