Merge pull request #626 from Rikrdoga:Dolby_Atmos

Fix(Atmos): Resolve Race Condition and Optimize Session Switching for Atmos Downloads

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,6 +128,7 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         if settings:
             self.settings = settings
 
+        if not self.is_atmos_session:
             self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
         self.session.video_quality = tidalapi.VideoQuality.high
 
@@ -162,33 +175,64 @@ class Tidal(BaseConfig, metaclass=SingletonMeta):
         self.set_option("expiry_time", self.session.expiry_time)
         self.save()
 
-    @contextmanager
+    def switch_to_atmos_session(self) -> bool:
-    def atmos_session_context(self):
+        """
+        Switches the shared session to Dolby Atmos credentials.
+        Only re-authenticates if not already in Atmos mode.
 
-        if not self.session.check_login():
+        Returns:
-            print("Not logged in.")
+            bool: True if successful or already in Atmos mode, False otherwise.
+        """
+        # If we are already in Atmos mode, do nothing.
+        if self.is_atmos_session:
+            return True
 
-        original_client_id = self.session.config.client_id
+        print("Switching session context to Dolby Atmos...")
-        original_client_secret = self.session.config.client_secret
-        original_audio_quality = self.session.audio_quality
-
-        try:
         self.session.config.client_id = ATMOS_CLIENT_ID
         self.session.config.client_secret = ATMOS_CLIENT_SECRET
         self.session.audio_quality = ATMOS_REQUEST_QUALITY
 
+        # Re-login with new credentials
         if not self.login_token(do_pkce=self.is_pkce):
-                print("Warning: Session restore failed.")
+            print("Warning: Atmos session authentication failed.")
+            # Try to switch back to normal to be safe
+            self.restore_normal_session(force=True)
+            return False
 
-            yield
+        self.is_atmos_session = True  # Set the flag
+        print("Session is now in Atmos mode.")
+        return True
 
-        finally:
+    def restore_normal_session(self, force: bool = False) -> bool:
-            self.session.config.client_id = original_client_id
+        """
-            self.session.config.client_secret = original_client_secret
+        Restores the shared session to the original user credentials.
-            self.session.audio_quality = original_audio_quality
+        Only re-authenticates if not already in Normal mode.
 
+        Args:
+            force: If True, forces restoration even if already in Normal mode.
+
+        Returns:
+            bool: True if successful or already in Normal mode, False otherwise.
+        """
+        # If we are already in Normal mode (and not forced), do nothing.
+        if not self.is_atmos_session and not force:
+            return True
+
+        print("Restoring session context to Normal...")
+        self.session.config.client_id = self.original_client_id
+        self.session.config.client_secret = self.original_client_secret
+
+        # Explicitly restore audio quality to user's configured setting
+        self.session.audio_quality = tidalapi.Quality(self.settings.data.quality_audio)
+
+        # 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()

tidal_dl_ng/download.py

@@ -16,6 +16,7 @@ import tempfile
 import time
 from collections.abc import Callable
 from concurrent import futures
+from dataclasses import dataclass
 from threading import Event
 from uuid import uuid4
 
@@ -75,6 +76,16 @@ from tidal_dl_ng.model.downloader import DownloadSegmentResult
 from tidal_dl_ng.model.gui_data import ProgressBars
 
 
+@dataclass
+class TrackStreamInfo:
+    """Container for track stream information."""
+
+    stream_manifest: StreamManifest | None
+    file_extension: str
+    requires_flac_extraction: bool
+    media_stream: Stream | None
+
+
 # TODO: Set appropriate client string and use it for video download.
 # https://github.com/globocom/m3u8#using-different-http-clients
 class RequestsClient:
@@ -134,7 +145,8 @@ class Download:
         Args:
             tidal_obj (Tidal): TIDAL configuration object. Required for:
                 - session: Main TIDAL API session
-                - atmos_session_context(): Dolby Atmos credential switching
+                - switch_to_atmos_session(): Dolby Atmos credential switching
+                - restore_normal_session(): Restore original session credentials
             path_base (str): Base path for downloads.
             fn_logger (Callable): Logger function or object.
             skip_existing (bool, optional): Whether to skip existing files. Defaults to False.
@@ -794,45 +806,117 @@ 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 (
+                if isinstance(media, Track):
-                    self.settings.data.download_dolby_atmos
+                    track_info = self._get_track_stream_info(media)
-                    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()
 
-                stream_manifest = media_stream.get_stream_manifest()
+                    if track_info.stream_manifest is None:
+                        return None, "", False, None
+
+                    stream_manifest = track_info.stream_manifest
+                    file_extension = track_info.file_extension
+                    do_flac_extract = track_info.requires_flac_extraction
+                    media_stream = track_info.media_stream
+
+                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
 
+        return stream_manifest, file_extension, do_flac_extract, media_stream
+
+    def _get_track_stream_info(self, media: Track) -> TrackStreamInfo:
+        """
+        Gets stream info for a Track, handling Atmos/Normal session switching.
+
+        Args:
+            media: The track to get stream information for.
+
+        Returns:
+            TrackStreamInfo: Container with stream manifest, file extension,
+                            FLAC extraction flag, and media stream object.
+                            Returns TrackStreamInfo with None/empty values if fails.
+        """
+        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 TrackStreamInfo(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 TrackStreamInfo(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
+        requires_flac_extraction = 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
+            requires_flac_extraction = True
-        elif isinstance(media, Video):
-            file_extension = AudioExtensions.MP4 if self.settings.data.video_convert_mp4 else VideoExtensions.TS
 
-        return stream_manifest, file_extension, do_flac_extract, media_stream
+        return TrackStreamInfo(
+            stream_manifest=stream_manifest,
+            file_extension=file_extension,
+            requires_flac_extraction=requires_flac_extraction,
+            media_stream=media_stream,
+        )
 
     def _perform_actual_download(
         self,