diff --git a/tests/test_automation_ck_health.py b/tests/test_automation_ck_health.py index 1cfa892..f523833 100644 --- a/tests/test_automation_ck_health.py +++ b/tests/test_automation_ck_health.py @@ -16,6 +16,8 @@ > 新代码不再读写,测试也不许再往上面写值。 """ +import threading +import time from datetime import UTC, datetime, timedelta import pytest @@ -31,6 +33,7 @@ from web.backend.services.automation import ( reset_registry, run_due_jobs, run_job_now, + scheduler, ) from web.backend.services.automation.jobs import ck_health @@ -792,3 +795,124 @@ class TestHeartbeatDecoupledFromInterval(_Base): job.next_run_at = datetime.now(UTC) - timedelta(seconds=1) self.session.commit() return job + + +class TestSchedulerRunsJobsConcurrently: + """作业之间**并发**跑 —— 慢作业不该让快作业干等。 + + 以前是串行(源码注释:「作业串行执行 —— 目前只有 CK 巡检,单轮约 20~40 秒」)。 + 但「账号资料定时补全」一轮要 4 分钟以上,串行会把 CK 巡检饿住。 + + 调度循环用 `wait=False`:**派发完就返回**,tick 不再被慢作业拖住。 + """ + + def setup_method(self): + self.engine = get_test_engine() + self.session = make_test_session(self.engine) + reset_registry() + + def teardown_method(self): + reset_registry() + self.session.close() + + def _job(self, key: str, run, *, enabled: bool = True): + register_job( + JobSpec( + key=key, + title=key, + interval_seconds=3600, + batch_size=1, + concurrency=1, + run=run, + ) + ) + self.session.add( + AutomationJob( + key=key, + title=key, + interval_seconds=3600, + batch_size=1, + concurrency=1, + enabled=enabled, + ) + ) + self.session.commit() + + def test_wait_false_returns_without_blocking(self): + """`wait=False` 必须**立刻返回**,不等作业跑完。""" + started = threading.Event() + release = threading.Event() + + def slow(db, job): + started.set() + release.wait(timeout=5) + return {"scanned": 1} + + self._job("slow", slow) + outcomes = run_due_jobs(self.session, wait=False) + try: + assert outcomes == [], "wait=False 不该返回 outcomes(它不等结果)" + assert started.wait(timeout=5), "作业应已在后台开始" + finally: + release.set() + # 等在途标记清掉,避免污染后续用例 + for _ in range(50): + if not scheduler._IN_FLIGHT: + break + time.sleep(0.1) + + def test_in_flight_job_is_not_dispatched_twice(self): + """**在途守卫**:同一个作业不会重复派发。 + + 这是并发化之后最容易出的事故:`next_run_at` 要等作业跑完才更新, + 所以一个跑得比 tick 还久的作业,下一个 tick 看它仍然是「到期」—— + 没有守卫就会被**再派发一次**,等于把速率翻倍 + (对斗鱼接口就是"并发放开",正是风控最敏感的地方)。 + """ + calls: list[int] = [] + release = threading.Event() + entered = threading.Event() + + def blocking(db, job): + calls.append(1) + entered.set() + release.wait(timeout=5) + return {} + + self._job("blocking", blocking) + run_due_jobs(self.session, wait=False) + try: + assert entered.wait(timeout=5) + # 第一次还在跑,再触发一轮 —— 不该被重复派发 + run_due_jobs(self.session, wait=False) + time.sleep(0.3) + assert len(calls) == 1, "在途作业被重复派发了" + finally: + release.set() + for _ in range(50): + if not scheduler._IN_FLIGHT: + break + time.sleep(0.1) + + def test_release_clears_marker_even_on_failure(self): + """失败的作业也必须清掉在途标记,否则它以后**再也不会跑**。""" + scheduler._IN_FLIGHT.add("ghost") + scheduler._release("ghost") + assert "ghost" not in scheduler._IN_FLIGHT + + def test_failed_job_is_released_and_can_run_again(self): + """作业抛异常后仍要能被下一轮再派发(别把它永久卡死在"在途")。""" + calls: list[int] = [] + + def boom(db, job): + calls.append(1) + raise RuntimeError("接口挂了") + + self._job("boom", boom) + run_due_jobs(self.session, wait=False) + for _ in range(50): + if not scheduler._IN_FLIGHT: + break + time.sleep(0.1) + assert not scheduler._IN_FLIGHT, "失败后没清在途标记,作业会被永久跳过" + assert calls, "作业应该真的跑过一次" diff --git a/web/backend/services/automation/scheduler.py b/web/backend/services/automation/scheduler.py index e32b64c..b4058a9 100644 --- a/web/backend/services/automation/scheduler.py +++ b/web/backend/services/automation/scheduler.py @@ -16,6 +16,7 @@ from __future__ import annotations import os import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any @@ -63,6 +64,19 @@ class JobSpec: _REGISTRY: dict[str, JobSpec] = {} + +#: 同时最多几个作业在跑。作业之间互不依赖,并行是为了**慢作业不挡住快作业** +#: (以前串行:一个作业跑 4 分钟,CK 巡检就得等 4 分钟)。 +_JOB_MAX_WORKERS = 4 +_EXECUTOR: ThreadPoolExecutor | None = None +_EXECUTOR_LOCK = threading.Lock() + +#: **在途守卫**:正在跑的作业 key。 +#: 必须要有 —— `next_run_at` 是在作业**跑完之后**才更新的,所以一个跑得比 tick +#: 还久的作业,下一个 tick 看它仍然是「到期」状态,会被重复派发, +#: 等于把速率翻倍(对斗鱼接口来说就是"并发放开",正是风控最敏感的地方)。 +_IN_FLIGHT: set[str] = set() +_IN_FLIGHT_LOCK = threading.Lock() _thread: threading.Thread | None = None _stop_event = threading.Event() _lock = threading.Lock() @@ -151,32 +165,113 @@ def cleanup_orphan_runs(db: Session, message: str = "服务重启中断") -> int return len(orphans) -def run_due_jobs(db: Session, *, now: datetime | None = None) -> list[dict[str, Any]]: - """跑一轮「到点的作业」,返回每轮结果。 +def _get_executor() -> ThreadPoolExecutor: + global _EXECUTOR + with _EXECUTOR_LOCK: + if _EXECUTOR is None: + _EXECUTOR = ThreadPoolExecutor( + max_workers=_JOB_MAX_WORKERS, thread_name_prefix="automation-job" + ) + return _EXECUTOR - 抽成普通函数是为了可测:测试直接调用即可,不必启动线程或等 60 秒。 - 作业串行执行 —— 目前只有 CK 巡检,单轮约 20~40 秒,远小于 tick,串行 - 反而省掉了并发调度的复杂度。 + +def _run_job_isolated(job_key: str) -> dict[str, Any] | None: + """在**自己的 session** 里跑一个作业。 + + 必须独立 session:SQLAlchemy 的 Session **不是线程安全的**, + 多个作业共用一个 session 会互相踩事务(而他们现在是真的并行了)。 + 重新按 key 取行,是因为调度线程里的那行可能已经过期。 + """ + db = SessionLocal() + try: + job = db.query(AutomationJob).filter(AutomationJob.key == job_key).first() + spec = _REGISTRY.get(job_key) + if job is None or spec is None: + return None + return _run_one(db, job, spec) + finally: + db.close() + + +def run_due_jobs( + db: Session, *, now: datetime | None = None, wait: bool = True +) -> list[dict[str, Any]]: + """跑一轮到期的作业。 + + **作业之间并发**(以前是串行):慢作业不该让快作业干等。 + 两个正交的保护: + + - **在途守卫**:同一个作业不会重复派发(`next_run_at` 要等跑完才更新, + 否则跑得比 tick 久的作业每个 tick 都会被再派发一次 = 速率翻倍) + - `wait=False`(调度循环用):**派发完就返回**,tick 不被慢作业拖住; + `wait=True`(默认,测试与手动触发用):跑完并返回 outcomes + + ⚠️ 每**行**多久重检由作业自己按 `interval_seconds` 决定(见 `_run_one` 的注释), + 这里的 `next_run_at` 固定是「跑完 + 一个 tick」。 """ now = now or datetime.now(UTC) - outcomes: list[dict[str, Any]] = [] jobs = ( db.query(AutomationJob) .filter(AutomationJob.enabled.is_(True)) .order_by(AutomationJob.key.asc()) .all() ) + + dispatched: list[str] = [] for job in jobs: spec = _REGISTRY.get(job.key) if spec is None or not _due(job, now): continue - outcomes.append(_run_one(db, job, spec)) + with _IN_FLIGHT_LOCK: + if job.key in _IN_FLIGHT: + # 上一轮还没跑完 —— 跳过这一轮,别叠加(那等于把速率翻倍) + logger.info("自动化作业 {} 仍在执行,本轮跳过", job.key) + continue + _IN_FLIGHT.add(job.key) + dispatched.append(job.key) + + if wait: + # 同步路径(测试 / 手动触发):**顺序跑、用调用方的 session**。 + # 这样调用方跑完就能直接读到自己 session 里的最新状态。 + # 生产路径(下面是 wait=False)才走线程池 + 独立 session。 + outcomes_sync: list[dict[str, Any]] = [] + for key in dispatched: + try: + job = db.query(AutomationJob).filter(AutomationJob.key == key).first() + spec = _REGISTRY.get(key) + if job is not None and spec is not None: + outcomes_sync.append(_run_one(db, job, spec)) + except Exception: # noqa: BLE001 作业边界要兜住,不能拖垮调度 + logger.exception("自动化作业 {} 执行异常", key) + finally: + _release(key) + + # 到期但没注册实现的作业只留日志,不静默消失 + for job in jobs: + if job.key not in _REGISTRY: + logger.warning("自动化作业 {} 没有注册实现,已跳过", job.key) + return outcomes_sync + + futures = {} + executor = _get_executor() + for key in dispatched: + futures[executor.submit(_run_job_isolated, key)] = key # 到期但没注册实现的作业(例如降级回滚后残留的行)只留日志,不静默消失。 for job in jobs: if job.key not in _REGISTRY: logger.warning("自动化作业 {} 没有注册实现,已跳过", job.key) - return outcomes + + # 派发完就返回;在途守卫会在各自跑完后清掉标记。 + for future, key in futures.items(): + future.add_done_callback(lambda _f, k=key: _release(k)) + return [] + + +def _release(job_key: str) -> None: + """清掉在途标记(无论成功失败都要清,否则这个作业以后再也不跑了)。""" + with _IN_FLIGHT_LOCK: + _IN_FLIGHT.discard(job_key) def _run_one(db: Session, job: AutomationJob, spec: JobSpec) -> dict[str, Any]: @@ -265,7 +360,8 @@ def _loop() -> None: while not _stop_event.wait(TICK_SECONDS): db = SessionLocal() try: - run_due_jobs(db) + # wait=False:派发完就返回 —— tick 不再被慢作业拖住 + run_due_jobs(db, wait=False) except Exception: # noqa: BLE001 单轮异常不能让调度线程死掉 db.rollback() logger.exception("自动化调度轮询异常") @@ -297,6 +393,12 @@ def stop_scheduler() -> None: if _thread is None: return _stop_event.set() + global _EXECUTOR + with _EXECUTOR_LOCK: + if _EXECUTOR is not None: + # 不 cancel_futures:让已经在跑的作业自然跑完,硬砍会让账号停在半路 + _EXECUTOR.shutdown(wait=False) + _EXECUTOR = None _thread.join(timeout=5) _thread = None