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/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..0aa47246 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: @@ -103,3 +106,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"))