From d9d362c9c936c9f76b30b012ea04b0517e20c97a Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sun, 9 Aug 2026 12:48:36 +0800 Subject: [PATCH] fix(download): self-heal aria2 transfers lost on daemon restart --- py/services/aria2_downloader.py | 86 ++++++-- tests/services/test_aria2_downloader.py | 254 +++++++++++++++++++++++- 2 files changed, 325 insertions(+), 15 deletions(-) diff --git a/py/services/aria2_downloader.py b/py/services/aria2_downloader.py index 9da16b43..4de53e65 100644 --- a/py/services/aria2_downloader.py +++ b/py/services/aria2_downloader.py @@ -24,6 +24,10 @@ from .settings_manager import get_settings_manager logger = logging.getLogger(__name__) +# Maximum times the download poll loop will re-schedule a transfer after it +# is lost (daemon restart / RPC outage) before failing the download. +MAX_TRANSFER_RECOVERY_ATTEMPTS = 2 + def _try_certifi_ca_path() -> str | None: """Return the certifi CA bundle path if available, else None.""" try: @@ -85,6 +89,7 @@ class Aria2Downloader: self._rpc_session: Optional[aiohttp.ClientSession] = None self._rpc_session_lock = asyncio.Lock() self._process_lock = asyncio.Lock() + self._register_lock = asyncio.Lock() self._transfers: Dict[str, Aria2Transfer] = {} self._poll_interval = 0.5 self._state_store = Aria2TransferStateStore() @@ -103,26 +108,58 @@ class Aria2Downloader: progress_callback=None, headers: Optional[Dict[str, str]] = None, ) -> Tuple[bool, str]: - """Download a file using aria2 RPC and wait for completion.""" + """Download a file using aria2 RPC and wait for completion. + + The poll loop is self-healing: when the in-memory transfer entry + disappears (e.g. another download restarted the daemon and + ``close()`` cleared ``_transfers``) or the RPC becomes unreachable, + the transfer is re-scheduled with ``continue=true`` so the download + resumes from the on-disk ``.aria2`` control file. Recovery is bounded + by ``MAX_TRANSFER_RECOVERY_ATTEMPTS``. + """ await self._ensure_process() save_path = os.path.abspath(save_path) - transfer = self._transfers.get(download_id) - if transfer is None or os.path.abspath(transfer.save_path) != save_path: - gid = await self._schedule_download( - url, - save_path, - download_id=download_id, - headers=headers, - ) - transfer = Aria2Transfer(gid=gid, save_path=save_path) - self._transfers[download_id] = transfer + async with self._register_lock: + transfer = self._transfers.get(download_id) + if transfer is None or os.path.abspath(transfer.save_path) != save_path: + transfer = await self._register_transfer( + url, + save_path, + download_id=download_id, + headers=headers, + ) + + recovery_attempts = 0 try: while True: - status = await self._get_status_with_retry(download_id) + try: + status = await self._get_status_with_retry(download_id) + except Aria2Error: + status = None + if status is None: - return False, "aria2 download not found" + if recovery_attempts >= MAX_TRANSFER_RECOVERY_ATTEMPTS: + return False, "aria2 download not found" + recovery_attempts += 1 + logger.warning( + "aria2 transfer %s lost; re-scheduling with resume " + "(attempt %d/%d)", + download_id, + recovery_attempts, + MAX_TRANSFER_RECOVERY_ATTEMPTS, + ) + await asyncio.sleep(1.0) + await self._ensure_process() + async with self._register_lock: + transfer = await self._register_transfer( + url, + save_path, + download_id=download_id, + headers=headers, + ) + continue snapshot = self._build_progress_snapshot(status) if progress_callback is not None: @@ -139,7 +176,9 @@ class Aria2Downloader: await asyncio.sleep(self._poll_interval) finally: - self._transfers.pop(download_id, None) + current = self._transfers.get(download_id) + if current is not None and current.gid == transfer.gid: + self._transfers.pop(download_id, None) async def _get_status_with_retry( self, download_id: str, *, max_retries: int = 4, retry_delay: float = 3.0 @@ -242,6 +281,25 @@ class Aria2Downloader: ) return gid + async def _register_transfer( + self, + url: str, + save_path: str, + *, + download_id: str, + headers: Optional[Dict[str, str]] = None, + ) -> Aria2Transfer: + """Schedule a download and track it in the in-memory transfer registry.""" + gid = await self._schedule_download( + url, + save_path, + download_id=download_id, + headers=headers, + ) + transfer = Aria2Transfer(gid=gid, save_path=os.path.abspath(save_path)) + self._transfers[download_id] = transfer + return transfer + async def get_status(self, download_id: str) -> Optional[Dict[str, Any]]: """Return the raw aria2 status payload for a known download.""" diff --git a/tests/services/test_aria2_downloader.py b/tests/services/test_aria2_downloader.py index bf6763e2..58b9958b 100644 --- a/tests/services/test_aria2_downloader.py +++ b/tests/services/test_aria2_downloader.py @@ -6,7 +6,12 @@ from unittest.mock import AsyncMock import pytest -from py.services.aria2_downloader import Aria2Downloader, Aria2Error, Aria2Transfer +from py.services.aria2_downloader import ( + Aria2Downloader, + Aria2Error, + Aria2Transfer, + MAX_TRANSFER_RECOVERY_ATTEMPTS, +) from py.services.aria2_transfer_state import Aria2TransferStateStore from py.services import aria2_transfer_state @@ -246,6 +251,253 @@ async def test_download_file_reuses_existing_transfer_without_add_uri( assert [call[0] for call in rpc_calls] == ["aria2.tellStatus", "aria2.tellStatus"] +@pytest.mark.asyncio +async def test_download_file_recovers_when_transfer_lost_mid_poll( + tmp_path, monkeypatch +): + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + add_uri_count = {"n": 0} + poll_count = {"n": 0} + + async def fake_rpc_call(method, params): + if method == "aria2.addUri": + add_uri_count["n"] += 1 + return "gid-1" if add_uri_count["n"] == 1 else "gid-2" + if method == "aria2.tellStatus": + poll_count["n"] += 1 + if poll_count["n"] == 1: + # Simulate a concurrent close() wiping the transfer mid-poll. + downloader._transfers.pop("download-1", None) + return { + "gid": "gid-1", + "status": "active", + "completedLength": "5", + "totalLength": "10", + "downloadSpeed": "25", + } + return { + "gid": "gid-2", + "status": "complete", + "completedLength": "10", + "totalLength": "10", + "downloadSpeed": "0", + "files": [{"path": str(save_path)}], + } + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + success, result = await downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + + assert success is True + assert result == str(save_path) + assert add_uri_count["n"] == 2 + assert downloader._transfers == {} + + +@pytest.mark.asyncio +async def test_download_file_recovers_when_rpc_fails_mid_poll(tmp_path, monkeypatch): + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + add_uri_count = {"n": 0} + poll_count = {"n": 0} + + async def fake_rpc_call(method, params): + if method == "aria2.addUri": + add_uri_count["n"] += 1 + return "gid-1" if add_uri_count["n"] == 1 else "gid-2" + raise AssertionError(f"Unexpected RPC method: {method}") + + async def fake_get_status_with_retry(download_id): + poll_count["n"] += 1 + if poll_count["n"] == 1: + raise Aria2Error( + "Failed to query aria2 download status after 4 attempts: boom" + ) + return { + "gid": "gid-2", + "status": "complete", + "completedLength": "10", + "totalLength": "10", + "downloadSpeed": "0", + "files": [{"path": str(save_path)}], + } + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr(downloader, "_get_status_with_retry", fake_get_status_with_retry) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + success, result = await downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + + assert success is True + assert result == str(save_path) + assert add_uri_count["n"] == 2 + assert downloader._transfers == {} + + +@pytest.mark.asyncio +async def test_download_file_fails_after_recovery_attempts_exhausted( + tmp_path, monkeypatch +): + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + add_uri_count = {"n": 0} + + async def fake_rpc_call(method, params): + if method == "aria2.addUri": + add_uri_count["n"] += 1 + return f"gid-{add_uri_count['n']}" + raise AssertionError(f"Unexpected RPC method: {method}") + + async def fake_get_status(download_id): + return None # transfer never tracked / always lost + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr(downloader, "get_status", fake_get_status) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + success, result = await downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + + assert success is False + assert result == "aria2 download not found" + assert add_uri_count["n"] == 1 + MAX_TRANSFER_RECOVERY_ATTEMPTS + assert downloader._transfers == {} + + +@pytest.mark.asyncio +async def test_download_file_concurrent_same_id_schedules_once(tmp_path, monkeypatch): + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + add_uri_count = {"n": 0} + poll_count = {"n": 0} + + async def fake_rpc_call(method, params): + if method == "aria2.addUri": + add_uri_count["n"] += 1 + return "gid-1" + if method == "aria2.tellStatus": + poll_count["n"] += 1 + if poll_count["n"] < 4: + return { + "gid": "gid-1", + "status": "active", + "completedLength": "5", + "totalLength": "10", + "downloadSpeed": "25", + } + return { + "gid": "gid-1", + "status": "complete", + "completedLength": "10", + "totalLength": "10", + "downloadSpeed": "0", + "files": [{"path": str(save_path)}], + } + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + results = await asyncio.gather( + downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ), + downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ), + ) + + assert all(success for success, _ in results) + assert all(result == str(save_path) for _, result in results) + assert add_uri_count["n"] <= 2 + assert downloader._transfers == {} + + +@pytest.mark.asyncio +async def test_download_file_cleanup_preserves_newer_registration(tmp_path, monkeypatch): + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + poll_count = {"n": 0} + + async def fake_rpc_call(method, params): + if method == "aria2.addUri": + return "gid-1" + if method == "aria2.tellStatus": + poll_count["n"] += 1 + if poll_count["n"] == 1: + # Simulate another invocation registering its own transfer. + downloader._transfers["download-1"] = Aria2Transfer( + gid="gid-new", save_path=str(save_path) + ) + return { + "gid": "gid-1", + "status": "active", + "completedLength": "5", + "totalLength": "10", + "downloadSpeed": "25", + } + return { + "gid": "gid-new", + "status": "complete", + "completedLength": "10", + "totalLength": "10", + "downloadSpeed": "0", + "files": [{"path": str(save_path)}], + } + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock()) + + success, result = await downloader.download_file( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + + assert success is True + assert result == str(save_path) + assert downloader._transfers["download-1"].gid == "gid-new" + + def test_build_progress_snapshot_normalizes_numeric_fields(): downloader = Aria2Downloader()