diff --git a/docker/crawler/match_live_stream.py b/docker/crawler/match_live_stream.py index a43a508..1d1bc9f 100644 --- a/docker/crawler/match_live_stream.py +++ b/docker/crawler/match_live_stream.py @@ -30,6 +30,11 @@ DEFAULT_HEADERS = { } +def _should_use_internal_lock() -> bool: + flag = str(os.environ.get("DQD_RUNNING_UNDER_TASK_RUNNER") or "").strip().lower() + return flag not in {"1", "true", "yes", "on"} + + def parse_match_live_detail_html(body: str) -> dict: match = DETAIL_JSONP_RE.search((body or "").strip()) if not match: @@ -233,14 +238,16 @@ class MatchLiveStreamRepository: def run(db_config: Dict[str, Any]) -> Dict[str, Any]: - lock = TaskFileLock("match_live_stream") - if not lock.acquire(): - return { - "success": True, - "candidate_count": 0, - "saved_count": 0, - "summary_text": "match_live_stream 已在执行中,跳过本次", - } + lock = None + if _should_use_internal_lock(): + lock = TaskFileLock("match_live_stream") + if not lock.acquire(): + return { + "success": True, + "candidate_count": 0, + "saved_count": 0, + "summary_text": "match_live_stream 已在执行中,跳过本次", + } repo = MatchLiveStreamRepository(db_config) now_ts = int(time.time()) @@ -265,4 +272,5 @@ def run(db_config: Dict[str, Any]) -> Dict[str, Any]: return _build_run_result(candidate_count, success_count, failed_count) finally: repo.close() - lock.release() + if lock is not None: + lock.release() diff --git a/docker/crawler/scripts/docker_task_runner.py b/docker/crawler/scripts/docker_task_runner.py index b0da69e..dc824d2 100644 --- a/docker/crawler/scripts/docker_task_runner.py +++ b/docker/crawler/scripts/docker_task_runner.py @@ -259,6 +259,7 @@ def run_main_action(action: str, timeout: int, output_callback: Callable[[str], start = time.monotonic() env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" + env["DQD_RUNNING_UNDER_TASK_RUNNER"] = "1" proc = subprocess.Popen( [sys.executable, "main.py", action], cwd=str(APP_DIR), diff --git a/docker/crawler/tests/test_match_live_stream.py b/docker/crawler/tests/test_match_live_stream.py index f05544f..3e24ba9 100644 --- a/docker/crawler/tests/test_match_live_stream.py +++ b/docker/crawler/tests/test_match_live_stream.py @@ -2,6 +2,7 @@ import json import sys import unittest from pathlib import Path +from unittest import mock ROOT = Path(__file__).resolve().parents[1] @@ -14,6 +15,7 @@ from match_live_stream import ( choose_next_fetch_at, parse_match_live_detail_html, MatchLiveStreamRepository, + run, ) @@ -188,6 +190,60 @@ class MatchLiveStreamTest(unittest.TestCase): ), ) + def test_run_does_not_skip_due_rows_when_invoked_by_docker_task_runner(self): + class FakeRepo: + def __init__(self): + self.closed = False + self.mark_success_calls = [] + + def load_due_rows(self, now_ts, limit=30): + return [ + { + "id": 1, + "source_url": "https://yyzb1.tv/room/190118?scheduleId=undefined", + "play_url_flv": "", + "play_url_hd_flv": "", + "play_url_m3u8": "", + "play_url_hd_m3u8": "", + } + ] + + def mark_success(self, row_id, stream_urls, now_ts): + self.mark_success_calls.append((row_id, stream_urls, now_ts)) + + def mark_failed(self, row_id, error_message, now_ts, has_existing_stream): + raise AssertionError("docker task runner path should not fall into mark_failed for this test") + + def close(self): + self.closed = True + + class FakeLock: + def acquire(self): + return False + + def release(self): + return None + + repo = FakeRepo() + stream_urls = { + "play_url_flv": "http://stream/flv", + "play_url_hd_flv": "http://stream/hdflv", + "play_url_m3u8": "http://stream/m3u8", + "play_url_hd_m3u8": "http://stream/hdm3u8", + } + + with mock.patch.dict("os.environ", {"DQD_RUNNING_UNDER_TASK_RUNNER": "1"}, clear=False): + with mock.patch("match_live_stream.MatchLiveStreamRepository", return_value=repo): + with mock.patch("match_live_stream.TaskFileLock", return_value=FakeLock()): + with mock.patch("match_live_stream._fetch_stream_payload", return_value=stream_urls): + result = run({"prefix": "la_"}) + + self.assertTrue(result["success"]) + self.assertEqual(result["candidate_count"], 1) + self.assertEqual(result["saved_count"], 1) + self.assertEqual(len(repo.mark_success_calls), 1) + self.assertTrue(repo.closed) + if __name__ == "__main__": unittest.main()