Merge pull request #227 from exislow/multithread-track

✨ Track segments are downloaded simultaneously, which results in a hu…

6 files changed, 152 insertions, 63 deletions

diff --git a/tidal_dl_ng/constants.py b/tidal_dl_ng/constants.py
index d3ea2b8..ad8d2ee 100644
--- a/tidal_dl_ng/constants.py
+++ b/tidal_dl_ng/constants.py
@@ -6,6 +6,9 @@ EXTENSION_LYRICS: str = ".lrc"
 UNIQUIFY_THRESHOLD: int = 99
 FILENAME_SANITIZE_PLACEHOLDER: str = "_"
 COVER_NAME: str = "cover.jpg"
+BLOCK_SIZE: int = 4096
+BLOCKS: int = 1024
+CHUNK_SIZE: int = BLOCK_SIZE * BLOCKS
 
 
 class QualityVideo(StrEnum):
diff --git a/tidal_dl_ng/download.py b/tidal_dl_ng/download.py
index b6d5f13..7876368 100644
--- a/tidal_dl_ng/download.py
+++ b/tidal_dl_ng/download.py
@@ -5,12 +5,13 @@ import shutil
 import tempfile
 import time
 from collections.abc import Callable
+from concurrent import futures
 from uuid import uuid4
 
 import ffmpeg
 import m3u8
 import requests
-from constants import COVER_NAME
+from constants import CHUNK_SIZE, COVER_NAME
 from requests.exceptions import HTTPError
 from rich.progress import Progress, TaskID
 from tidalapi import Album, Mix, Playlist, Session, Track, UserPlaylist, Video
@@ -20,7 +21,7 @@ from tidal_dl_ng.config import Settings
 from tidal_dl_ng.constants import EXTENSION_LYRICS, REQUESTS_TIMEOUT_SEC, MediaType, QualityVideo, SkipExisting
 from tidal_dl_ng.helper.decryption import decrypt_file, decrypt_security_token
 from tidal_dl_ng.helper.exceptions import MediaMissing
-from tidal_dl_ng.helper.path import check_file_exists, format_path_media, path_file_sanitize
+from tidal_dl_ng.helper.path import check_file_exists, format_path_media, path_file_sanitize, url_to_filename
 from tidal_dl_ng.helper.tidal import (
     instantiate_media,
     items_results_all,
@@ -30,6 +31,7 @@ from tidal_dl_ng.helper.tidal import (
     name_builder_title,
 )
 from tidal_dl_ng.metadata import Metadata
+from tidal_dl_ng.model.downloader import DownloadSegmentResult
 from tidal_dl_ng.model.gui_data import ProgressBars
 
 
@@ -88,10 +90,13 @@ class Download:
     def _download(
         self,
         media: Track | Video,
-        path_file: str,
-    ) -> str:
+        path_file: pathlib.Path,
+    ) -> pathlib.Path:
         media_name: str = name_builder_item(media)
         urls: [str]
+        path_base: pathlib.Path = path_file.parent
+        result_segments: bool = True
+        dl_segment_results: [DownloadSegmentResult] = []
 
         # Get urls for media.
         if isinstance(media, Track):
@@ -106,9 +111,9 @@ class Download:
 
         # Set the correct progress output channel.
         if self.progress_gui is None:
-            progress_stdout: bool = True
+            progress_to_stdout: bool = True
         else:
-            progress_stdout: bool = False
+            progress_to_stdout: bool = False
             # Send signal to GUI with media name
             self.progress_gui.item_name.emit(media_name[:30])
 
@@ -119,8 +124,11 @@ class Download:
             progress_total: int = urls_count
             block_size: int | None = None
         elif urls_count == 1:
-            # Will be computed later.
-            progress_total: float = None
+            # Get file size and compute progress steps
+            r = requests.head(urls[0], timeout=REQUESTS_TIMEOUT_SEC)
+            total_size_in_bytes: int = int(r.headers.get("content-length", 0))
+            block_size: int | None = 1048576
+            progress_total: float = total_size_in_bytes / block_size
         else:
             raise ValueError
 
@@ -128,58 +136,100 @@ class Download:
         p_task: TaskID = self.progress.add_task(
             f"[blue]Item '{media_name[:30]}'",
             total=progress_total,
-            visible=progress_stdout,
+            visible=progress_to_stdout,
         )
 
-        # Write content to file until progress is finished.
+        # Download segments until progress is finished.
         while not self.progress.tasks[p_task].finished:
-            with open(path_file, "wb") as f:
-                for url in urls:
-                    try:
-                        # Create the request object with stream=True, so the content won't be loaded into memory at once.
-                        r = requests.get(url, stream=True, timeout=REQUESTS_TIMEOUT_SEC)
+            with futures.ThreadPoolExecutor(
+                max_workers=self.settings.data.downloads_simultaneous_per_track_max
+            ) as executor:
+                # Dispatch all download tasks to worker threads
+                l_futures: [any] = [
+                    executor.submit(self._download_segment, url, path_base, block_size, p_task, progress_to_stdout)
+                    for url in urls
+                ]
+                # Report results as they become available
+                for future in futures.as_completed(l_futures):
+                    # Retrieve result
+                    result_dl_segment: DownloadSegmentResult = future.result()
 
-                        r.raise_for_status()
+                    dl_segment_results.append(result_dl_segment)
 
-                        # Compute progress iterations based on the file size and update task details.
-                        if not progress_total:
-                            # Get file size and compute progress steps
-                            total_size_in_bytes: int = int(r.headers.get("content-length", 0))
-                            block_size: int | None = 1048576
-                            progress_total: float = total_size_in_bytes / block_size
+                    # check for a link that was skipped
+                    if not result_dl_segment.result and (result_dl_segment.url is not urls[-1]):
+                        # Sometimes it happens, if a track is very short (< 8 seconds or so), that the last URL in `urls` is
+                        # invalid (HTTP Error 500) and not necessary. File won't be corrupt.
+                        # If this is NOT the case, but any other URL has resulted in an error,
+                        # mark the whole thing as corrupt.
+                        result_segments = False
 
-                            self.progress.update(p_task, total=progress_total)
+        # TODO: Implement error handling on corrupt segments (also in the following method.
+        reult_merge: bool = self._segments_merge(path_file, dl_segment_results)
+        tmp_path_file_decrypted: pathlib.Path = pathlib.Path("")
 
-                        # Write the content to disk. If `chunk_size` is set to `None` the whole file will be written at once.
-                        for data in r.iter_content(chunk_size=block_size):
-                            f.write(data)
-                            # Advance progress bar.
-                            self.progress.advance(p_task)
-                    except HTTPError as e:
-                        if url is urls[-1]:
-                            # It happens, if a track is very short (< 8 seconds or so), that the last URL in `urls` is
-                            # invalid (HTTP Error 500) and not necessary. File won't be corrupt.
-                            # Thus, advance progress bar to avoid infinity loops.
-                            self.progress.advance(p_task)
-                        else:
-                            # Finish downloading early and report error.
-                            # TODO: The track should somehow be marked as corrupt.
-                            self.progress.update(p_task, completed=progress_total)
-                            self.fn_logger.error(e)
-                    finally:
-                        # To send the progress to the GUI, we need to emit the percentage.
-                        if not progress_stdout:
-                            self.progress_gui.item.emit(self.progress.tasks[p_task].percentage)
-
-        if isinstance(media, Track) and stream_manifest.is_encrypted:
-            key, nonce = decrypt_security_token(stream_manifest.encryption_key)
-            tmp_path_file_decrypted = path_file + "_decrypted"
-            decrypt_file(path_file, tmp_path_file_decrypted, key, nonce)
-        else:
-            tmp_path_file_decrypted = path_file
+        if reult_merge:
+            if isinstance(media, Track) and stream_manifest.is_encrypted:
+                key, nonce = decrypt_security_token(stream_manifest.encryption_key)
+                tmp_path_file_decrypted = path_file + "_decrypted"
+                decrypt_file(path_file, tmp_path_file_decrypted, key, nonce)
+            else:
+                tmp_path_file_decrypted = path_file
 
         return tmp_path_file_decrypted
 
+    def _segments_merge(self, path_file, dl_segment_results) -> bool:
+        result: bool
+
+        # Copy the content of all segments into one file.
+        try:
+            with path_file.open("wb") as f_target:
+                for dl_segment_result in dl_segment_results:
+                    if dl_segment_result.result:
+                        with dl_segment_result.path_segment.open("rb") as f_segment:
+                            # Read and write junks, which gives better HDD write performance
+                            while segment := f_segment.read(CHUNK_SIZE):
+                                f_target.write(segment)
+
+                        # Delete segment from HDD
+                        dl_segment_result.path_segment.unlink()
+
+            result = True
+        except:
+            result = False
+
+        return result
+
+    def _download_segment(
+        self, url: str, path_base: pathlib.Path, block_size: int | None, p_task: TaskID, progress_to_stdout: bool
+    ) -> DownloadSegmentResult:
+        result: bool = False
+        path_segment: pathlib.Path = path_base / url_to_filename(url)
+
+        try:
+            # Create the request object with stream=True, so the content won't be loaded into memory at once.
+            r = requests.get(url, stream=True, timeout=REQUESTS_TIMEOUT_SEC)
+
+            r.raise_for_status()
+
+            # Write the content to disk. If `chunk_size` is set to `None` the whole file will be written at once.
+            with path_segment.open("wb") as f:
+                for data in r.iter_content(chunk_size=block_size):
+                    f.write(data)
+                    # Advance progress bar.
+                    self.progress.advance(p_task)
+
+            result = True
+        except HTTPError:
+            # TODO: Maybe return e as well?
+            self.progress.advance(p_task)
+        finally:
+            # To send the progress to the GUI, we need to emit the percentage.
+            if not progress_to_stdout:
+                self.progress_gui.item.emit(self.progress.tasks[p_task].percentage)
+
+            return DownloadSegmentResult(result, url, path_segment)
+
     def item(
         self,
         file_template: str,
@@ -221,6 +271,7 @@ class Download:
 
         # Get extension.
         file_extension: str
+        do_flac_extract = False
 
         if isinstance(media, Track):
             # If a quality is explicitly set, change it.
@@ -232,8 +283,6 @@ class Download:
             file_extension = AudioExtensions.M4A if file_extension == AudioExtensions.MP4 else file_extension
 
             if self.settings.data.extract_flac:
-                do_flac_extract = False
-
                 if (
                     media.get_stream().get_stream_manifest().codecs.upper() == Codec.FLAC
                     and file_extension != AudioExtensions.FLAC
@@ -266,7 +315,11 @@ class Download:
         if not file_exists:
             # Create a temp directory and file.
             with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp_path_dir:
-                tmp_path_file = os.path.join(tmp_path_dir, str(uuid4()))
+                tmp_path_file: pathlib.Path = pathlib.Path(tmp_path_dir) / str(uuid4())
+
+                # Create empty file
+                tmp_path_file.touch()
+
                 # Download media.
                 tmp_path_file = self._download(media=media, path_file=tmp_path_file)
 
@@ -403,7 +456,7 @@ class Download:
         return result
 
     def metadata_write(
-        self, track: Track, path_media: str, is_parent_album: str
+        self, track: Track, path_media: pathlib.Path, is_parent_album: bool
     ) -> (bool, pathlib.Path | None, pathlib.Path | None):
         result: bool = False
         path_lyrics: pathlib.Path | None = None
@@ -433,14 +486,14 @@ class Download:
                 print(f"Could not retrieve lyrics for `{name_builder_item(track)}`.")
 
         if lyrics and self.settings.data.lyrics_file:
-            path_lyrics = self.lyrics_to_file(pathlib.Path(path_media).parent, lyrics)
+            path_lyrics = self.lyrics_to_file(path_media.parent, lyrics)
 
         if self.settings.data.metadata_cover_embed or (self.settings.data.cover_album_file and is_parent_album):
             url_cover = track.album.image(int(self.settings.data.metadata_cover_dimension))
             cover_data = self.cover_data(url=url_cover)
 
         if cover_data and self.settings.data.cover_album_file and is_parent_album:
-            path_cover = self.cover_to_file(pathlib.Path(path_media).parent, cover_data)
+            path_cover = self.cover_to_file(path_media.parent, cover_data)
 
         # `None` values are not allowed.
         m: Metadata = Metadata(
@@ -533,8 +586,8 @@ class Download:
                 if not progress_stdout:
                     self.progress_gui.list_item.emit(self.progress.tasks[p_task1].percentage)
 
-    def _video_convert(self, path_file: str) -> str:
-        path_file_out = path_file + AudioExtensions.MP4
+    def _video_convert(self, path_file: pathlib.Path) -> pathlib.Path:
+        path_file_out: pathlib.Path = path_file.with_suffix(AudioExtensions.MP4)
         result, _ = (
             ffmpeg.input(path_file)
             .output(path_file_out, map=0, c="copy", loglevel="quiet")
@@ -543,8 +596,8 @@ class Download:
 
         return path_file_out
 
-    def _extract_flac(self, path_media_src: str) -> str:
-        path_media_out = path_media_src + AudioExtensions.FLAC
+    def _extract_flac(self, path_media_src: pathlib.Path) -> pathlib.Path:
+        path_media_out = path_media_src.with_suffix(AudioExtensions.FLAC)
         result, _ = (
             ffmpeg.input(path_media_src)
             .output(
diff --git a/tidal_dl_ng/helper/decryption.py b/tidal_dl_ng/helper/decryption.py
index ebb946a..17ea19b 100644
--- a/tidal_dl_ng/helper/decryption.py
+++ b/tidal_dl_ng/helper/decryption.py
@@ -1,4 +1,5 @@
 import base64
+import pathlib
 
 from Crypto.Cipher import AES
 from Crypto.Util import Counter
@@ -35,7 +36,7 @@ def decrypt_security_token(security_token: str) -> (str, str):
     return key, nonce
 
 
-def decrypt_file(path_file_encrypted: str, path_file_destination: str, key: str, nonce: str) -> None:
+def decrypt_file(path_file_encrypted: pathlib.Path, path_file_destination: pathlib.Path, key: str, nonce: str) -> None:
     """
     Decrypts an encrypted MQA file given the file, key and nonce.
     TODO: Is it really only necessary for MQA of for all other formats, too?
@@ -46,9 +47,9 @@ def decrypt_file(path_file_encrypted: str, path_file_destination: str, key: str,
     decryptor = AES.new(key, AES.MODE_CTR, counter=counter)
 
     # Open and decrypt
-    with open(path_file_encrypted, "rb") as f_src:
+    with path_file_encrypted.open("rb") as f_src:
         audio_decrypted = decryptor.decrypt(f_src.read())
 
         # Replace with decrypted file
-        with open(path_file_destination, "wb") as f_dst:
+        with path_file_destination.open("wb") as f_dst:
             f_dst.write(audio_decrypted)
diff --git a/tidal_dl_ng/helper/path.py b/tidal_dl_ng/helper/path.py
index cc8d5b6..e6f5b14 100644
--- a/tidal_dl_ng/helper/path.py
+++ b/tidal_dl_ng/helper/path.py
@@ -1,8 +1,10 @@
 import math
 import os
+import posixpath
 import re
 import sys
 from pathlib import Path, PosixPath
+from urllib.parse import unquote, urlsplit
 
 from pathvalidate import sanitize_filename, sanitize_filepath
 from pathvalidate.error import ValidationError
@@ -295,3 +297,22 @@ def resource_path(relative_path):
         base_path = os.path.abspath(".")
 
     return os.path.join(base_path, relative_path)
+
+
+def url_to_filename(url: str) -> str:
+    """Return basename corresponding to url.
+    >>> print(url_to_filename('http://example.com/path/to/file%C3%80?opt=1'))
+    fileÀ
+    >>> print(url_to_filename('http://example.com/slash%2fname')) # '/' in name
+    Taken from https://gist.github.com/zed/c2168b9c52b032b5fb7d
+    Traceback (most recent call last):
+    ...
+    ValueError
+    """
+    urlpath: str = urlsplit(url).path
+    basename: str = posixpath.basename(unquote(urlpath))
+
+    if os.path.basename(basename) != basename or unquote(posixpath.basename(urlpath)) != basename:
+        raise ValueError  # reject '%2f' or 'dir%5Cbasename.ext' on Windows
+
+    return basename
diff --git a/tidal_dl_ng/model/cfg.py b/tidal_dl_ng/model/cfg.py
index fae3590..6e9b991 100644
--- a/tidal_dl_ng/model/cfg.py
+++ b/tidal_dl_ng/model/cfg.py
@@ -37,6 +37,7 @@ class Settings:
     metadata_cover_embed: bool = True
     cover_album_file: bool = True
     extract_flac: bool = True
+    downloads_simultaneous_per_track_max: int = 20
 
 
 @dataclass_json
@@ -82,6 +83,7 @@ class HelpSettings:
     metadata_cover_embed: str = "Embed album cover into file."
     cover_album_file: str = "Save cover to 'cover.jpg', if an album is downloaded."
     extract_flac: str = "Extract FLAC audio tracks from MP4 containers and save them as `*.flac` (uses FFmpeg)."
+    downloads_simultaneous_per_track_max: str = "Maximum number of simultaneous chunk downloads per track."
 
 
 @dataclass_json
diff --git a/tidal_dl_ng/model/downloader.py b/tidal_dl_ng/model/downloader.py
new file mode 100644
index 0000000..b8c5348
--- /dev/null
+++ b/tidal_dl_ng/model/downloader.py
@@ -0,0 +1,9 @@
+import pathlib
+from dataclasses import dataclass
+
+
+@dataclass
+class DownloadSegmentResult:
+    result: bool
+    url: str
+    path_segment: pathlib.Path