diff --git a/streamrip/client/downloadable.py b/streamrip/client/downloadable.py index 3b29b504..e51771ee 100644 --- a/streamrip/client/downloadable.py +++ b/streamrip/client/downloadable.py @@ -37,29 +37,37 @@ def generate_temp_path(url: str): ) -async def fast_async_download(path, url, headers, callback): - """Synchronous download with yield for every 1MB read. +def _do_download(path, url, headers, callback, loop): + """Blocking download using requests. Runs in a thread pool via fast_async_download. + + Large chunk size avoids the CPU-bound problem caused by yielding to the event loop + on every small read (the original aiohttp/aiofiles approach capped total speed at ~10MB/s). - Using aiofiles/aiohttp resulted in a yield to the event loop for every 1KB, - which made file downloads CPU-bound. This resulted in a ~10MB max total download - speed. This fixes the issue by only yielding to the event loop for every 1MB read. + callback is dispatched back onto the event loop (thread-safe) since rich's Live + display must only be updated from the main thread. """ chunk_size: int = 2**17 # 131 KB - counter = 0 - yield_every = 8 # 1 MB - with open(path, "wb") as file: # noqa: ASYNC101 - with requests.get( # noqa: ASYNC100 + with open(path, "wb") as file: + with requests.get( url, headers=headers, allow_redirects=True, stream=True, ) as resp: + resp.raise_for_status() for chunk in resp.iter_content(chunk_size=chunk_size): file.write(chunk) - callback(len(chunk)) - if counter % yield_every == 0: - await asyncio.sleep(0) - counter += 1 + loop.call_soon_threadsafe(callback, len(chunk)) + + +async def fast_async_download(path, url, headers, callback): + """Run the blocking download in a thread pool so the event loop stays free. + + Keeping requests + large chunks for throughput, but offloading to a thread + so concurrent downloads are not frozen while one connection is active. + """ + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, _do_download, path, url, headers, callback, loop) @dataclass(slots=True) diff --git a/streamrip/media/album.py b/streamrip/media/album.py index 3ca79a29..4c7903f6 100644 --- a/streamrip/media/album.py +++ b/streamrip/media/album.py @@ -1,4 +1,3 @@ -import asyncio import logging import os from dataclasses import dataclass @@ -12,7 +11,7 @@ from ..metadata import AlbumMetadata from ..metadata.util import get_album_track_ids from .artwork import download_artwork -from .media import Media, Pending +from .media import Media, Pending, rip_pending_items_in_order from .track import PendingTrack logger = logging.getLogger("streamrip") @@ -31,22 +30,7 @@ async def preprocess(self): progress.add_title(self.meta.album) async def download(self): - async def _resolve_and_download(pending: Pending): - try: - track = await pending.resolve() - if track is None: - return - await track.rip() - except Exception as e: - logger.error(f"Error downloading track: {e}") - - results = await asyncio.gather( - *[_resolve_and_download(p) for p in self.tracks], return_exceptions=True - ) - - for result in results: - if isinstance(result, Exception): - logger.error(f"Album track processing error: {result}") + await rip_pending_items_in_order(self.tracks, self.config, logger) async def postprocess(self): progress.remove_title(self.meta.album) diff --git a/streamrip/media/media.py b/streamrip/media/media.py index 78ad3945..210c49cc 100644 --- a/streamrip/media/media.py +++ b/streamrip/media/media.py @@ -1,3 +1,5 @@ +import asyncio +import logging from abc import ABC, abstractmethod @@ -30,3 +32,56 @@ class Pending(ABC): async def resolve(self) -> Media | None: """Fetch metadata and resolve into a downloadable `Media` object.""" raise NotImplementedError + + +def _active_rip_limit(config, total: int) -> int: + downloads = config.session.downloads + if not getattr(downloads, "concurrency", False): + return 1 + + max_connections = getattr(downloads, "max_connections", total) + if isinstance(max_connections, int) and max_connections > 0: + return max_connections + + return max(total, 1) + + +async def rip_pending_items_in_order( + pending_items: list[Pending], + config, + logger: logging.Logger, + error_message: str = "Error downloading track", +): + active: set[asyncio.Task] = set() + limit = _active_rip_limit(config, len(pending_items)) + + async def _rip(media: Media): + try: + await media.rip() + except Exception as e: + logger.error("%s: %s", error_message, e) + + async def _wait_for_one(): + nonlocal active + done, active = await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED) + for task in done: + await task + + for pending in pending_items: + try: + media = await pending.resolve() + except Exception as e: + logger.error("%s: %s", error_message, e) + continue + + if media is None: + continue + + active.add(asyncio.create_task(_rip(media))) + await asyncio.sleep(0) + + if len(active) >= limit: + await _wait_for_one() + + if active: + await asyncio.gather(*active) diff --git a/streamrip/media/playlist.py b/streamrip/media/playlist.py index 383f2465..514117af 100644 --- a/streamrip/media/playlist.py +++ b/streamrip/media/playlist.py @@ -26,7 +26,7 @@ ) from ..utils.ssl_utils import get_aiohttp_connector_kwargs from .artwork import download_artwork -from .media import Media, Pending +from .media import Media, Pending, rip_pending_items_in_order from .track import Track logger = logging.getLogger("streamrip") @@ -118,28 +118,7 @@ async def postprocess(self): progress.remove_title(self.name) async def download(self): - track_resolve_chunk_size = 20 - - async def _resolve_download(item: PendingPlaylistTrack): - try: - track = await item.resolve() - if track is None: - return - await track.rip() - except Exception as e: - logger.error(f"Error downloading track: {e}") - - batches = self.batch( - [_resolve_download(track) for track in self.tracks], - track_resolve_chunk_size, - ) - - for batch in batches: - results = await asyncio.gather(*batch, return_exceptions=True) - - for result in results: - if isinstance(result, Exception): - logger.error(f"Batch processing error: {result}") + await rip_pending_items_in_order(self.tracks, self.config, logger) @staticmethod def batch(iterable, n=1): diff --git a/streamrip/media/track.py b/streamrip/media/track.py index bc3207d5..312434ef 100644 --- a/streamrip/media/track.py +++ b/streamrip/media/track.py @@ -40,37 +40,40 @@ async def preprocess(self): async def download(self): # TODO: progress bar description async with global_download_semaphore(self.config.session.downloads): - with get_progress_callback( - self.config.session.cli.progress_bars, - await self.downloadable.size(), - f"Track {self.meta.tracknumber}", - ) as callback: - try: + try: + with get_progress_callback( + self.config.session.cli.progress_bars, + await self.downloadable.size(), + f"Track {self.meta.tracknumber}", + ) as callback: await self.downloadable.download(self.download_path, callback) - retry = False - except Exception as e: - logger.error( - f"Error downloading track '{self.meta.title}', retrying: {e}" - ) - retry = True - - if not retry: - return - - with get_progress_callback( - self.config.session.cli.progress_bars, - await self.downloadable.size(), - f"Track {self.meta.tracknumber} (retry)", - ) as callback: - try: + return + except Exception as e: + logger.error( + "Error downloading track '%s', retrying: %s\nDownload URL: %s", + self.meta.title, + e, + self.downloadable.url, + ) + + try: + with get_progress_callback( + self.config.session.cli.progress_bars, + await self.downloadable.size(), + f"Track {self.meta.tracknumber} (retry)", + ) as callback: await self.downloadable.download(self.download_path, callback) - except Exception as e: - logger.error( - f"Persistent error downloading track '{self.meta.title}', skipping: {e}" - ) - self.db.set_failed( - self.downloadable.source, "track", self.meta.info.id - ) + except Exception as e: + logger.error( + "Persistent error downloading track '%s', skipping: %s\nDownload URL: %s", + self.meta.title, + e, + self.downloadable.url, + ) + self.db.set_failed( + self.downloadable.source, "track", self.meta.info.id + ) + raise async def postprocess(self): if self.is_single: diff --git a/streamrip/rip/cli.py b/streamrip/rip/cli.py index 2b3f532e..41e897b4 100644 --- a/streamrip/rip/cli.py +++ b/streamrip/rip/cli.py @@ -188,9 +188,7 @@ async def url(ctx, urls): version_coro = None async with Main(cfg) as main: - await main.add_all(urls) - await main.resolve() - await main.rip() + await main.rip_urls(urls) if version_coro is not None: latest_version, notes = await version_coro @@ -253,7 +251,8 @@ async def file(ctx, path): console.print( f"Detected list of urls. Loading [yellow]{len(items)}[/yellow] items" ) - await main.add_all(items) + await main.rip_urls(items) + return await main.resolve() await main.rip() diff --git a/streamrip/rip/main.py b/streamrip/rip/main.py index 562773fc..e5ea69bc 100644 --- a/streamrip/rip/main.py +++ b/streamrip/rip/main.py @@ -131,6 +131,48 @@ async def add_all(self, urls: list[str]): ) self.pending.extend(pendings) + async def rip_urls(self, urls: list[str]): + """Resolve and rip URLs as soon as each one is ready.""" + parsed_urls = [] + for i, parsed in enumerate(parse_url(url) for url in urls): + if parsed is None: + console.print( + f"[red]Found invalid url [cyan]{urls[i]}[/cyan], skipping.", + ) + continue + parsed_urls.append((urls[i], parsed)) + + clients = { + source: await self.get_logged_in_client(source) + for source in {parsed.source for _, parsed in parsed_urls} + } + + async def _resolve_and_rip(raw_url: str, parsed): + try: + pending = await parsed.into_pending( + clients[parsed.source], self.config, self.database + ) + media = await pending.resolve() + if media is None: + return None + await media.rip() + except Exception as e: + logger.error("Error processing url %s: %s", raw_url, e) + return e + + results = await asyncio.gather( + *[_resolve_and_rip(raw_url, parsed) for raw_url, parsed in parsed_urls], + return_exceptions=True, + ) + + failed_items = sum(1 for result in results if isinstance(result, Exception)) + if failed_items > 0: + logger.info( + "Download completed with %s failed URLs out of %s total URLs.", + failed_items, + len(parsed_urls), + ) + async def get_logged_in_client(self, source: str): """Return a functioning client instance for `source`.""" client = self.clients.get(source) diff --git a/tests/test_error_handling.py b/tests/test_error_handling.py index 67f6a6e7..2a6e0e08 100644 --- a/tests/test_error_handling.py +++ b/tests/test_error_handling.py @@ -1,10 +1,13 @@ +import asyncio import json +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from streamrip.media.album import Album from streamrip.media.playlist import Playlist +from streamrip.media.track import Track class TestErrorHandling: @@ -70,6 +73,49 @@ async def test_album_handles_failed_track(self): mock_track_success.resolve.return_value.rip.assert_called_once() mock_track_failure.resolve.assert_called_once() + @pytest.mark.asyncio + async def test_album_starts_parallel_downloads_in_track_order(self): + events = [] + mock_config = MagicMock() + mock_config.session.downloads.concurrency = True + mock_config.session.downloads.max_connections = 2 + mock_db = MagicMock() + mock_meta = MagicMock() + + class ResolvedTrack: + def __init__(self, name): + self.name = name + + async def rip(self): + events.append(f"rip {self.name}") + await asyncio.sleep(0) + + class PendingTrack: + def __init__(self, name, delay): + self.name = name + self.delay = delay + + async def resolve(self): + events.append(f"resolve {self.name}") + await asyncio.sleep(self.delay) + return ResolvedTrack(self.name) + + album = Album( + meta=mock_meta, + config=mock_config, + tracks=[ + PendingTrack("01", 0.01), + PendingTrack("02", 0), + ], + folder="/test/folder", + db=mock_db, + ) + + await album.download() + + assert events.index("rip 01") < events.index("resolve 02") + assert events.index("rip 01") < events.index("rip 02") + @pytest.mark.asyncio async def test_main_rip_handles_failed_media(self): """Test that the Main.rip method handles failed media items.""" @@ -103,3 +149,93 @@ async def test_main_rip_handles_failed_media(self): mock_media_success.rip.assert_called_once() mock_media_failure.rip.assert_called_once() + + @pytest.mark.asyncio + async def test_track_download_failure_is_not_marked_downloaded(self): + """Persistent download failures should not be recorded as successful.""" + config = MagicMock() + config.session.cli.progress_bars = False + config.session.downloads.concurrency = False + + meta = MagicMock() + meta.title = "Failed Track" + meta.tracknumber = 1 + meta.info = SimpleNamespace(id="track-id") + + downloadable = MagicMock() + downloadable.url = "https://example.test/file.flac?token=secret&expires=1" + downloadable.source = "qobuz" + downloadable.size = AsyncMock(return_value=100) + downloadable.download = AsyncMock(side_effect=Exception("network error")) + + database = MagicMock() + + track = Track( + meta=meta, + downloadable=downloadable, + config=config, + folder="/tmp", + cover_path=None, + db=database, + download_path="/tmp/failed.flac", + ) + + with pytest.raises(Exception, match="network error"): + await track.download() + + database.set_failed.assert_called_once_with("qobuz", "track", "track-id") + database.set_downloaded.assert_not_called() + + @pytest.mark.asyncio + async def test_rip_urls_starts_ready_url_before_slow_url_resolves(self): + """A ready URL should begin ripping before every URL has resolved.""" + from streamrip.rip.main import Main + + events = [] + config = MagicMock() + config.session.downloads.requests_per_minute = 0 + config.session.database.downloads_enabled = False + config.session.database.failed_downloads_enabled = False + + class Parsed: + source = "qobuz" + + def __init__(self, name): + self.name = name + + async def into_pending(self, *_): + return Pending(self.name) + + class Pending: + def __init__(self, name): + self.name = name + + async def resolve(self): + if self.name == "slow": + await asyncio.sleep(0.05) + events.append((self.name, "resolved")) + return Media(self.name) + + class Media: + def __init__(self, name): + self.name = name + + async def rip(self): + events.append((self.name, "ripped")) + + with ( + patch("streamrip.rip.main.QobuzClient") as qobuz_client, + patch("streamrip.rip.main.TidalClient"), + patch("streamrip.rip.main.DeezerClient"), + patch("streamrip.rip.main.SoundcloudClient"), + patch( + "streamrip.rip.main.parse_url", + side_effect=[Parsed("slow"), Parsed("fast")], + ), + ): + qobuz_client.return_value.logged_in = True + main = Main(config) + + await main.rip_urls(["slow-url", "fast-url"]) + + assert events.index(("fast", "ripped")) < events.index(("slow", "resolved"))