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:
+61
-17
@@ -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
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user