diff --git a/moon_engine.py b/moon_engine.py index 54eb1e0..ca063ca 100644 --- a/moon_engine.py +++ b/moon_engine.py @@ -132,6 +132,8 @@ def __init__(self): self._alive = True self._thread = None + self._loop = None + self._gate = None def _inc(self, attr, delta=1): with self._lock: setattr(self, attr, getattr(self, attr) + delta) @@ -335,11 +337,17 @@ async def snap_task(): # fuckingfast is pure HTTP: no browser, no profile, not even the Playwright # driver. Chrome opens on the first datanodes link and not before - a # fuckingfast-only batch never launches one. + browser_started = False + def _chrome_starting(): + nonlocal browser_started + browser_started = True self._inc("_browsers") self.log(" datanodes: starting Chrome...", "dim") gate = BrowserGate(LAUNCH_ARGS, on_open=_chrome_starting) + with self._lock: + self._loop, self._gate = asyncio.get_running_loop(), gate async def _launch(wid): await self._browser_worker( @@ -350,7 +358,7 @@ async def _launch(wid): try: await asyncio.gather(*[asyncio.create_task(_launch(i)) for i in range(n_workers)]) finally: - if gate.opened: + if browser_started: self._inc("_browsers", -1) await gate.aclose() @@ -515,14 +523,33 @@ def _guarded_run(self, urls, n, d, r): # would otherwise vanish silently instead of surfacing to the GUI/CLI. self.log(f"✗ engine crash: {traceback.format_exc(limit=3)}", "fail") self._on_done() + finally: + with self._lock: + self._loop = None + self._gate = None + + def stop(self, timeout: float = 1.5) -> dict: + running = self._get("_running") + if running: + with self._lock: + self._stop_flag = True + self._state = "stopping" + self.log("⏹ stop requested — finishing the downloads in flight...", "warn") - def stop(self) -> dict: - if not self._get("_running"): - return {"ok": True} with self._lock: - self._stop_flag = True - self._state = "stopping" - self.log("⏹ stop requested — finishing the downloads in flight...", "warn") + loop, gate, thread = self._loop, self._gate, self._thread + deadline = time.monotonic() + max(0.0, timeout) + + 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) + future.result(timeout=max(0.0, deadline - time.monotonic())) + except Exception: + pass + + if (running and thread is not None and thread is not threading.current_thread() + and thread.is_alive()): + thread.join(timeout=max(0.0, deadline - time.monotonic())) return {"ok": True} def _files_payload(self) -> list[dict]: diff --git a/moon_extract.py b/moon_extract.py index 2cb2c1b..a7cd635 100644 --- a/moon_extract.py +++ b/moon_extract.py @@ -1136,7 +1136,8 @@ class BrowserGate: inside a single asyncio.run()/engine run, never hand one across two.* """ - __slots__ = ("_args", "_headless", "_on_open", "_lock", "_pw", "_browser", "_shared") + __slots__ = ("_args", "_headless", "_on_open", "_lock", "_pw", "_browser", "_shared", + "_closed") def __init__(self, launch_args: list[str], *, headless: bool | None = None, on_open=None) -> None: @@ -1147,6 +1148,7 @@ def __init__(self, launch_args: list[str], *, headless: bool | None = None, self._pw = None self._browser = None self._shared = False + self._closed = False @property def opened(self) -> bool: @@ -1154,9 +1156,13 @@ def opened(self) -> bool: async def get(self): """The shared browser, launching Playwright + Chrome on the first call.""" + if self._closed: + raise RuntimeError("browser gate is closed") if self._browser is not None: return self._browser async with self._lock: + if self._closed: + raise RuntimeError("browser gate is closed") if self._browser is None: if self._on_open is not None: self._on_open() @@ -1171,6 +1177,7 @@ async def aclose(self) -> None: """Tear down browser then driver, in that order. No-op if nothing opened.""" async with self._lock: browser, pw, shared = self._browser, self._pw, self._shared + self._closed = True self._browser = None self._pw = None if browser is not None: diff --git a/tests/test_exit_cleanup.py b/tests/test_exit_cleanup.py new file mode 100644 index 0000000..3083eae --- /dev/null +++ b/tests/test_exit_cleanup.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import asyncio +import time + +import moon_engine + + +def test_engine_stop_closes_browser_before_returning(browser_calls, monkeypatch, tmp_path): + release = False + + async def blocked_extract(browser, url): + nonlocal release + while not release and browser_calls["shutdown_chrome"] == 0: + await asyncio.sleep(0.01) + if browser_calls["shutdown_chrome"]: + raise RuntimeError("browser closed") + return url + "?fake", "" + + monkeypatch.setattr(moon_engine, "extract_datanodes", blocked_extract) + engine = moon_engine.Engine() + result = engine.start( + { + "links": ["https://datanodes.to/example/file.zip"], + "mode": "links", + "out_folder": str(tmp_path), + "workers": 2, + "dl_streams": 2, + "retries": 0, + } + ) + assert result == {"ok": True, "proxies": 0, "effective": result["effective"]} + + deadline = time.monotonic() + 2 + while browser_calls["open_browser"] == 0 and time.monotonic() < deadline: + time.sleep(0.01) + assert browser_calls["open_browser"] == 1 + + try: + started = time.monotonic() + assert engine.stop(timeout=1.5) == {"ok": True} + elapsed = time.monotonic() - started + closed_before_return = browser_calls["shutdown_chrome"] + thread_alive_before_return = engine._thread.is_alive() + finally: + release = True + engine.stop() + engine._thread.join(timeout=2) + + assert closed_before_return == 1 + assert not thread_alive_before_return + assert elapsed < 1.5