Merge pull request #626 from Rikrdoga:Dolby_Atmos

Fix(Atmos): Resolve Race Condition and Optimize Session Switching for Atmos Downloads
This commit is contained in:
exislow
2025-11-13 08:32:26 +01:00
committed by GitHub
2 changed files with 177 additions and 49 deletions
+61 -17
View File
@@ -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 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 not self.session.check_login():
print("Not logged in.")
Returns:
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
original_client_secret = self.session.config.client_secret
original_audio_quality = self.session.audio_quality
try:
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
# 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:
self.session.config.client_id = original_client_id
self.session.config.client_secret = original_client_secret
self.session.audio_quality = original_audio_quality
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.
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()
+103 -19
View File
@@ -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 (
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):
track_info = self._get_track_stream_info(media)
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
elif isinstance(media, Video):
file_extension = AudioExtensions.MP4 if self.settings.data.video_convert_mp4 else VideoExtensions.TS
requires_flac_extraction = True
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,