Always run crypto-related functions in the dedicated thread

delivrance · delivrance · commit a48d27f501b7 · 2021-01-02T18:45:43.000+01:00
diff --git a/pyrogram/__init__.py b/pyrogram/__init__.py
@@ -46,4 +46,4 @@ class ContinuePropagation(StopAsyncIteration):
 
 CRYPTO_EXECUTOR_SIZE_THRESHOLD = 512
 
-crypto_executor = ThreadPoolExecutor(2, thread_name_prefix="CryptoWorker")
+crypto_executor = ThreadPoolExecutor(1, thread_name_prefix="CryptoWorker")
diff --git a/pyrogram/connection/transport/tcp/tcp_abridged_o.py b/pyrogram/connection/transport/tcp/tcp_abridged_o.py
@@ -18,11 +18,9 @@
 
 import logging
 import os
-
 from typing import Optional
 
-from pyrogram import utils
-
+import pyrogram
 from pyrogram.crypto import aes
 from .tcp import TCP
 
@@ -60,7 +58,7 @@ async def connect(self, address: tuple):
     async def send(self, data: bytes, *args):
         length = len(data) // 4
         data = (bytes([length]) if length <= 126 else b"\x7f" + length.to_bytes(3, "little")) + data
-        payload = await utils.maybe_run_in_executor(aes.ctr256_encrypt, data, len(data), self.loop, *self.encrypt)
+        payload = await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_encrypt, data, *self.encrypt)
 
         await super().send(payload)
 
@@ -85,4 +83,4 @@ async def recv(self, length: int = 0) -> Optional[bytes]:
         if data is None:
             return None
 
-        return await utils.maybe_run_in_executor(aes.ctr256_decrypt, data, len(data), self.loop, *self.decrypt)
+        return await self.loop.run_in_executor(pyrogram.crypto_executor, aes.ctr256_decrypt, data, *self.decrypt)
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
@@ -25,7 +25,7 @@
 from io import BytesIO
 
 import pyrogram
-from pyrogram import __copyright__, __license__, __version__, utils
+from pyrogram import __copyright__, __license__, __version__
 from pyrogram import raw
 from pyrogram.connection import Connection
 from pyrogram.crypto import mtproto
@@ -217,8 +217,10 @@ async def restart(self):
         await self.start()
 
     async def handle_packet(self, packet):
-        data = await utils.maybe_run_in_executor(
-            mtproto.unpack, BytesIO(packet), len(packet), self.loop,
+        data = await self.loop.run_in_executor(
+            pyrogram.crypto_executor,
+            mtproto.unpack,
+            BytesIO(packet),
             self.session_id,
             self.auth_key,
             self.auth_key_id
@@ -360,8 +362,10 @@ async def _send(self, data: TLObject, wait_response: bool = True, timeout: float
         log.debug(f"Sent:")
         log.debug(message)
 
-        payload = await utils.maybe_run_in_executor(
-            mtproto.pack, message, len(message), self.loop,
+        payload = await self.loop.run_in_executor(
+            pyrogram.crypto_executor,
+            mtproto.pack,
+            message,
             self.current_salt.salt,
             self.session_id,
             self.auth_key,
diff --git a/pyrogram/utils.py b/pyrogram/utils.py
@@ -309,11 +309,3 @@ async def parse_text_entities(
         "message": text,
         "entities": entities
     }
-
-
-async def maybe_run_in_executor(func, data, length, loop, *args):
-    return (
-        func(data, *args)
-        if length <= pyrogram.CRYPTO_EXECUTOR_SIZE_THRESHOLD
-        else await loop.run_in_executor(pyrogram.crypto_executor, func, data, *args)
-    )

Original file line number	Diff line number	Diff line change
`@@ -46,4 +46,4 @@ class ContinuePropagation(StopAsyncIteration):`
`46`	`46`
`47`	`47`	`CRYPTO_EXECUTOR_SIZE_THRESHOLD = 512`
`48`	`48`
`49`		`-crypto_executor = ThreadPoolExecutor(2, thread_name_prefix="CryptoWorker")`
	`49`	`+crypto_executor = ThreadPoolExecutor(1, thread_name_prefix="CryptoWorker")`