feat(core): use thread pool to limit number of active downloads

Signed-off-by: Jordan Shatford <[email protected]>

core/ydcore/__init__.py

@@ -11,7 +11,6 @@
 from .models import Video
 from .models import VideoFormat
 from .search import YouTubeSearch
-from .thread import YoutubeDownloadThread
 
 __all__ = [
     'AudioFormat',
@@ -27,5 +26,4 @@
     'Video',
     'VideoFormat',
     'YouTubeSearch',
-    'YoutubeDownloadThread',
 ]

core/ydcore/downloader.py

@@ -0,0 +1,20 @@
+from yt_dlp import YoutubeDL  # type: ignore
+
+from .config import DownloadConfig
+from .models import DownloadState
+from .models import DownloadStatus
+
+
+def run_downloader(config: DownloadConfig) -> None:
+    downloader = YoutubeDL(config.as_ytdlp_params)
+    config.on_status_update(
+        DownloadStatus(state=DownloadState.DOWNLOADING),
+    )
+    try:
+        downloader.download(  # type: ignore
+            [str(config.download.video.url)],
+        )
+    except Exception:
+        config.on_status_update(
+            DownloadStatus(state=DownloadState.ERROR),
+        )

core/ydcore/manager.py

@@ -1,23 +1,26 @@
 import os
 import time
+from multiprocessing.pool import ThreadPool
 
 from .config import DownloadConfig
 from .config import StatusHook
+from .downloader import run_downloader
 from .models import Download
 from .models import DownloadFile
 from .models import DownloadInput
-from .thread import YoutubeDownloadThread
 
 
 class DownloadManager:
     def __init__(
         self, output_dir: str,
         status_hook: StatusHook | None = None,
+        *,
+        num_threads: int | None = 2,
     ):
         self._status_hook = status_hook
-        self._configs: dict[str, DownloadConfig] = {}
-        self._downloads: dict[str, YoutubeDownloadThread] = {}
+        self._downloads: dict[str, DownloadConfig] = {}
         self._output_dir = output_dir
+        self._pool = ThreadPool(num_threads)
 
     def __contains__(self, download_id: str) -> bool:
         return download_id in self._downloads
@@ -27,29 +30,26 @@ def add(self, download: DownloadInput) -> Download:
             config = DownloadConfig(download, self._output_dir)
             if self._status_hook:
                 config.add_status_hook(self._status_hook)
-            self._configs[download.video.id] = config
-            thread = YoutubeDownloadThread(self._configs[download.video.id])
-            self._downloads[download.video.id] = thread
-            thread.start()
-        return self._configs[download.video.id].download
+            self._downloads[download.video.id] = config
+            self._pool.apply_async(run_downloader, (config,))
+        return self._downloads[download.video.id].download
 
     def remove(self, download_id: str) -> None:
         if download_id in self._downloads:
-            path = self._configs[download_id].path
+            path = self._downloads[download_id].path
             if os.path.exists(path):
                 os.remove(path)
             self._downloads.pop(download_id, None)
-            self._configs.pop(download_id, None)
 
     def get(self, download_id: str) -> Download | None:
-        config = self._configs.get(download_id, None)
+        config = self._downloads.get(download_id, None)
         return None if config is None else config.download
 
     def get_all(self) -> list[Download]:
-        return [config.download for config in self._configs.values()]
+        return [config.download for config in self._downloads.values()]
 
     def get_file(self, download_id: str) -> DownloadFile | None:
-        config = self._configs.get(download_id, None)
+        config = self._downloads.get(download_id, None)
         if config is None or not os.path.exists(config.path):
             return None
         return DownloadFile(name=config.filename, path=config.path)
@@ -59,8 +59,7 @@ def send_status_update(self, update: Download) -> None:
             self._status_hook(update)
 
     def wait(self) -> None:
-        for download in self._downloads.values():
-            download.join()
-        self._configs = {}
+        self._pool.close()
+        self._pool.join()
         self._downloads = {}
         time.sleep(1)