Resolve Race Condition and Optimize Session Switching for Atmos Downloads

This commit is contained in:
Rikrdoga
2025-11-04 12:36:27 -06:00
parent 6e665bcd45
commit af58b0a83d
2 changed files with 134 additions and 45 deletions
+58 -23
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,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()
+76 -22
View File
@@ -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