From 16afdcef07b8bd58e1cf2e1ea755fd938f0ef064 Mon Sep 17 00:00:00 2001 From: Rixuan Shao <2023311022@bipt.edu.cn> Date: Tue, 14 Jul 2026 15:23:29 +0800 Subject: [PATCH] fix: harden scheduler and login desktop runtime --- .github/workflows/schedule.yml | 6 +-- DouYinSparkFlow/core/tasks.py | 19 ++++++--- DouYinSparkFlow/login_desktop_server.py | 16 ++++---- .../tests/test_deployment_contract.py | 5 +++ DouYinSparkFlow/tests/test_send_state.py | 25 ++++++++++++ DouYinSparkFlow/tests/test_webui_safety.py | 35 +++++++++++++++- DouYinSparkFlow/webui/app.py | 40 +++++++++++++------ 7 files changed, 116 insertions(+), 30 deletions(-) diff --git a/.github/workflows/schedule.yml b/.github/workflows/schedule.yml index be5c999..905c811 100644 --- a/.github/workflows/schedule.yml +++ b/.github/workflows/schedule.yml @@ -22,9 +22,9 @@ jobs: runs-on: ubuntu-latest environment: user-data steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + - uses: actions/checkout@93cb6efe18208431cddfb8368fd83d5badbf9bfd # v5 - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5 + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 with: python-version: "3.11" - name: Install dependencies @@ -46,7 +46,7 @@ jobs: SPARKFLOW_MANUAL_RUN: "1" PYTHONUNBUFFERED: "1" run: python main.py --doTask - - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4 + - uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6 if: ${{ !cancelled() }} with: name: run-logs diff --git a/DouYinSparkFlow/core/tasks.py b/DouYinSparkFlow/core/tasks.py index f068570..d310a38 100644 --- a/DouYinSparkFlow/core/tasks.py +++ b/DouYinSparkFlow/core/tasks.py @@ -2748,11 +2748,18 @@ async def runTasks(): if not runnable_user_data: return - with task_run_lock(): - protocol_user_data, browser_user_data = _split_sender_modes(active_config, runnable_user_data) - if protocol_user_data: - await run_protocol_tasks(active_config, protocol_user_data, build_message) - await run_browser_tasks(active_config, browser_user_data) + try: + with task_run_lock(): + protocol_user_data, browser_user_data = _split_sender_modes(active_config, runnable_user_data) + if protocol_user_data: + await run_protocol_tasks(active_config, protocol_user_data, build_message) + await run_browser_tasks(active_config, browser_user_data) + except TaskRunAlreadyInProgress: + logger.warning("Skipping task run because another task run is already in progress") + + +class TaskRunAlreadyInProgress(RuntimeError): + """Raised when a live task process already owns the global run lock.""" @contextmanager @@ -2791,7 +2798,7 @@ def task_run_lock(): pass continue - raise RuntimeError("another task run is already in progress") from exc + raise TaskRunAlreadyInProgress("another task run is already in progress") from exc try: handle.write(f"{os.getpid()}\n") diff --git a/DouYinSparkFlow/login_desktop_server.py b/DouYinSparkFlow/login_desktop_server.py index 8b6348b..1100faf 100644 --- a/DouYinSparkFlow/login_desktop_server.py +++ b/DouYinSparkFlow/login_desktop_server.py @@ -2,6 +2,7 @@ import asyncio import os import shutil import time +from contextlib import asynccontextmanager from pathlib import Path from fastapi import FastAPI, HTTPException, Request, Response @@ -414,18 +415,19 @@ async def collect_www_login_result(page, context): manager = LoginDesktopManager() -app = FastAPI(title="Douyin Login Desktop") -@app.on_event("startup") -async def startup(): +@asynccontextmanager +async def lifespan(_app: FastAPI): await manager.start_idle_monitor() + try: + yield + finally: + await manager.stop_idle_monitor() + await manager.stop(clear_profile=False) -@app.on_event("shutdown") -async def shutdown(): - await manager.stop_idle_monitor() - await manager.stop(clear_profile=False) +app = FastAPI(title="Douyin Login Desktop", lifespan=lifespan) @app.get("/health") diff --git a/DouYinSparkFlow/tests/test_deployment_contract.py b/DouYinSparkFlow/tests/test_deployment_contract.py index 8bf0e6e..7531cca 100644 --- a/DouYinSparkFlow/tests/test_deployment_contract.py +++ b/DouYinSparkFlow/tests/test_deployment_contract.py @@ -125,6 +125,11 @@ class DeploymentContractTests(unittest.TestCase): self.assertIn('"/creator-micro/" in page.url', server) self.assertIn('"qr_ready": qr_ready', server) + def test_login_desktop_uses_fastapi_lifespan(self): + server = (SOURCE_ROOT / "login_desktop_server.py").read_text(encoding="utf-8") + self.assertIn("lifespan=lifespan", server) + self.assertNotIn("@app.on_event", server) + def test_login_desktop_exposes_cropped_qr_endpoint(self): server = (SOURCE_ROOT / "login_desktop_server.py").read_text(encoding="utf-8") self.assertIn('@app.get("/qr")', server) diff --git a/DouYinSparkFlow/tests/test_send_state.py b/DouYinSparkFlow/tests/test_send_state.py index 03c5f56..a15e6b1 100644 --- a/DouYinSparkFlow/tests/test_send_state.py +++ b/DouYinSparkFlow/tests/test_send_state.py @@ -1,3 +1,4 @@ +import asyncio import os import unittest from datetime import datetime, timezone @@ -126,6 +127,30 @@ class SendStateTests(unittest.TestCase): self.assertEqual(["confirmed", "pending"], prepared[0]["targets"]) + def test_overlapping_task_run_is_skipped_without_traceback(self): + config = { + "multiTask": False, + "taskCount": 1, + "sendStrategy": {}, + "messageTemplate": "", + "hitokotoTypes": [], + } + user = {"enabled": True, "username": "demo", "targets": ["friend"]} + with ( + patch.object(tasks, "get_config", return_value=config), + patch.object(tasks, "get_userData", return_value=[user]), + patch.object(tasks, "_prepare_active_users_for_run", return_value=[user]), + patch.object( + tasks, + "task_run_lock", + side_effect=tasks.TaskRunAlreadyInProgress("already running"), + ), + patch.object(tasks, "run_browser_tasks") as run_browser, + ): + asyncio.run(tasks.runTasks()) + + run_browser.assert_not_called() + def test_message_choice_avoids_previous_and_last_when_possible(self): with patch.object( msg_builder, diff --git a/DouYinSparkFlow/tests/test_webui_safety.py b/DouYinSparkFlow/tests/test_webui_safety.py index d7ad325..b846e27 100644 --- a/DouYinSparkFlow/tests/test_webui_safety.py +++ b/DouYinSparkFlow/tests/test_webui_safety.py @@ -1,4 +1,5 @@ -import errno +import asyncio +import errno import os import tempfile import time @@ -104,6 +105,38 @@ class WebUiSafetyTests(unittest.TestCase): self.assertEqual(200, response.status_code, path) self.assertEqual("no-store", response.headers["cache-control"]) + def test_login_desktop_timeout_is_wrapped_as_runtime_error(self): + with patch.object(app_module.urllib.request, "urlopen", side_effect=TimeoutError("timed out")): + with self.assertRaisesRegex(RuntimeError, "login-desktop unavailable: timed out"): + app_module.call_login_desktop("/open-login", method="POST", payload={}) + + def test_login_desktop_open_uses_extended_startup_timeout(self): + client = TestClient(app_module.app) + with ( + patch.object(app_module, "current_user", return_value="admin"), + patch.object(app_module, "validate_csrf", return_value=True), + patch.object(app_module, "call_login_desktop", return_value={}) as call_login, + ): + response = client.post("/login-desktop/open", data={"csrf_token": "test"}) + + self.assertEqual(200, response.status_code) + call_login.assert_called_once_with("/open-login", method="POST", payload={}, timeout=90) + + def test_websocket_relay_cleans_up_pending_tasks(self): + cleaned_up = [] + + async def completes(): + return None + + async def waits_forever(): + try: + await asyncio.Event().wait() + finally: + cleaned_up.append(True) + + asyncio.run(app_module._run_websocket_relays(completes(), waits_forever())) + self.assertEqual([True], cleaned_up) + def test_login_desktop_urls_honor_container_environment(self): with ( patch.dict( diff --git a/DouYinSparkFlow/webui/app.py b/DouYinSparkFlow/webui/app.py index 03c59da..90608f5 100644 --- a/DouYinSparkFlow/webui/app.py +++ b/DouYinSparkFlow/webui/app.py @@ -259,8 +259,28 @@ def call_login_desktop(path: str, *, method: str = "GET", payload: dict | None = except urllib.error.HTTPError as exc: body = exc.read().decode("utf-8", errors="replace") raise RuntimeError(f"login-desktop API error {exc.code}: {body}") from exc - except urllib.error.URLError as exc: - raise RuntimeError(f"login-desktop unavailable: {exc.reason}") from exc + except (urllib.error.URLError, TimeoutError) as exc: + reason = getattr(exc, "reason", exc) + raise RuntimeError(f"login-desktop unavailable: {reason}") from exc + + +async def _run_websocket_relays(*coroutines): + tasks = {asyncio.create_task(coroutine) for coroutine in coroutines} + try: + _, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in pending: + task.cancel() + results = await asyncio.gather(*tasks, return_exceptions=True) + for result in results: + if isinstance(result, (ConnectionClosed, WebSocketDisconnect, asyncio.CancelledError)): + continue + if isinstance(result, BaseException): + raise result + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) def save_exported_login_result(login_result: dict, *, relogin_unique_id: str = "", display_name: str = "") -> tuple[dict, str]: @@ -980,16 +1000,10 @@ def create_app(): else: await websocket.send_text(message) - relays = { - asyncio.create_task(client_to_upstream()), - asyncio.create_task(upstream_to_client()), - } - done, pending = await asyncio.wait(relays, return_when=asyncio.FIRST_COMPLETED) - for task in pending: - task.cancel() - await asyncio.gather(*pending, return_exceptions=True) - for task in done: - task.result() + await _run_websocket_relays( + client_to_upstream(), + upstream_to_client(), + ) except (ConnectionClosed, WebSocketDisconnect): pass except Exception as exc: @@ -1054,7 +1068,7 @@ def create_app(): if not validate_csrf(request, str(form.get("csrf_token", ""))): return JSONResponse({"ok": False, "error": "Invalid CSRF token"}, status_code=403) try: - call_login_desktop("/open-login", method="POST", payload={}) + call_login_desktop("/open-login", method="POST", payload={}, timeout=90) return JSONResponse({"ok": True, "public_url": login_desktop_public_url(request)}) except RuntimeError as exc: return JSONResponse({"ok": False, "error": str(exc)}, status_code=503)