chongming-utils-task/tests/unit/test_thread_executor.py

"""线程执行器测试。

这些用例全部使用**真实线程与真实时间**：执行器的全部价值就在于线程能不能被
停下、能不能被暂停、放弃之后进程还能不能退出。用假的线程模型测不出这些。

为避免偶发失败，协作路径给足超时（5 秒），而"不配合"的用例刻意把任务做成
有限时长（约 5 秒后自行结束），这样即使被放弃，残留的守护线程也不会长期占用资源。
"""

from __future__ import annotations

import asyncio
import threading
import time
from typing import Any

import pytest

from chongming_task.models import FailureKind
from chongming_task.runtime.executors.base import ExecutionRequest
from chongming_task.runtime.executors.thread import ThreadExecutor
from chongming_task.testing import FakeClock

_COOPERATIVE_WINDOW = 5.0
_ABANDON_WINDOW = 0.2


def _request(task_id: str, target: Any, *args: Any) -> ExecutionRequest:
    return ExecutionRequest(task_id=task_id, name=task_id, target=target, args=args)


# ---------------------------------------------------------------------------
# 上下文注入契约
# ---------------------------------------------------------------------------


def test_context_is_injected_with_the_handle_token_and_metadata() -> None:
    """令牌必须是句柄那一个，否则控制面会停错对象。

    这也是"令牌由执行器注入、而非被任务闭包捕获"的直接体现——令牌每次尝试才诞生，
    闭包捕获到的只可能是上一次尝试留下的那一个。
    """
    seen: dict[str, Any] = {}

    def work(ctx: Any) -> str:
        seen["token"] = ctx.token
        seen["task_id"] = ctx.task_id
        seen["name"] = ctx.name
        seen["attempt"] = ctx.attempt
        return "ok"

    handle = ThreadExecutor().launch(
        ExecutionRequest(task_id="t-ctx", name="抽取", target=work, attempt=2)
    )
    assert handle.wait(timeout=_COOPERATIVE_WINDOW) is not None
    assert seen["token"] is handle.token
    assert seen["task_id"] == "t-ctx"
    assert seen["name"] == "抽取"
    assert seen["attempt"] == 2


def test_arguments_are_forwarded_after_the_context() -> None:
    def work(ctx: Any, left: int, right: int, *, scale: int = 1) -> int:
        del ctx
        return (left + right) * scale

    handle = ThreadExecutor().launch(
        ExecutionRequest(
            task_id="t-args",
            name="t-args",
            target=work,
            args=(2, 3),
            kwargs={"scale": 10},
        )
    )
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert outcome.value == 50


def test_checkpoint_refreshes_the_heartbeat() -> None:
    """checkpoint 就是"我还在推进"的信号，心跳绑在它上面是有意的。"""
    clock = FakeClock()
    observed: list[float] = []

    def work(ctx: Any) -> str:
        for _ in range(3):
            clock.advance(30)
            observed.append(ctx.heartbeat.idle_seconds())
            ctx.checkpoint()
        return "ok"

    handle = ThreadExecutor(clock=clock).launch(_request("t-beat", work))
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert outcome.ok
    assert observed == [30.0, 30.0, 30.0], "每次 checkpoint 都应把静默时长清零"
    assert handle.heartbeat.idle_seconds() == 0.0


# ---------------------------------------------------------------------------
# async 任务
# ---------------------------------------------------------------------------


def test_async_target_is_actually_awaited() -> None:
    """``async def`` 任务必须真的被执行完，而不是把协程对象当返回值。

    后者曾经是静默的：没有存储时它看起来"成功"了，返回值却是个 ``coroutine``；
    有存储时才会以"返回值不可序列化"失败。
    """

    async def work(ctx: Any) -> str:
        await asyncio.sleep(0)
        return "done"

    handle = ThreadExecutor().launch(_request("t-async", work))
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert outcome.ok
    assert outcome.value == "done"


def test_async_target_cooperates_with_stop() -> None:
    """协作式停止对 async 任务同样有效：``await ctx.acheckpoint()``。"""

    async def work(ctx: Any) -> str:
        for _ in range(10_000):
            await ctx.acheckpoint()
            await asyncio.sleep(0.005)
        return "never"

    handle = ThreadExecutor().launch(_request("t-async-stop", work))
    outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
    assert not outcome.ok, "停止请求之后不该报成功"
    assert not outcome.abandoned, "它会到达检查点，应当优雅收尾而不是被放弃"
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.CANCELLED


def test_async_target_failure_is_reported_across_the_loop() -> None:
    """协程里抛的异常照常转成结构化失败，不能被事件循环吞掉。"""

    async def work(ctx: Any) -> str:
        await asyncio.sleep(0)
        raise ValueError("协程里炸了")

    handle = ThreadExecutor().launch(_request("t-async-boom", work))
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert not outcome.ok
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.EXCEPTION
    assert "协程里炸了" in outcome.failure.message


# ---------------------------------------------------------------------------
# 正常路径
# ---------------------------------------------------------------------------


def test_launch_returns_immediately_without_waiting_for_the_task() -> None:
    """调度器契约：launch 必须立刻返回，否则调度循环会被单个任务卡住。"""
    release = threading.Event()

    def work(ctx: Any) -> str:
        del ctx
        release.wait(5)
        return "done"

    handle = ThreadExecutor().launch(_request("t-async-launch", work))
    try:
        assert handle.running is True
        assert handle.wait(timeout=0.05) is None, "任务尚未结束，wait 应返回 None"
    finally:
        release.set()
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert outcome.ok
    assert outcome.value == "done"


def test_successful_task_returns_its_value() -> None:
    handle = ThreadExecutor().launch(_request("t-ok", lambda ctx: 6 * 7))
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert outcome.ok
    assert outcome.value == 42
    assert handle.running is False


def test_failing_task_reports_exception_with_traceback() -> None:
    def boom(ctx: Any) -> None:
        del ctx
        raise ValueError("坏掉了")

    handle = ThreadExecutor().launch(_request("t-boom", boom))
    outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
    assert outcome is not None
    assert not outcome.ok
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.EXCEPTION
    assert "ValueError" in outcome.failure.message
    assert outcome.failure.traceback is not None
    assert "boom" in outcome.failure.traceback
    assert outcome.failure.retryable is True


# ---------------------------------------------------------------------------
# 协作式停止
# ---------------------------------------------------------------------------


def test_stop_completes_cooperatively_when_task_checks_the_token() -> None:
    started = threading.Event()
    steps: list[int] = []

    def work(ctx: Any) -> str:
        started.set()
        for index in range(100_000):
            ctx.checkpoint()
            steps.append(index)
            time.sleep(0.001)
        return "ran to completion"

    handle = ThreadExecutor().launch(_request("t-coop", work))
    assert started.wait(_COOPERATIVE_WINDOW)

    outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW, reason="发布窗口")
    assert outcome.abandoned is False
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.CANCELLED
    assert outcome.failure.retryable is False, "取消不是失败，绝不能被重试"
    assert handle.running is False
    assert len(steps) < 100_000, "任务应当在收到停止请求后提前退出"


def test_cancellation_is_not_swallowed_by_a_broad_except_in_the_task_body() -> None:
    """这正是 ``TaskCancelled`` 继承 ``BaseException`` 的理由。

    若它继承 ``Exception``，下面这种极常见的写法会把停止信号吃掉，任务永远停不掉。
    """
    started = threading.Event()
    swallowed: list[BaseException] = []

    def work(ctx: Any) -> str:
        started.set()
        try:
            while True:
                ctx.checkpoint()
                time.sleep(0.001)
        except Exception as exc:  # 用户代码里最常见的写法
            swallowed.append(exc)
            return "signal was swallowed"

    handle = ThreadExecutor().launch(_request("t-swallow", work))
    assert started.wait(_COOPERATIVE_WINDOW)

    outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
    assert swallowed == [], "取消信号被 except Exception 吞掉了"
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.CANCELLED


def test_stop_is_idempotent() -> None:
    def work(ctx: Any) -> str:
        del ctx
        return "quick"

    handle = ThreadExecutor().launch(_request("t-idem", work))
    first = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
    second = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
    assert first.ok
    assert second.ok
    assert handle.running is False


# ---------------------------------------------------------------------------
# 协作式暂停（D6 定案的语义，端到端验证）
# ---------------------------------------------------------------------------


def test_pause_suspends_a_running_task_and_resume_continues_it() -> None:
    started = threading.Event()
    progress: list[int] = []

    def work(ctx: Any) -> str:
        started.set()
        for index in range(100_000):
            ctx.checkpoint()
            progress.append(index)
            time.sleep(0.001)
        return "ran to completion"

    handle = ThreadExecutor().launch(_request("t-pause", work))
    assert started.wait(_COOPERATIVE_WINDOW)

    assert handle.pause(timeout=_COOPERATIVE_WINDOW) is True
    frozen = len(progress)
    assert frozen > 0, "暂停前任务应当已经推进过"
    time.sleep(0.05)
    assert len(progress) == frozen, "暂停生效后任务不应继续推进"

    handle.resume()
    time.sleep(0.05)
    assert len(progress) > frozen, "恢复后任务应继续推进"

    outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.CANCELLED


def test_pause_reports_false_when_task_never_reaches_a_checkpoint() -> None:
    """这正是"暂停没生效必须回滚为 RUNNING"的依据：如实返回 False。"""
    started = threading.Event()

    def work(ctx: Any) -> str:
        del ctx
        started.set()
        for _ in range(500):
            time.sleep(0.01)
        return "ignored"

    handle = ThreadExecutor().launch(_request("t-nopause", work))
    assert started.wait(_COOPERATIVE_WINDOW)
    assert handle.pause(timeout=0.2) is False
    handle.stop(graceful_timeout=_ABANDON_WINDOW)


def test_stop_wakes_a_paused_task() -> None:
    """暂停中的任务必须能被停止，否则暂停就成了不可撤销的挂起。"""
    started = threading.Event()

    def work(ctx: Any) -> str:
        started.set()
        for _ in range(100_000):
            ctx.checkpoint()
            time.sleep(0.001)
        return "ran to completion"

    handle = ThreadExecutor().launch(_request("t-pause-stop", work))
    assert started.wait(_COOPERATIVE_WINDOW)
    assert handle.pause(timeout=_COOPERATIVE_WINDOW) is True

    outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW, reason="运维介入")
    assert outcome.abandoned is False
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.CANCELLED
    assert "运维介入" in outcome.failure.message


# ---------------------------------------------------------------------------
# 无法停止：放弃
# ---------------------------------------------------------------------------


def _uncooperative(ctx: Any) -> str:
    """刻意不检查令牌；约 5 秒后自行结束，避免残留线程长期占用资源。"""
    del ctx
    for _ in range(500):
        time.sleep(0.01)
    return "ignored the stop request"


def test_uncooperative_task_is_abandoned_and_reported_as_lost() -> None:
    """线程无法强杀：必须如实上报，不能假装已停止。"""
    handle = ThreadExecutor().launch(_request("t-stubborn", _uncooperative))
    outcome = handle.stop(graceful_timeout=_ABANDON_WINDOW)
    assert outcome.abandoned is True
    assert outcome.ok is False
    assert outcome.failure is not None
    assert outcome.failure.kind is FailureKind.LOST
    assert "无法强制终止" in outcome.failure.message


def test_abandoned_task_still_reports_running() -> None:
    """它确实还在跑。谎报为 False 会让上层以为资源已经释放。"""
    handle = ThreadExecutor().launch(_request("t-still-running", _uncooperative))
    handle.stop(graceful_timeout=_ABANDON_WINDOW)
    assert handle.running is True


def test_abandoned_outcome_is_sticky() -> None:
    """已记为 LOST 的结论不应被后台线程后来的返回悄悄改写，否则审计自相矛盾。"""
    handle = ThreadExecutor().launch(_request("t-sticky", _uncooperative))
    first = handle.stop(graceful_timeout=_ABANDON_WINDOW)
    assert handle.wait(timeout=_ABANDON_WINDOW) is first
    assert handle.stop(graceful_timeout=_ABANDON_WINDOW) is first


def test_worker_threads_are_daemon_so_an_abandoned_task_cannot_block_shutdown() -> None:
    """``ThreadPoolExecutor`` 的工作线程会在解释器退出时被 join，一个不响应取消
    的任务会让整个进程无法退出。这里用守护线程正是为了避开这一点。"""
    handle = ThreadExecutor().launch(_request("t-daemon", _uncooperative))
    try:
        assert handle.thread.daemon is True
        assert handle.thread.name == "task-t-daemon"
    finally:
        handle.stop(graceful_timeout=_ABANDON_WINDOW)


@pytest.mark.parametrize("graceful_timeout", [0.0, 0.05])
def test_zero_or_tiny_graceful_window_still_yields_a_conclusion(
    graceful_timeout: float,
) -> None:
    """优雅窗口再短也必须给出结论，不能抛异常或永远阻塞。"""
    handle = ThreadExecutor().launch(_request(f"t-window-{graceful_timeout}", _uncooperative))
    outcome = handle.stop(graceful_timeout=graceful_timeout)
    assert outcome.abandoned is True
    assert outcome.failure is not None