fix(session): ensure safe closure of connection and cancellation of receive task

HitaloM · HitaloM · commit 58c173064faf · 2025-07-25T21:08:39.000-03:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -20,8 +20,8 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.7
+    rev: v0.12.5
     hooks:
+      - id: ruff-check
+        args: [--fix, --ignore=E501, --unsafe-fixes, --exit-non-zero-on-fix]
       - id: ruff-format
-      - id: ruff
-        args: [--fix, --ignore=E501, --exit-non-zero-on-fix]
diff --git a/hydrogram/session/session.py b/hydrogram/session/session.py
@@ -111,6 +111,7 @@ def __init__(
         self.recv_task = None
 
         self.is_started = asyncio.Event()
+        self.restart_lock = asyncio.Lock()
 
         self.last_reconnect_attempt = None
 
@@ -182,12 +183,16 @@ async def stop(self):
 
         self.ping_task_event.clear()
 
-        await self.connection.close()
+        if self.connection:
+            await self.connection.close()
 
-        if self.recv_task:
+        if self.recv_task and not self.recv_task.done():
             self.recv_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self.recv_task
+
+            with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError, RuntimeError):
+                await asyncio.wait_for(self.recv_task, timeout=1.0)
+
+            self.recv_task = None
 
         if not self.is_media and callable(self.client.disconnect_handler):
             try:
@@ -198,17 +203,18 @@ 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()
+        async with self.restart_lock:
+            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()
 
     async def handle_packet(self, packet):
         data = await self.client.loop.run_in_executor(
diff --git a/uv.lock b/uv.lock
