Download/Upload speed improvement

Author: Surendra9123

pyrogram/client.py

@@ -26,6 +26,8 @@
 import shutil
 import sys
 import time
+import math
+import tempfile
 from concurrent.futures.thread import ThreadPoolExecutor
 from datetime import datetime, timedelta
 from hashlib import sha256
@@ -198,6 +200,10 @@ class Client(Methods):
             Defaults to False, because ``getpass`` (the library used) is known to be problematic in some
             terminal environments.
 
+        crypto_executor_workers (``int``, *optional*):
+            Set the maximum crypto executor workers value for a session.
+            Defaults to ``workers if workers < 5 else 5``.
+
         max_concurrent_transmissions (``int``, *optional*):
             Set the maximum amount of concurrent transmissions (uploads & downloads).
             A value that is too high may result in network related issues.
@@ -260,6 +266,7 @@ class Client(Methods):
     UPGRADED_GIFT_RE = re.compile(r"^(?:https?://)?(?:www\.)?(?:t(?:elegram)?\.(?:org|me|dog)/(?:nft/|\+))([\w-]+)$")
     CHANNEL_MESSAGE_LINK_RE = re.compile(r"^(?:https?://)?(?:www\.)?(?:t(?:elegram)?\.(?:org|me|dog)/(?:c/)?)([\w]+)(?:.+)?$")
     WORKERS = min(32, (os.cpu_count() or 0) + 4)  # os.cpu_count() can be None
+    CRYPTO_EXECUTOR_WORKERS = WORKERS if WORKERS < 5 else 5
     WORKDIR = PARENT_DIR
 
     # Interval of seconds in which the updates watchdog will kick in
@@ -301,6 +308,7 @@ def __init__(
         takeout: Optional[bool] = None,
         sleep_threshold: int = Session.SLEEP_THRESHOLD,
         hide_password: Optional[bool] = False,
+        crypto_executor_workers: int = CRYPTO_EXECUTOR_WORKERS,
         max_concurrent_transmissions: int = MAX_CONCURRENT_TRANSMISSIONS,
         max_message_cache_size: int = MAX_MESSAGE_CACHE_SIZE,
         max_topic_cache_size: int = MAX_TOPIC_CACHE_SIZE,
@@ -345,6 +353,7 @@ def __init__(
         self.takeout = takeout
         self.sleep_threshold = sleep_threshold
         self.hide_password = hide_password
+        self.crypto_executor_workers = crypto_executor_workers
         self.max_concurrent_transmissions = max_concurrent_transmissions
         self.max_message_cache_size = max_message_cache_size
         self.max_topic_cache_size = max_topic_cache_size
@@ -1072,12 +1081,13 @@ def load_plugins(self):
     async def handle_download(self, packet):
         file_id, directory, file_name, in_memory, file_size, progress, progress_args = packet
 
+        max_workers = self.crypto_executor_workers
         os.makedirs(directory, exist_ok=True) if not in_memory else None
         temp_file_path = os.path.abspath(re.sub("\\\\", "/", os.path.join(directory, file_name))) + ".temp"
         file = BytesIO() if in_memory else open(temp_file_path, "wb")
 
         try:
-            async for chunk in self.get_file(file_id, file_size, 0, 0, progress, progress_args):
+            async for chunk in self.get_file(file_id, file_size, 0, 0, progress, progress_args, max_workers, in_memory):
                 file.write(chunk)
         except BaseException as e:
             if not in_memory:
@@ -1108,7 +1118,9 @@ async def get_file(
         limit: int = 0,
         offset: int = 0,
         progress: Callable = None,
-        progress_args: tuple = ()
+        progress_args: tuple = (),
+        max_workers: int = 1,
+        in_memory: bool = False
     ) -> AsyncGenerator[bytes, None]:
         async with self.get_file_semaphore:
             file_type = file_id.file_type
@@ -1151,15 +1163,39 @@ async def get_file(
                 )
 
             current = 0
+            temporary = True
             total = abs(limit) or (1 << 31) - 1
             chunk_size = 1024 * 1024
             offset_bytes = abs(offset) * chunk_size
 
             dc_id = file_id.dc_id
 
+            async def handle_progress(downloaded_bytes):
+                if not progress:
+                    return
+                func = functools.partial(
+                    progress,
+                    min(downloaded_bytes, file_size) if file_size != 0 else downloaded_bytes,
+                    file_size,
+                    *progress_args
+                )
+                if inspect.iscoroutinefunction(progress):
+                    asyncio.create_task(func())
+                else:
+                    self.loop.run_in_executor(self.executor, func)
+          
             try:
-                session = await self.get_session(dc_id, is_media=True)
-
+                try:
+                   session = await self.get_session(dc_id, is_media=True, temporary=True)
+                except Exception as e:
+                   temporary = False
+                   try:
+                       await session.stop()
+                   except:
+                       pass
+                   log.info(f"Reusing cached media session due to {e}")
+                   session = await self.get_session(dc_id, is_media=True)
+                   
                 r = await session.invoke(
                     raw.functions.upload.GetFile(
                         location=location,
@@ -1170,40 +1206,119 @@ async def get_file(
                 )
 
                 if isinstance(r, raw.types.upload.File):
-                    while True:
+                    if file_size and file_size > 0:
+                        total_parts = math.ceil(file_size / chunk_size)
+                    else:
+                        if limit:
+                            total_parts = abs(limit)
+                        else:
+                            total_parts = None
+                    if total_parts is None:
                         chunk = r.bytes
-
                         yield chunk
-
                         current += 1
                         offset_bytes += chunk_size
+                        await handle_progress(offset_bytes)
 
-                        if progress:
-                            func = functools.partial(
-                                progress,
-                                min(offset_bytes, file_size)
-                                if file_size != 0
-                                else offset_bytes,
-                                file_size,
-                                *progress_args
+                        while True:
+                            if len(chunk) < chunk_size or current >= total:
+                                break
+    
+                            r = await session.invoke(
+                                raw.functions.upload.GetFile(
+                                    location=location,
+                                    offset=offset_bytes,
+                                    limit=chunk_size
+                                ),
+                                sleep_threshold=30
                             )
-
-                            if inspect.iscoroutinefunction(progress):
-                                await func()
+                            if not isinstance(r, raw.types.upload.File):
+                                break
+                            chunk = r.bytes
+                            yield chunk
+                            current += 1
+                            offset_bytes += chunk_size
+                            await handle_progress(offset_bytes)
+                    else:
+                        concurrency = min(max_workers, total_parts)
+                        sem = asyncio.Semaphore(concurrency)
+                        downloaded_bytes = 0
+                        exceptions: list[Exception] = []
+                        parts_lock = asyncio.Lock()
+
+                        parts: list[dict] = []
+                        for i in range(total_parts):
+                            part = {
+                                "part_no": i,
+                                "is_completed": asyncio.Future(),
+                                "bytes": None
+                            }
+                            parts.append(part)
+
+                        if not in_memory:
+                            tempdir = tempfile.mkdtemp(prefix=f"pyro_getfile_{self.rnd_id()}")
+                            for part in parts:
+                                part["bytes"] = os.path.join(tempdir, f"part_{part['part_no']}")
+
+                            def write_file(data: bytes, path: str):
+                                with open(path, "wb") as f:
+                                    f.write(data)
+
+                            def cleanup_tempdir():
+                                for part in parts:
+                                    try:
+                                        os.remove(part["bytes"])
+                                    except:
+                                        pass
+                                try:
+                                    os.rmdir(tempdir)
+                                except:
+                                    pass
+
+                        async def download_part(part: dict):
+                            nonlocal downloaded_bytes
+                            part_no = part["part_no"]
+                            local_offset = offset_bytes + part_no * chunk_size
+                            async with sem:
+                                try:
+                                    r = await session.invoke(
+                                        raw.functions.upload.GetFile(location=location, offset=local_offset, limit=chunk_size),
+                                        sleep_threshold=30
+                                    )
+                                    if not isinstance(r, raw.types.upload.File):
+                                        raise RuntimeError(f"Unexpected response for part {part_no}: {type(r)}")
+
+                                    if in_memory:
+                                        part["bytes"] = r.bytes
+                                    else:
+                                        await self.loop.run_in_executor(self.executor, write_file, r.bytes, part["bytes"])
+
+                                    async with parts_lock:
+                                        downloaded_bytes += len(r.bytes)
+                                        await handle_progress(min(downloaded_bytes, file_size))
+
+                                    part["is_completed"].set_result(True)
+                                except Exception as e:
+                                    if not part["is_completed"].done():
+                                        part["is_completed"].set_exception(e)
+                                    exceptions.append(e)
+
+                        tasks = [asyncio.create_task(download_part(part)) for part in parts]
+
+                        for part in parts:
+                            await part["is_completed"]
+                            if in_memory:
+                                yield part["bytes"]
                             else:
-                                await self.loop.run_in_executor(self.executor, func)
+                                data = await self.loop.run_in_executor(self.executor, lambda p=part["bytes"]: open(p, "rb").read())
+                                yield data
 
-                        if len(chunk) < chunk_size or current >= total:
-                            break
+                        await asyncio.gather(*tasks, return_exceptions=True)
+                        if not in_memory:
+                            await self.loop.run_in_executor(self.executor, cleanup_tempdir)
 
-                        r = await session.invoke(
-                            raw.functions.upload.GetFile(
-                                location=location,
-                                offset=offset_bytes,
-                                limit=chunk_size
-                            ),
-                            sleep_threshold=30
-                        )
+                        if exceptions:
+                            raise exceptions[0]
 
                 elif isinstance(r, raw.types.upload.FileCdnRedirect):
 
@@ -1291,6 +1406,9 @@ def _check_all_hashes():
                 raise
             except Exception as e:
                 log.exception(e)
+            finally:
+                if temporary:
+                    await session.stop()
 
     async def get_session(
         self,
@@ -1367,6 +1485,10 @@ async def get_session(
 
         if is_media:
             auth_key = (await self.get_session(dc_id)).auth_key
+
+            if temporary:
+                await self.get_session(dc_id, is_media=True)
+
         else:
             if not is_current_dc:
                 auth_key = await Auth(
@@ -1386,7 +1508,8 @@ async def get_session(
             port,
             auth_key,
             await self.storage.test_mode(),
-            is_media=is_media
+            is_media=is_media,
+            crypto_executor_workers=self.crypto_executor_workers
         )
 
         if not temporary: