diff --git a/py/services/aria2_downloader.py b/py/services/aria2_downloader.py index d4f6a0ec..88e5eb42 100644 --- a/py/services/aria2_downloader.py +++ b/py/services/aria2_downloader.py @@ -161,6 +161,11 @@ class Aria2Downloader: (typically an expired CivitAI signed URL): a fresh URL is resolved and the partial download continues. Recovery is bounded by ``MAX_TRANSFER_RECOVERY_ATTEMPTS``. + + Cancellation never leaks daemon transfers: the gid is tracked in + ``_transfers`` before any post-``addUri`` await, and a gid accepted + by the daemon while the caller is being cancelled is removed again + before the ``CancelledError`` propagates. """ await self._ensure_process() @@ -251,7 +256,11 @@ class Aria2Downloader: await asyncio.sleep(self._poll_interval) finally: current = self._transfers.get(download_id) - if current is not None and current.gid == transfer.gid: + if ( + transfer is not None + and current is not None + and current.gid == transfer.gid + ): self._transfers.pop(download_id, None) async def _get_status_with_retry( @@ -339,21 +348,43 @@ class Aria2Downloader: resolved_url != url, ) + # Shield the addUri RPC from cancellation: the daemon may accept the + # download even when the caller is cancelled while the request is in + # flight. On cancellation, wait for the RPC result so the freshly + # created gid can be removed instead of leaking an untracked + # download that keeps running in the daemon. + add_task = asyncio.ensure_future( + self._rpc_call("aria2.addUri", [[resolved_url], options]) + ) try: - gid = await self._rpc_call("aria2.addUri", [[resolved_url], options]) + gid = await asyncio.shield(add_task) + except asyncio.CancelledError: + leaked_gid: Any = None + try: + leaked_gid = await add_task + except Exception: + leaked_gid = None + if isinstance(leaked_gid, str) and leaked_gid: + logger.info( + "Removing aria2 gid %s accepted while download %s was " + "being cancelled", + leaked_gid, + download_id, + ) + try: + await self._rpc_call("aria2.forceRemove", [leaked_gid]) + except Exception as exc: + logger.warning( + "Failed to remove leaked aria2 gid %s for download %s: %s", + leaked_gid, + download_id, + exc, + ) + raise except Exception as exc: raise Aria2Error(f"Failed to schedule aria2 download: {exc}") from exc logger.debug("aria2 accepted download %s with gid %s", download_id, gid) - await self._state_store.upsert( - download_id, - { - "gid": gid, - "save_path": save_path, - "status": "downloading", - "url": url, - }, - ) return gid async def _register_transfer( @@ -372,7 +403,46 @@ class Aria2Downloader: headers=headers, ) transfer = Aria2Transfer(gid=gid, save_path=os.path.abspath(save_path)) + # Register the transfer before any further await: once the daemon + # holds the gid, cancel_download() must be able to find it. An await + # in between would open a window where a concurrent cancel reports + # "Download task not found" and the daemon keeps downloading + # untracked. self._transfers[download_id] = transfer + try: + await self._state_store.upsert( + download_id, + { + "gid": gid, + "save_path": transfer.save_path, + "status": "downloading", + "url": url, + }, + ) + except asyncio.CancelledError: + # The task was cancelled while persisting state and the + # coordinator's cancel ran before the transfer was registered + # above. Remove the daemon transfer unless it was deliberately + # paused (skip_download preserves paused transfers for resume). + status = None + try: + status = await self.get_status(download_id) + except Exception: + status = None + if status is not None and status.get("status") != "paused": + try: + await self._rpc_call("aria2.forceRemove", [gid]) + except Exception as exc: + logger.warning( + "Failed to remove aria2 gid %s for cancelled download %s: %s", + gid, + download_id, + exc, + ) + current = self._transfers.get(download_id) + if current is not None and current.gid == gid: + self._transfers.pop(download_id, None) + raise return transfer async def get_status(self, download_id: str) -> Optional[Dict[str, Any]]: diff --git a/tests/services/test_aria2_downloader.py b/tests/services/test_aria2_downloader.py index 453ff611..dff0aedc 100644 --- a/tests/services/test_aria2_downloader.py +++ b/tests/services/test_aria2_downloader.py @@ -889,6 +889,215 @@ async def test_cancel_download_tolerates_missing_gid(monkeypatch): assert await downloader._state_store.get("download-1") is None +@pytest.mark.asyncio +async def test_register_transfer_tracks_gid_before_state_persist_completes( + tmp_path, monkeypatch +): + """A cancel arriving while the state store write is still in flight must + already find the transfer — otherwise the gid leaks and the daemon keeps + downloading.""" + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + rpc_calls = [] + + async def fake_rpc_call(method, params, **_kwargs): + rpc_calls.append((method, params)) + if method == "aria2.addUri": + return "gid-1" + if method == "aria2.forceRemove": + return "OK" + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_ensure_process", AsyncMock()) + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + persist_started = asyncio.Event() + persist_release = asyncio.Event() + + class BlockingStore: + async def upsert(self, download_id, payload): + persist_started.set() + await persist_release.wait() + + async def remove(self, download_id): + return None + + monkeypatch.setattr(downloader, "_state_store", BlockingStore()) + + register_task = asyncio.create_task( + downloader._register_transfer( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + ) + await asyncio.wait_for(persist_started.wait(), timeout=1.0) + + transfer = downloader._transfers.get("download-1") + assert transfer is not None and transfer.gid == "gid-1" + + result = await downloader.cancel_download("download-1") + assert result["success"] is True + assert ("aria2.forceRemove", ["gid-1"]) in rpc_calls + + persist_release.set() + registered = await register_task + assert registered.gid == "gid-1" + + +@pytest.mark.asyncio +async def test_schedule_download_removes_gid_accepted_while_cancelled( + tmp_path, monkeypatch +): + """Cancelling while the addUri RPC is in flight must remove the gid the + daemon accepted, instead of leaking an untracked download.""" + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + add_uri_started = asyncio.Event() + force_removed = [] + + async def fake_rpc_call(method, params, **_kwargs): + if method == "aria2.addUri": + add_uri_started.set() + # The daemon processes the request while the client is cancelled. + await asyncio.sleep(0.05) + return "gid-leaked" + if method == "aria2.forceRemove": + force_removed.append(params[0]) + return "OK" + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + schedule_task = asyncio.create_task( + downloader._schedule_download( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + ) + await asyncio.wait_for(add_uri_started.wait(), timeout=1.0) + schedule_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await schedule_task + + assert force_removed == ["gid-leaked"] + assert "download-1" not in downloader._transfers + + +@pytest.mark.asyncio +async def test_register_transfer_cancelled_during_persist_removes_active_gid( + tmp_path, monkeypatch +): + """Cancellation landing after the gid is registered but before the state + store write finishes must remove the still-active daemon transfer.""" + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + persist_started = asyncio.Event() + force_removed = [] + + async def fake_rpc_call(method, params, **_kwargs): + if method == "aria2.addUri": + return "gid-2" + if method == "aria2.tellStatus": + return {"gid": "gid-2", "status": "active"} + if method == "aria2.forceRemove": + force_removed.append(params[0]) + return "OK" + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + class BlockingStore: + async def upsert(self, download_id, payload): + persist_started.set() + await asyncio.Event().wait() + + async def remove(self, download_id): + return None + + monkeypatch.setattr(downloader, "_state_store", BlockingStore()) + + register_task = asyncio.create_task( + downloader._register_transfer( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + ) + await asyncio.wait_for(persist_started.wait(), timeout=1.0) + register_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await register_task + + assert force_removed == ["gid-2"] + assert "download-1" not in downloader._transfers + + +@pytest.mark.asyncio +async def test_register_transfer_cancelled_during_persist_preserves_paused_gid( + tmp_path, monkeypatch +): + """skip_download pauses the daemon transfer before cancelling the task; + the unwind cleanup must not remove a deliberately paused gid.""" + downloader = Aria2Downloader() + downloader._rpc_url = "http://127.0.0.1/jsonrpc" + downloader._rpc_secret = "secret" + + save_path = tmp_path / "downloads" / "model.safetensors" + persist_started = asyncio.Event() + force_removed = [] + + async def fake_rpc_call(method, params, **_kwargs): + if method == "aria2.addUri": + return "gid-3" + if method == "aria2.tellStatus": + return {"gid": "gid-3", "status": "paused"} + if method == "aria2.forceRemove": + force_removed.append(params[0]) + return "OK" + raise AssertionError(f"Unexpected RPC method: {method}") + + monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call) + + class BlockingStore: + async def upsert(self, download_id, payload): + persist_started.set() + await asyncio.Event().wait() + + async def remove(self, download_id): + return None + + monkeypatch.setattr(downloader, "_state_store", BlockingStore()) + + register_task = asyncio.create_task( + downloader._register_transfer( + "https://example.com/model.safetensors", + str(save_path), + download_id="download-1", + ) + ) + await asyncio.wait_for(persist_started.wait(), timeout=1.0) + register_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await register_task + + assert force_removed == [] + assert downloader._transfers["download-1"].gid == "gid-3" + + @pytest.mark.asyncio async def test_rpc_call_suppresses_error_log_when_log_errors_false( monkeypatch, caplog