From 39e6872eee284f88f54c8004f4f430e2d3afda36 Mon Sep 17 00:00:00 2001 From: Bai_yang Date: Thu, 6 Aug 2026 09:51:28 +0800 Subject: [PATCH] Stop in-flight downloads when stopping engine --- moon_engine.py | 26 ++++++++++++--- tests/test_exit_cleanup.py | 67 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) diff --git a/moon_engine.py b/moon_engine.py index 5036368..5a19789 100644 --- a/moon_engine.py +++ b/moon_engine.py @@ -138,6 +138,7 @@ def __init__(self): self._thread = None self._loop = None self._gate = None + self._active_kill_events = set() self.proxy_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "proxies.txt") self._proxy_mtime = 0.0 @@ -175,9 +176,15 @@ async def _do_dl(self, proxy_url, cookies, filename, orig_url, rec, kc = kill_counts.get(orig_url, 0) kill_evt = asyncio.Event() - ok, msg, bytes_done = await download_file( - proxy_url, cookies, dest, rec, self._bytes_acc, kill_evt, kc, - telem=telem, on_event=self.log) + with self._lock: + self._active_kill_events.add(kill_evt) + try: + ok, msg, bytes_done = await download_file( + proxy_url, cookies, dest, rec, self._bytes_acc, kill_evt, kc, + telem=telem, on_event=self.log) + finally: + with self._lock: + self._active_kill_events.discard(kill_evt) rec.dl_s = max(time.monotonic()-rec.dl_start, 0.001) if ok: @@ -185,6 +192,10 @@ async def _do_dl(self, proxy_url, cookies, filename, orig_url, rec, spd = f" ({rec.avg_mbs:.1f} MB/s)" if rec.avg_mbs > 0 else "" self.log(f" ✓ Saved: {filename}{spd}", "ok") rec.status="ok"; mark_done_fn(); self._inc("_dl_done") + elif msg == "stall_killed" and self._get("_stop_flag"): + rec.status = "stopped" + self.log(f" stop: {filename} interrupted; partial file kept", "warn") + mark_done_fn(); self._inc("_dl_done") elif msg == "stall_killed": done_mb = bytes_done//(1<<20) new_kc = kc + 1; kill_counts[orig_url] = new_kc @@ -518,6 +529,7 @@ def start(self, cfg: dict) -> dict: self._bytes_acc.clear(); self._t0 = time.monotonic() self._t_end = 0.0 self._tracked.clear() + self._active_kill_events.clear() with self._log_lock: self._log_ring.clear(); self._log_total = 0 @@ -563,16 +575,22 @@ def _guarded_run(self, urls, n, d, r): def stop(self, timeout: float = 1.5) -> dict: running = self._get("_running") + kill_events = [] if running: with self._lock: self._stop_flag = True self._state = "stopping" - self.log("⏹ stop requested — finishing the downloads in flight...", "warn") + kill_events = list(self._active_kill_events) + self.log("⏹ stop requested — stopping downloads in flight...", "warn") with self._lock: loop, gate, thread = self._loop, self._gate, self._thread deadline = time.monotonic() + max(0.0, timeout) + if running and loop is not None and loop.is_running(): + for kill_evt in kill_events: + loop.call_soon_threadsafe(kill_evt.set) + if running and loop is not None and gate is not None and loop.is_running(): try: future = asyncio.run_coroutine_threadsafe(gate.aclose(), loop) diff --git a/tests/test_exit_cleanup.py b/tests/test_exit_cleanup.py index 3083eae..6b71531 100644 --- a/tests/test_exit_cleanup.py +++ b/tests/test_exit_cleanup.py @@ -1,6 +1,8 @@ from __future__ import annotations import asyncio +import pathlib +import threading import time import moon_engine @@ -50,3 +52,68 @@ async def blocked_extract(browser, url): assert closed_before_return == 1 assert not thread_alive_before_return assert elapsed < 1.5 + + +def test_engine_stop_aborts_inflight_download_without_counting_failure( + browser_calls, monkeypatch, tmp_path +): + download_started = threading.Event() + release_download = threading.Event() + failed_links = pathlib.Path(moon_engine.__file__).with_name("failed_links.txt") + failed_links.unlink(missing_ok=True) + + async def fake_download( + proxy_url, + cookies, + dest, + rec, + bytes_acc, + kill_evt, + kills_so_far, + telem=None, + on_event=None, + ): + download_started.set() + partial_path = pathlib.Path(dest + ".tmp") + partial_path.write_bytes(b"partial") + while not release_download.is_set(): + if kill_evt.is_set(): + return False, "stall_killed", 4096 + await asyncio.sleep(0.01) + return True, "ok", 4096 + + monkeypatch.setattr(moon_engine, "download_file", fake_download) + + engine = moon_engine.Engine() + result = engine.start( + { + "links": ["https://fuckingfast.co/example/file.zip"], + "mode": "download", + "out_folder": str(tmp_path), + "workers": 1, + "dl_streams": 1, + "retries": 0, + } + ) + assert result == {"ok": True, "proxies": 0, "effective": result["effective"]} + assert download_started.wait(timeout=2) + + try: + started = time.monotonic() + assert engine.stop(timeout=1.5) == {"ok": True} + elapsed = time.monotonic() - started + thread_alive_before_cleanup = engine._thread.is_alive() + snapshot = engine.snapshot() + partial_files = list(tmp_path.glob("*.tmp")) + finally: + release_download.set() + engine.stop() + engine._thread.join(timeout=2) + + assert not thread_alive_before_cleanup + assert elapsed < 1.5 + assert snapshot["state"] == "done" + assert snapshot["metrics"]["fail"] == 0 + assert snapshot["metrics"]["kills"] == 0 + assert partial_files + assert not failed_links.exists()