Refactor connection

Author: KurimuzonAkuma

pyrogram/__init__.py

@@ -38,5 +38,3 @@ class ContinuePropagation(StopAsyncIteration):
 from . import raw, types, filters, handlers, enums
 from .client import Client
 from .sync import idle, compose
-
-crypto_executor = ThreadPoolExecutor(1, thread_name_prefix="CryptoWorker")

pyrogram/connection/connection.py

@@ -37,6 +37,7 @@ def __init__(
         proxy: dict,
         media: bool = False,
         protocol_factory: Type[TCP] = TCPAbridged,
+        crypto_executor_workers: int = 1,
         loop: Optional[asyncio.AbstractEventLoop] = None
     ) -> None:
         self.dc_id = dc_id
@@ -45,6 +46,7 @@ def __init__(
         self.proxy = proxy
         self.media = media
         self.protocol_factory = protocol_factory
+        self.crypto_executor_workers = crypto_executor_workers
 
         self.address = DataCenter(dc_id, test_mode, ipv6, media)
         self.protocol: Optional[TCP] = None
@@ -56,7 +58,7 @@ def __init__(
 
     async def connect(self) -> None:
         for i in range(Connection.MAX_CONNECTION_ATTEMPTS):
-            self.protocol = self.protocol_factory(ipv6=self.ipv6, proxy=self.proxy, loop=self.loop)
+            self.protocol = self.protocol_factory(ipv6=self.ipv6, proxy=self.proxy, crypto_executor_workers=self.crypto_executor_workers, loop=self.loop)
 
             try:
                 log.info("Connecting...")

pyrogram/connection/transport/tcp/tcp.py

@@ -21,7 +21,7 @@
 import logging
 import socket
 from concurrent.futures import ThreadPoolExecutor
-from typing import Tuple, Dict, TypedDict, Optional
+from typing import Dict, Optional, Tuple, TypedDict
 
 import socks
 
@@ -45,13 +45,24 @@ class Proxy(TypedDict):
 class TCP:
     TIMEOUT = 10
 
-    def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
+    def __init__(
+        self,
+        ipv6: bool = False,
+        proxy: Proxy = None,
+        crypto_executor_workers: int = 1,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
         self.ipv6 = ipv6
         self.proxy = proxy
 
+        self.crypto_executor_workers = crypto_executor_workers
+        self.crypto_executor = ThreadPoolExecutor(max_workers=self.crypto_executor_workers, thread_name_prefix="CryptoWorker")
+
         self.reader: Optional[asyncio.StreamReader] = None
         self.writer: Optional[asyncio.StreamWriter] = None
 
+        self.marker_event = asyncio.Event()
+
         self.lock = asyncio.Lock()
 
         if isinstance(loop, asyncio.AbstractEventLoop):
@@ -95,8 +106,7 @@ async def _connect_via_proxy(
         )
         sock.settimeout(TCP.TIMEOUT)
 
-        with ThreadPoolExecutor() as executor:
-            await self.loop.run_in_executor(executor, sock.connect, destination)
+        await self.loop.run_in_executor(self.crypto_executor, sock.connect, destination)
 
         sock.setblocking(False)
 
@@ -124,27 +134,40 @@ 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)
+            await asyncio.wait_for(self._connect(address), timeout=TCP.TIMEOUT)
         except asyncio.TimeoutError:  # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11
             raise TimeoutError("Connection timed out")
 
     async def close(self) -> None:
-        if self.writer is None:
-            return None
+        async with self.lock:
+            if self.writer is None or self.writer.is_closing():
+                return None
 
-        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.writer = None
-
-    async def send(self, data: bytes) -> None:
+            try:
+                if self.writer.transport is not None:
+                    self.writer.transport.abort()
+
+                self.writer.close()
+
+                await asyncio.wait_for(self.writer.wait_closed(), timeout=TCP.TIMEOUT)
+            except asyncio.TimeoutError:
+                log.warning("Disconnect timed out")
+            except Exception as e:
+                log.info("Close exception: %s %s", type(e).__name__, e)
+            finally:
+                self.writer = None
+
+    async def send(self, data: bytes, wait_for_marker: bool = True) -> None:
         async with self.lock:
             if self.writer is None or self.writer.is_closing():
                 return None
 
+            if wait_for_marker:
+                try:
+                    await asyncio.wait_for(self.marker_event.wait(), timeout=TCP.TIMEOUT)
+                except asyncio.TimeoutError:
+                    raise TimeoutError
+
             try:
                 self.writer.write(data)
                 await self.writer.drain()
@@ -162,7 +185,7 @@ async def recv(self, length: int = 0) -> Optional[bytes]:
             try:
                 chunk = await asyncio.wait_for(
                     self.reader.read(length - len(data)),
-                    TCP.TIMEOUT
+                    timeout=TCP.TIMEOUT
                 )
             except (OSError, asyncio.TimeoutError):
                 return None

pyrogram/connection/transport/tcp/tcp_abridged.py

@@ -26,12 +26,20 @@
 
 
 class TCPAbridged(TCP):
-    def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
-        super().__init__(ipv6, proxy, loop)
+    def __init__(
+        self,
+        ipv6: bool = False,
+        proxy: Proxy = None,
+        crypto_executor_workers: int = 1,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
+        super().__init__(ipv6, proxy, crypto_executor_workers, loop)
 
     async def connect(self, address: Tuple[str, int]) -> None:
+        self.marker_event.clear()
         await super().connect(address)
-        await super().send(b"\xef")
+        await super().send(b"\xef", wait_for_marker=False)
+        self.marker_event.set()
 
     async def send(self, data: bytes, *args) -> None:
         length = len(data) // 4

pyrogram/connection/transport/tcp/tcp_abridged_o.py

@@ -31,13 +31,20 @@
 class TCPAbridgedO(TCP):
     RESERVED = (b"HEAD", b"POST", b"GET ", b"OPTI", b"\xee" * 4)
 
-    def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
-        super().__init__(ipv6, proxy, loop)
+    def __init__(
+        self,
+        ipv6: bool = False,
+        proxy: Proxy = None,
+        crypto_executor_workers: int = 1,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
+        super().__init__(ipv6, proxy, crypto_executor_workers, loop)
 
         self.encrypt = None
         self.decrypt = None
 
     async def connect(self, address: Tuple[str, int]) -> None:
+        self.marker_event.clear()
         await super().connect(address)
 
         while True:
@@ -54,12 +61,13 @@ async def connect(self, address: Tuple[str, int]) -> None:
 
         nonce[56:64] = aes.ctr256_encrypt(nonce, *self.encrypt)[56:64]
 
-        await super().send(nonce)
+        await super().send(nonce, wait_for_marker=False)
+        self.marker_event.set()
 
     async def send(self, data: bytes, *args) -> None:
         length = len(data) // 4
         data = (bytes([length]) if length <= 126 else b"\x7f" + length.to_bytes(3, "little")) + data
-        payload = await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_encrypt, data, *self.encrypt)
+        payload = await self.loop.run_in_executor(self.crypto_executor, aes.ctr256_encrypt, data, *self.encrypt)
 
         await super().send(payload)
 
@@ -84,4 +92,4 @@ async def recv(self, length: int = 0) -> Optional[bytes]:
         if data is None:
             return None
 
-        return await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_decrypt, data, *self.decrypt)
+        return await self.loop.run_in_executor(self.crypto_executor, aes.ctr256_decrypt, data, *self.decrypt)

pyrogram/connection/transport/tcp/tcp_full.py

@@ -28,21 +28,29 @@
 
 
 class TCPFull(TCP):
-    def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
-        super().__init__(ipv6, proxy, loop)
+    def __init__(
+        self,
+        ipv6: bool = False,
+        proxy: Proxy = None,
+        crypto_executor_workers: int = 1,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
+        super().__init__(ipv6, proxy, crypto_executor_workers, loop)
 
-        self.seq_no: Optional[int] = None
+        self.seq_no: int = 0
 
     async def connect(self, address: Tuple[str, int]) -> None:
         await super().connect(address)
         self.seq_no = 0
 
     async def send(self, data: bytes, *args) -> None:
+        self.marker_event.clear()
         data = pack("<II", len(data) + 12, self.seq_no) + data
         data += pack("<I", crc32(data))
         self.seq_no += 1
 
-        await super().send(data)
+        await super().send(data, wait_for_marker=False)
+        self.marker_event.set()
 
     async def recv(self, length: int = 0) -> Optional[bytes]:
         length = await super().recv(4)

pyrogram/connection/transport/tcp/tcp_intermediate.py

@@ -31,8 +31,10 @@ def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEven
         super().__init__(ipv6, proxy, loop)
 
     async def connect(self, address: Tuple[str, int]) -> None:
+        self.marker_event.clear()
         await super().connect(address)
-        await super().send(b"\xee" * 4)
+        await super().send(b"\xee" * 4, wait_for_marker=False)
+        self.marker_event.set()
 
     async def send(self, data: bytes, *args) -> None:
         await super().send(pack("<i", len(data)) + data)

pyrogram/connection/transport/tcp/tcp_intermediate_o.py

@@ -38,6 +38,7 @@ def __init__(self, ipv6: bool, proxy: Proxy, loop: Optional[asyncio.AbstractEven
         self.decrypt = None
 
     async def connect(self, address: Tuple[str, int]) -> None:
+        self.marker_event.clear()
         await super().connect(address)
 
         while True:
@@ -54,7 +55,8 @@ async def connect(self, address: Tuple[str, int]) -> None:
 
         nonce[56:64] = aes.ctr256_encrypt(nonce, *self.encrypt)[56:64]
 
-        await super().send(nonce)
+        await super().send(nonce, wait_for_marker=False)
+        self.marker_event.set()
 
     async def send(self, data: bytes, *args) -> None:
         await super().send(

pyrogram/session/internals/msg_factory.py

@@ -16,23 +16,49 @@
 #  You should have received a copy of the GNU Lesser General Public License
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
+import asyncio
+import time
+
 from pyrogram.raw.core import Message, MsgContainer, TLObject
 from pyrogram.raw.functions import Ping
-from pyrogram.raw.types import MsgsAck, HttpWait
-from .msg_id import MsgId
-from .seq_no import SeqNo
-
-not_content_related = (Ping, HttpWait, MsgsAck, MsgContainer)
+from pyrogram.raw.types import HttpWait, MsgsAck
 
 
 class MsgFactory:
     def __init__(self):
-        self.seq_no = SeqNo()
-
-    def __call__(self, body: TLObject) -> Message:
-        return Message(
-            body,
-            MsgId(),
-            self.seq_no(not isinstance(body, not_content_related)),
-            len(body)
-        )
+        self._last_msg_id = 0
+
+        self._msg_id_lock = asyncio.Lock()
+        self._seq_no_lock = asyncio.Lock()
+
+        self._content_related_messages_sent = 0
+
+    async def allocate_message_identity(self) -> int:
+        async with self._msg_id_lock:
+            now = time.time()
+
+            base_msg_id = int(now * (2**32)) & ~0b11
+
+            if base_msg_id <= self._last_msg_id:
+                base_msg_id = self._last_msg_id + 4
+
+            self._last_msg_id = base_msg_id
+
+            return base_msg_id
+
+    async def allocate_message_sequence(self, is_content_related: bool) -> int:
+        async with self._seq_no_lock:
+            seq_no = (self._content_related_messages_sent * 2) + (1 if is_content_related else 0)
+
+            if is_content_related:
+                self._content_related_messages_sent += 1
+
+            return seq_no
+
+    async def create(self, body: TLObject) -> Message:
+        msg_id = await self.allocate_message_identity()
+
+        is_content_related = not isinstance(body, (Ping, HttpWait, MsgsAck, MsgContainer))
+        seq_no = await self.allocate_message_sequence(is_content_related)
+
+        return Message(body, msg_id, seq_no, len(body))

pyrogram/session/internals/msg_id.py

@@ -23,13 +23,15 @@
 
 
 class MsgId:
-    last_time = 0
-    offset = 0
+    _last_msg_id = 0
 
     def __new__(cls) -> int:
-        now = int(time.time())
-        cls.offset = (cls.offset + 4) if now == cls.last_time else 0
-        msg_id = (now * 2 ** 32) + cls.offset
-        cls.last_time = now
+        now = time.time()
+        base_msg_id = int(now * (2**32)) & ~0b11
 
-        return msg_id
+        if base_msg_id <= cls._last_msg_id:
+           base_msg_id = cls._last_msg_id + 4
+
+        cls._last_msg_id = base_msg_id
+
+        return base_msg_id