From af58b0a83db5f1ae6d5ca4a74e280f97b8ec6b13 Mon Sep 17 00:00:00 2001 From: Rikrdoga Date: Tue, 4 Nov 2025 12:36:27 -0600 Subject: [PATCH] Resolve Race Condition and Optimize Session Switching for Atmos Downloads --- tidal_dl_ng/config.py | 81 ++++++++++++++++++++++++---------- tidal_dl_ng/download.py | 98 ++++++++++++++++++++++++++++++++--------- 2 files changed, 134 insertions(+), 45 deletions(-) diff --git a/tidal_dl_ng/config.py b/tidal_dl_ng/config.py index 1ba8c8c..134f3ff 100644 --- a/tidal_dl_ng/config.py +++ b/tidal_dl_ng/config.py @@ -2,10 +2,9 @@ import json import os import shutil from collections.abc import Callable -from contextlib import contextmanager from json import JSONDecodeError from pathlib import Path -from threading import Event +from threading import Event, Lock from typing import Any import tidalapi @@ -103,6 +102,19 @@ class Tidal(BaseConfig, metaclass=SingletonMeta): self.cls_model = ModelToken tidal_config: tidalapi.Config = tidalapi.Config(item_limit=10000) self.session = tidalapi.Session(tidal_config) + self.original_client_id = self.session.config.client_id + self.original_client_secret = self.session.config.client_secret + # Lock to ensure session-switching is thread-safe. + # This lock protects against a race condition where one thread + # changes the session credentials while another is using them. + # It is intentionally held by Download._get_stream_info + # for the *entire* duration of the credential switch AND + # the get_stream() call. + self.stream_lock = Lock() + # State-tracking flag to prevent redundant, expensive + # session re-authentication when the session is already in the + # correct mode (Atmos or Normal). + self.is_atmos_session = False # self.session.config.client_id = "km8T1xS355y7dd3H" # self.session.config.client_secret = "vcmeGW1OuZ0fWYMCSZ6vNvSLJlT3XEpW0ambgYt5ZuI=" self.file_path = path_file_token() @@ -116,7 +128,8 @@ class Tidal(BaseConfig, metaclass=SingletonMeta): if settings: self.settings = settings - self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio) + if not self.is_atmos_session: + self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio) self.session.video_quality = tidalapi.VideoQuality.high return True @@ -162,33 +175,55 @@ class Tidal(BaseConfig, metaclass=SingletonMeta): self.set_option("expiry_time", self.session.expiry_time) self.save() - @contextmanager - def atmos_session_context(self): + def switch_to_atmos_session(self) -> bool: + """ + Switches the shared session to Dolby Atmos credentials. + Only re-authenticates if not already in Atmos mode. + """ + # If we are already in Atmos mode, do nothing. + if self.is_atmos_session: + return True - if not self.session.check_login(): - print("Not logged in.") + print("Switching session context to Dolby Atmos...") + self.session.config.client_id = ATMOS_CLIENT_ID + self.session.config.client_secret = ATMOS_CLIENT_SECRET + self.session.audio_quality = ATMOS_REQUEST_QUALITY - original_client_id = self.session.config.client_id - original_client_secret = self.session.config.client_secret - original_audio_quality = self.session.audio_quality + # Re-login with new credentials + if not self.login_token(do_pkce=self.is_pkce): + print("Warning: Atmos session authentication failed.") + # Try to switch back to normal to be safe + self.restore_normal_session(force=True) + return False - try: - self.session.config.client_id = ATMOS_CLIENT_ID - self.session.config.client_secret = ATMOS_CLIENT_SECRET - self.session.audio_quality = ATMOS_REQUEST_QUALITY + self.is_atmos_session = True # Set the flag + print("Session is now in Atmos mode.") + return True - if not self.login_token(do_pkce=self.is_pkce): - print("Warning: Session restore failed.") + def restore_normal_session(self, force: bool = False) -> bool: + """ + Restores the shared session to the original user credentials. + Only re-authenticates if not already in Normal mode. + """ + # If we are already in Normal mode (and not forced), do nothing. + if not self.is_atmos_session and not force: + return True - yield + print("Restoring session context to Normal...") + self.session.config.client_id = self.original_client_id + self.session.config.client_secret = self.original_client_secret - finally: - self.session.config.client_id = original_client_id - self.session.config.client_secret = original_client_secret - self.session.audio_quality = original_audio_quality + # Re-apply user's quality setting + self.settings_apply() - if not self.login_token(do_pkce=self.is_pkce): - print("Warning: Restoring the original session context failed. Please restart the application.") + # Re-login with original credentials + if not self.login_token(do_pkce=self.is_pkce): + print("Warning: Restoring the original session context failed. Please restart the application.") + return False + + self.is_atmos_session = False # Set the flag + print("Session is now in Normal mode.") + return True def login(self, fn_print: Callable) -> bool: is_token = self.login_token() diff --git a/tidal_dl_ng/download.py b/tidal_dl_ng/download.py index 829d23f..f5b7608 100644 --- a/tidal_dl_ng/download.py +++ b/tidal_dl_ng/download.py @@ -794,43 +794,97 @@ class Download: stream_manifest: StreamManifest | None = None media_stream: Stream | None = None do_flac_extract: bool = False + file_extension: str = "" - if isinstance(media, Track): + # CRITICAL: This lock is intentionally broad and serializes all + # stream-fetching (Phase 1) to prevent a critical race condition. + # + # THE PROBLEM: + # The single, shared session (self.tidal.session) must change its + # credentials to switch between Atmos and Hi-Res/Normal streams. + # + # THE RACE CONDITION IT FIXES: + # If this lock is released *before* get_stream() is called, + # another thread could change the session (e.g., back to "Normal") + # right after this thread switched it to "Atmos". This would + # cause this thread to call get_stream() with the wrong credentials, + # resulting in the API returning AAC 320 instead of Atmos. + # + # THE TRADEOFF: + # This creates a "tollbooth" bottleneck, serializing the get_stream() + # calls. However, the *actual* segment downloads (Phase 2) + # still run in parallel, governed by `downloads_concurrent_max`. + # + # DO NOT "OPTIMIZE" THIS by making the lock more granular. + # Correctness > Performance. + + with self.tidal.stream_lock: try: - if ( - self.settings.data.download_dolby_atmos - and hasattr(media, "audio_modes") - and AudioMode.dolby_atmos.value in media.audio_modes - ): - with self.tidal.atmos_session_context(): - atmos_track = self.session.track(media.id) - media_stream = atmos_track.get_stream() - else: - media_stream = media.get_stream() + if isinstance(media, Track): + stream_manifest, file_extension, do_flac_extract, media_stream = self._get_track_stream_info(media) + if stream_manifest is None: + return None, "", False, None - stream_manifest = media_stream.get_stream_manifest() + elif isinstance(media, Video): + # Videos always require the normal session + if not self.tidal.restore_normal_session(): + self.fn_logger.error(f"Failed to restore normal session for video: {media.id}") + return None, "", False, None + + file_extension = AudioExtensions.MP4 if self.settings.data.video_convert_mp4 else VideoExtensions.TS + + stream_manifest = None + media_stream = None + do_flac_extract = False + + else: + self.fn_logger.error(f"Unknown media type for stream info: {type(media)}") + return None, "", False, None except TooManyRequests: self.fn_logger.exception( f"Too many requests against TIDAL backend. Skipping '{name_builder_item(media)}'. " f"Consider to activate delay between downloads." ) - return None, "", False, None + except Exception: self.fn_logger.exception(f"Something went wrong. Skipping '{name_builder_item(media)}'.") - return None, "", False, None - file_extension = stream_manifest.file_extension + return stream_manifest, file_extension, do_flac_extract, media_stream - if self.settings.data.extract_flac and ( - stream_manifest.codecs.upper() == Codec.FLAC and file_extension != AudioExtensions.FLAC - ): - file_extension = AudioExtensions.FLAC - do_flac_extract = True - elif isinstance(media, Video): - file_extension = AudioExtensions.MP4 if self.settings.data.video_convert_mp4 else VideoExtensions.TS + def _get_track_stream_info(self, media: Track) -> tuple[StreamManifest | None, str, bool, Stream | None]: + """ + Gets stream info for a Track, handling Atmos/Normal session switching. + This is a helper for _get_stream_info to reduce complexity. + """ + want_atmos = ( + self.settings.data.download_dolby_atmos + and hasattr(media, "audio_modes") + and AudioMode.dolby_atmos.value in media.audio_modes + ) + + if want_atmos: + if not self.tidal.switch_to_atmos_session(): + self.fn_logger.error(f"Failed to switch to Atmos session for track: {media.id}") + return None, "", False, None + else: + if not self.tidal.restore_normal_session(): + self.fn_logger.error(f"Failed to restore normal session for track: {media.id}") + return None, "", False, None + + media_stream = self.session.track(media.id).get_stream() if want_atmos else media.get_stream() + + stream_manifest = media_stream.get_stream_manifest() + file_extension = stream_manifest.file_extension + do_flac_extract = False + + if self.settings.data.extract_flac and ( + stream_manifest.codecs.upper() == Codec.FLAC and file_extension != AudioExtensions.FLAC + ): + file_extension = AudioExtensions.FLAC + do_flac_extract = True return stream_manifest, file_extension, do_flac_extract, media_stream