Stop in-flight downloads when stopping engine

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,16 +176,26 @@ 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:
                 self._inc("_ok")
                 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)

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()