From a48d27f501b73eee842384723c33db298ce1e6c3 Mon Sep 17 00:00:00 2001
From: Dan <14043624+delivrance@users.noreply.github.com>
Date: Sat, 2 Jan 2021 18:45:43 +0100
Subject: [PATCH] Always run crypto-related functions in the dedicated thread

---
 pyrogram/__init__.py                               |  2 +-
 .../connection/transport/tcp/tcp_abridged_o.py     |  8 +++-----
 pyrogram/session/session.py                        | 14 +++++++++-----
 pyrogram/utils.py                                  |  8 --------
 4 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/pyrogram/__init__.py b/pyrogram/__init__.py
index 987b7848..7733e0f7 100644
--- a/pyrogram/__init__.py
+++ b/pyrogram/__init__.py
@@ -46,4 +46,4 @@ main_event_loop = asyncio.get_event_loop()
 
 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
index b05cad16..9db148bf 100644
--- 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 @@ class TCPAbridgedO(TCP):
     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 @@ class TCPAbridgedO(TCP):
         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
index 6ce342ea..686531e3 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -25,7 +25,7 @@ from hashlib import sha1
 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 @@ class Session:
         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 @@ class Session:
         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
index bc169ded..1082dade 100644
--- 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)
-    )