fix(session): handle fatal errors and improve session concurrency

Author: SychO3

pyrogram/client.py

@@ -17,6 +17,7 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
+from contextlib import suppress
 import functools
 import inspect
 import logging
@@ -395,6 +396,7 @@ def __init__(
         self.sessions = {}
         self.media_sessions = {}
         self.sessions_lock = asyncio.Lock()
+        self._session_futures = {}
 
         self.save_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions)
         self.get_file_semaphore = asyncio.Semaphore(self.max_concurrent_transmissions)
@@ -1371,67 +1373,121 @@ async def get_session(
             return self.session
 
         sessions = self.media_sessions if is_media else self.sessions
+        session_key = (dc_id, is_media)
 
-        if not temporary and sessions.get(dc_id):
-            return sessions[dc_id]
+        creator = False
 
-        if not server_address or not port:
-            dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6, is_cdn=is_cdn)
+        if not temporary:
+            async with self.sessions_lock:
+                if sessions.get(dc_id):
+                    return sessions[dc_id]
 
-            server_address = server_address or dc_option.ip_address
-            port = port or dc_option.port
+                pending_session = self._session_futures.get(session_key)
 
-        if is_media:
-            auth_key = (await self.get_session(dc_id)).auth_key
-        else:
-            if not is_current_dc:
-                auth_key = await Auth(
-                    self,
-                    dc_id,
-                    server_address,
-                    port,
-                    await self.storage.test_mode()
-                ).create()
-            else:
-                auth_key = await self.storage.auth_key()
-
-        session = Session(
-            self,
-            dc_id,
-            server_address,
-            port,
-            auth_key,
-            await self.storage.test_mode(),
-            is_media=is_media
-        )
+                if pending_session is None:
+                    pending_session = self.loop.create_future()
+                    pending_session.add_done_callback(
+                        lambda future: None if future.cancelled() else future.exception()
+                    )
+                    self._session_futures[session_key] = pending_session
+                    creator = True
 
-        if not temporary:
-            sessions[dc_id] = session
+            if not creator:
+                return await pending_session
 
-        await session.start()
+        session = None
 
-        if not is_current_dc and export_authorization:
-            for _ in range(3):
-                exported_auth = await self.invoke(
-                    raw.functions.auth.ExportAuthorization(
-                        dc_id=dc_id
-                    )
-                )
+        try:
+            if not server_address or not port:
+                dc_option = await self.get_dc_option(dc_id, is_media=is_media, ipv6=self.ipv6, is_cdn=is_cdn)
 
-                try:
-                    await session.invoke(
-                        raw.functions.auth.ImportAuthorization(
-                            id=exported_auth.id,
-                            bytes=exported_auth.bytes
+                server_address = server_address or dc_option.ip_address
+                port = port or dc_option.port
+
+            if is_media:
+                auth_key = (await self.get_session(dc_id)).auth_key
+            else:
+                if not is_current_dc:
+                    auth_key = await Auth(
+                        self,
+                        dc_id,
+                        server_address,
+                        port,
+                        await self.storage.test_mode()
+                    ).create()
+                else:
+                    auth_key = await self.storage.auth_key()
+
+            session = Session(
+                self,
+                dc_id,
+                server_address,
+                port,
+                auth_key,
+                await self.storage.test_mode(),
+                is_media=is_media
+            )
+
+            await session.start()
+
+            try:
+                await asyncio.wait_for(session.is_started.wait(), Session.WAIT_TIMEOUT)
+            except asyncio.TimeoutError as e:
+                with suppress(Exception):
+                    await session.stop()
+                session = None
+                raise ConnectionError(f"Failed to start session for DC{dc_id}") from e
+
+            if not is_current_dc and export_authorization:
+                for _ in range(3):
+                    exported_auth = await self.invoke(
+                        raw.functions.auth.ExportAuthorization(
+                            dc_id=dc_id
                         )
                     )
-                except AuthBytesInvalid:
-                    continue
+
+                    try:
+                        await session.invoke(
+                            raw.functions.auth.ImportAuthorization(
+                                id=exported_auth.id,
+                                bytes=exported_auth.bytes
+                            )
+                        )
+                    except AuthBytesInvalid:
+                        continue
+                    else:
+                        break
                 else:
-                    break
-            else:
-                await session.stop()
-                raise AuthBytesInvalid
+                    await session.stop()
+                    session = None
+                    raise AuthBytesInvalid
+
+            if not temporary:
+                async with self.sessions_lock:
+                    cached_session = sessions.get(dc_id)
+
+                    if cached_session is None:
+                        sessions[dc_id] = session
+                    else:
+                        session = cached_session
+
+                    pending_session = self._session_futures.pop(session_key, None)
+
+                    if pending_session is not None and not pending_session.done():
+                        pending_session.set_result(session)
+        except Exception as e:
+            if session is not None:
+                with suppress(Exception):
+                    await session.stop()
+
+            if not temporary:
+                async with self.sessions_lock:
+                    pending_session = self._session_futures.pop(session_key, None)
+
+                    if pending_session is not None and not pending_session.done():
+                        pending_session.set_exception(e)
+
+            raise
 
         return session
 

pyrogram/methods/auth/send_phone_number_code.py

@@ -22,7 +22,7 @@
 
 import pyrogram
 from pyrogram import enums, raw, types
-from pyrogram.errors import NetworkMigrate, PhoneMigrate
+from pyrogram.errors import NetworkMigrateX, PhoneMigrateX  
 
 log = logging.getLogger(__name__)
 
@@ -165,7 +165,7 @@ async def send_phone_number_code(
                     )
 
                 r = await self.invoke(rpc, recaptcha_token=recaptcha_token)
-            except (PhoneMigrate, NetworkMigrate) as e:
+            except (PhoneMigrateX, NetworkMigrateX) as e:
                 dc_option = await self.get_dc_option(e.value, ipv6=self.ipv6)
                 await self.session.stop()
 

pyrogram/session/session.py

@@ -140,6 +140,7 @@ def __init__(
 
         self.is_started = asyncio.Event()
         self.restart_lock = asyncio.Lock()
+        self.fatal_error: Optional[BaseException] = None
 
     @property
     def state(self) -> SessionState:
@@ -160,6 +161,7 @@ async def start(self):
             return
 
         await self._set_state(SessionState.STARTING)
+        self.fatal_error = None
 
         self.connection = self.client.connection_factory(
             dc_id=self.dc_id,
@@ -289,6 +291,16 @@ async def restart(self):
             await self.stop()
             await self.start()
 
+    def _fail_pending_results(self, error: BaseException) -> None:
+        for result in self.results.values():
+            if result.value is None:
+                result.value = error
+                result.event.set()
+
+    def _set_fatal_error(self, error: BaseException) -> None:
+        self.fatal_error = error
+        self._fail_pending_results(error)
+
     async def handle_packet(self, packet):
         try:
             data = await self.client.loop.run_in_executor(
@@ -299,7 +311,7 @@ async def handle_packet(self, packet):
                 self.auth_key,
                 self.auth_key_id
             )
-        except ValueError as e:
+        except (ConnectionError, SecurityCheckMismatch, ValueError) as e:
             log.debug(e)
             log.info("Restarting session due to - %s - %s", e.__class__.__name__, e)
             self.client.loop.create_task(self.restart())
@@ -443,9 +455,10 @@ async def recv_worker(self):
                 if packet:
                     error_code = -Int.read(BytesIO(packet))
                     error_msg = "unknown error"
+                    transport_error = None
 
                     if error_code == 404:
-                        raise AuthKeyNotFound(
+                        transport_error = AuthKeyNotFound(
                             "Auth key not found in the system. Try again or delete your session file "
                             "and log in again with your phone number or bot token."
                         )
@@ -460,10 +473,19 @@ async def recv_worker(self):
                                 "Invalid data center. Please check your configuration."
                             )
                     except TransportError as e:
+                        transport_error = e
                         error_msg = str(e)
+                    else:
+                        if transport_error is not None:
+                            error_msg = str(transport_error)
 
                     log.warning("Server sent transport error: %s (%s)", error_code, error_msg)
 
+                    if isinstance(transport_error, AuthKeyNotFound):
+                        self._set_fatal_error(transport_error)
+                        self.client.loop.create_task(self.stop())
+                        break
+
 
                 if self.is_started.is_set():
                     if packet:
@@ -483,6 +505,9 @@ async def recv_worker(self):
     async def send(
         self, data: TLObject, wait_response: bool = True, timeout: float = WAIT_TIMEOUT
     ):
+        if self.fatal_error is not None:
+            raise self.fatal_error
+
         message = await self.msg_factory.create(data)
         msg_id = message.msg_id
 
@@ -518,6 +543,9 @@ async def send(
             if result is None:
                 raise TimeoutError("Request timed out")
 
+            if isinstance(result, BaseException):
+                raise result
+
             if isinstance(result, raw.types.RpcError):
                 if isinstance(
                     data, (raw.functions.InvokeWithoutUpdates, raw.functions.InvokeWithTakeout)