PolyGateway/tests/unit/test_backpressure.py

"""背压 stall 双条件判死 + poll jitter + 记账侧降级(M2 设计 §4/§10)。

保真蓝本 CHS governance.py:200-285: local_waited(本地 monotonic,调用级
累计不重置)与 progress_age(全局活性)**同时**超 stall_window 才判死;
poll 间隔带 [0.5p, 1.0p] jitter 防惊群。
"""

import asyncio

import pytest

from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.errors import (
    AllSourcesExhausted,
    GatewayUnavailableError,
    GovernanceBackendError,
    SourceNotConfiguredError,
    TransientError,
)
from polygateway.middleware.ratelimit import QuotaGate
from polygateway.middleware.retry import RetryMW, backoff_delay
from polygateway.sources import RoundRobinSelector, SourceCooldownMemo
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    ChatRequest,
    GlobalLimits,
    RetryPolicy,
)
from tests.contracts.conftest import FakeClock, make_source
from tests.unit.test_retry import FakeTransport, _ok

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}])
_STALL = 300.0


class BoundedSleep:
    """每次 sleep 执行注入的副作用;超过上限仍在轮询 = 判死逻辑失效,炸出而非死循环。"""

    def __init__(self, side_effect=None, limit: int = 10):
        self.delays: list[float] = []
        self._side_effect = side_effect
        self._limit = limit

    async def __call__(self, seconds: float) -> None:
        self.delays.append(seconds)
        if len(self.delays) > self._limit:
            raise RuntimeError(f"超过 {self._limit} 次轮询仍未判死/未获 permit")
        if self._side_effect is not None:
            await self._side_effect(len(self.delays))


def _mw(
    sources,
    limiter,
    script,
    *,
    clock,
    sleep,
    rng=lambda: 0.0,
    quota_full="wait",
    gate=None,
    transport=None,
    emitter=None,
):
    return RetryMW(
        scope="llm",
        sources=sources,
        selector=RoundRobinSelector(),
        limiter=limiter,
        gate=gate or InMemoryGate(config=_BREAKER, now=clock),
        transport=transport or FakeTransport(script),
        retry=RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0),
        backpressure=BackpressurePolicy(stall_window_s=_STALL, poll_interval_s=0.01),
        quota_full=quota_full,
        cooldown_memo=SourceCooldownMemo(now=clock),
        emitter=emitter,
        now=clock,
        sleep=sleep,
        rng=rng,
    )


def _blocked_limiter(clock):
    """单并发源被外部占满 → RetryMW 进入 wait 轮询分支。"""
    src = make_source(max_concurrency=1)
    limiter = InMemoryLimiter(
        scope="llm",
        sources={"s1": src},
        global_limits=_NO_GLOBAL,
        lease_ttl_s=10_000.0,
        now=clock,
    )
    return src, limiter


class TestStallQuadrants:
    async def test_both_windows_exceeded_raises_stalled(self):
        """双超窗: 本地等待与全局无进展同时 > stall_window → stalled。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        _held = await limiter.try_acquire("s1", 0)

        async def advance(_n):
            clock.advance(_STALL + 100)

        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(advance))
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "stalled"
        assert ei.value.per_source_reasons == {"s1": "rate_limited"}

    async def test_local_exceeded_but_global_fresh_keeps_waiting(self):
        """仅本地超窗: 别人一直在出餐 → 不判死,等到 permit 后正常成功。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        held = await limiter.try_acquire("s1", 0)

        async def advance_and_feed(n):
            clock.advance(_STALL + 100)  # 本地远超窗
            await limiter.mark_progress()  # 但全局刚出过餐
            if n >= 3:
                await held.release()  # 第 3 轮让出 permit

        mw = _mw([src], limiter, [_ok()], clock=clock, sleep=BoundedSleep(advance_and_feed))
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_global_stale_but_local_fresh_keeps_waiting(self):
        """仅全局超窗(从未出餐 age=inf): 本地才刚开始等 → 不判死。

        `inf` 语义在 issue #8 后未变;变的是"本地"的口径——它现在度量的是
        非生产性等待累计,不再是墙钟总耗时(见 TestStallBudget)。
        """
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        held = await limiter.try_acquire("s1", 0)

        async def release_soon(n):
            if n >= 2:  # 本地累计 poll 极短,未超窗
                await held.release()

        mw = _mw([src], limiter, [_ok()], clock=clock, sleep=BoundedSleep(release_soon))
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_poll_jitter_bounds(self):
        """wait 轮询间隔 ∈ [0.5p, 1.0p](CHS governance.py:283-285)。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        held = await limiter.try_acquire("s1", 0)

        async def release_soon(n):
            if n >= 2:
                await held.release()

        for rng_v, expected in ((0.0, 0.005), (1.0, 0.01)):
            clock2 = FakeClock()
            src2, limiter2 = _blocked_limiter(clock2)
            held2 = await limiter2.try_acquire("s1", 0)

            async def release2(n, _h=held2):
                if n >= 2:
                    await _h.release()

            sleep = BoundedSleep(release2)
            mw = _mw([src2], limiter2, [_ok()], clock=clock2, sleep=sleep, rng=lambda v=rng_v: v)
            await mw(_REQ)
            assert sleep.delays[0] == pytest.approx(expected)

    async def test_fail_fast_unaffected(self):
        """fail_fast 路径回归: 不进入 stall 判定,立即 quota_exhausted。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        _held = await limiter.try_acquire("s1", 0)
        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(), quota_full="fail_fast")
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "quota_exhausted"

    async def test_cancellation_pierces_wait_loop(self):
        """stall 等待中的取消穿透(铁律): sleep 可取消,任务立即终止。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        _held = await limiter.try_acquire("s1", 0)
        mw = _mw([src], limiter, [], clock=clock, sleep=asyncio.sleep)
        task = asyncio.create_task(mw(_REQ))
        await asyncio.sleep(0.03)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task


class ClockAdvancingTransport:
    """按脚本 [(推进秒数, 动作), ...] 执行: 在一次尝试内部推进时钟, 模拟真实耗时。

    动作语义同 `FakeTransport`(异常即抛、"hang" 即挂起、其余为返回值)。
    stall 口径的关键区分在于"时间花在哪", 故必须能让时钟只在 transport 内前进。
    """

    def __init__(self, script, clock):
        self.script = list(script)
        self.clock = clock
        self.calls = []

    async def complete(self, *, messages, source, stream, overlay, call_id):
        self.calls.append((source.name, call_id))
        advance, action = self.script.pop(0)
        self.clock.advance(advance)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        return action


class _SlowEmitter:
    """遥测收尾中推进时钟: 钉住"遥测耗时属生产性"(设计 §3.1 边界声明)。"""

    def __init__(self, clock, advance):
        self._clock = clock
        self._advance = advance

    async def emit_attempt(self, *args, **kwargs):
        self._clock.advance(self._advance)


class TestStallBudget:
    """stall 预算只计非生产性等待(issue #8 设计 §3.1)。

    根因是两个预算重叠计费: 真实尝试的耗时同时烧重试预算与 stall 预算,
    而 stall 预算更小必然先耗尽, 于是 max_attempts 在超时场景下永不生效。
    """

    def _free_limiter(self, clock):
        src = make_source()
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        return src, limiter

    async def test_single_timeout_does_not_exhaust_stall_budget(self):
        """timeout_s == stall_window_s 时, 一次超时不得判死——重试预算须真实可用。"""
        clock = FakeClock()
        src, limiter = self._free_limiter(clock)
        # 第一次尝试耗满 300s 超时后失败, 第二次立即成功
        transport = ClockAdvancingTransport(
            [(_STALL + 1, TransientError("timeout", status_code=504)), (0.0, _ok())], clock
        )
        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(), transport=transport)
        resp = await mw(_REQ)
        assert resp.content == "ok"
        assert len(transport.calls) == 2  # 第二次尝试确实发出了

    async def test_productive_time_excluded_from_stall(self):
        """连续多次长尝试也不烧 stall 预算: 它们烧的是重试预算。"""
        clock = FakeClock()
        src, limiter = self._free_limiter(clock)
        transport = ClockAdvancingTransport(
            [
                (_STALL + 100, TransientError("slow", status_code=500)),
                (_STALL + 100, TransientError("slow", status_code=500)),
                (0.0, _ok()),
            ],
            clock,
        )
        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(), transport=transport)
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_telemetry_time_counts_as_productive(self):
        """遥测收尾属 `_attempt` 边界内: 遥测抖动不得参与判死(设计 §3.1)。

        必须走**失败**路径才有判别力: 成功后直接 return, 循环开头的 stall
        判定根本不会再执行。此处让首次尝试快速失败、而遥测收尾慢得超窗,
        下一轮循环开头即检验遥测耗时有没有被算进 stall 账。
        """
        clock = FakeClock()
        src, limiter = self._free_limiter(clock)
        transport = ClockAdvancingTransport(
            [(0.1, TransientError("boom", status_code=500)), (0.0, _ok())], clock
        )
        mw = _mw(
            [src],
            limiter,
            [],
            clock=clock,
            sleep=BoundedSleep(),
            transport=transport,
            emitter=_SlowEmitter(clock, _STALL + 100),
        )
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_nonproductive_wait_still_triggers_stall(self):
        """兜底未被削弱: 纯轮询等待累满窗口仍判死。"""
        clock = FakeClock()
        src, limiter = _blocked_limiter(clock)
        _held = await limiter.try_acquire("s1", 0)

        async def advance(_n):
            clock.advance(_STALL + 100)

        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(advance))
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "stalled"

    async def test_saturation_429_still_stalls(self):
        """429 免预算不烧 fails, 主循环兜底须仍能判死而非无限循环(设计 §3.5)。"""
        clock = FakeClock()
        src, limiter = self._free_limiter(clock)
        # 429 往返本身极快(生产性可忽略), 退避 sleep 才是非生产性的大头
        transport = ClockAdvancingTransport(
            [(0.1, TransientError("429", status_code=429)) for _ in range(10)], clock
        )

        async def advance(_n):
            clock.advance(_STALL)

        mw = _mw(
            [src], limiter, [], clock=clock, sleep=BoundedSleep(advance), transport=transport
        )
        with pytest.raises(AllSourcesExhausted) as ei:
            await mw(_REQ)
        assert ei.value.reason == "stalled"  # 不是 retry_exhausted: 429 确实没烧重试预算

    async def test_cancel_inside_attempt_pierces(self):
        """取消发生在 `attempting()` 包裹内仍逐字穿透(库铁律)。"""
        clock = FakeClock()
        src, limiter = self._free_limiter(clock)
        transport = ClockAdvancingTransport([(0.0, "hang")], clock)
        mw = _mw([src], limiter, [], clock=clock, sleep=asyncio.sleep, transport=transport)
        task = asyncio.create_task(mw(_REQ))
        while not transport.calls:
            await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert (await limiter.source_stats("s1")).inflight == 0  # permit 在 finally 释放

    async def test_clock_is_per_call_not_per_instance(self):
        """StallClock 必须是**调用级**局部状态,不得提升为 RetryMW 实例属性。

        生产形态是一个长寿命 RetryMW 跑成千上万次调用。若 clock 成了实例属性,
        `_entered_at` 会固定在进程启动时刻, 每次调用的 stalled_s() 随进程运行
        时长单调增长, 最终所有调用被误判 stalled——这是本用例要拦的灾难。

        判别力的关键是**复用同一个 mw**: 两个 mw 实例天然隔离, 抓不到实例共享。
        """
        clock = FakeClock()
        src = make_source()
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        transport = ClockAdvancingTransport([(0.0, _ok("first")), (0.0, _ok("second"))], clock)
        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(), transport=transport)
        first = await mw(_REQ)
        clock.advance(_STALL + 100)  # 两次调用之间进程空转远超窗
        second = await mw(_REQ)
        assert (first.content, second.content) == ("first", "second")

    async def test_concurrent_calls_do_not_share_clock(self):
        """并发两路共用同一个 mw: 一路长尝试不得污染另一路的 stall 账。"""
        clock = FakeClock()
        src = make_source(max_concurrency=2)
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        transport = ClockAdvancingTransport(
            [(_STALL + 100, _ok("slow")), (0.0, _ok("fast"))], clock
        )
        mw = _mw([src], limiter, [], clock=clock, sleep=BoundedSleep(), transport=transport)
        results = await asyncio.gather(mw(_REQ), mw(_REQ))
        assert {r.content for r in results} == {"slow", "fast"}


class _GateSuccessBroken(InMemoryGate):
    async def record_success(self, entry):
        raise GovernanceBackendError("redis 抖动", scope="llm")


class _GateFailureBroken(InMemoryGate):
    async def record_failure(self, entry, reason, force_open):
        raise GovernanceBackendError("redis 抖动", scope="llm")


class _LimiterProgressBroken(InMemoryLimiter):
    async def mark_progress(self):
        raise GovernanceBackendError("redis 抖动", scope="llm")


class _GateSuccessMisconfigured(InMemoryGate):
    # 签名须与端口一致(含 count_attempt),否则抛的是 TypeError 而非本类要测的异常
    async def record_success(self, entry, *, count_attempt: bool = True):
        raise SourceNotConfiguredError("未知源 's1'(scope=llm)")


class TestAccountingDegradation:
    """记账侧降级(设计 §10,ARCH §7.3 勘误): 调用已完成,写回失败不冒泡。"""

    async def test_record_success_failure_does_not_lose_response(self):
        clock = FakeClock()
        src = make_source()
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        gate = _GateSuccessBroken(config=_BREAKER, now=clock)
        mw = _mw([src], limiter, [_ok()], clock=clock, sleep=BoundedSleep(), gate=gate)
        resp = await mw(_REQ)
        assert resp.content == "ok"  # 真实成功响应不因记账失败被丢弃

    async def test_assembly_defect_on_accounting_path_also_degrades(self):
        """记账侧降级按"路径性质"而非异常类型: 装配缺陷同样不得毁掉已完成的调用。

        `SourceNotConfiguredError` 被放行穿透闸门包装器(issue #7 §T6)后,若
        `_record_quietly` 只降级 `GovernanceBackendError`,它就会从记账侧冒泡、
        销毁一个真实成功的响应——反转本类钉住的既有行为。当前无后端会从记账
        方法抛它,此用例是为将来加了源名校验的后端守住这条不变式。
        """
        clock = FakeClock()
        src = make_source()
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        gate = _GateSuccessMisconfigured(config=_BREAKER, now=clock)
        mw = _mw([src], limiter, [_ok()], clock=clock, sleep=BoundedSleep(), gate=gate)
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_mark_progress_failure_does_not_lose_response(self):
        clock = FakeClock()
        src = make_source()
        limiter = _LimiterProgressBroken(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        mw = _mw([src], limiter, [_ok()], clock=clock, sleep=BoundedSleep())
        resp = await mw(_REQ)
        assert resp.content == "ok"

    async def test_record_failure_failure_does_not_mask_retry(self):
        """失败记账挂掉 → 原始尝试异常不被掩盖,重试照常换发并成功。"""
        clock = FakeClock()
        src = make_source()
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": src}, global_limits=_NO_GLOBAL, now=clock
        )
        gate = _GateFailureBroken(config=_BREAKER, now=clock)
        script = [TransientError("boom", status_code=500), _ok()]
        mw = _mw([src], limiter, script, clock=clock, sleep=BoundedSleep(), gate=gate)
        resp = await mw(_REQ)
        assert resp.content == "ok"


class TestBackoffExtraction:
    """T5 提取的模块级纯函数(T9 Embedding 复用)。"""

    def test_formula_matches_policy(self):
        policy = RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0)
        exc = TransientError("x", status_code=500)
        assert backoff_delay(policy, 1, exc, lambda: 0.0) == pytest.approx(1.0)  # 2*0.5
        assert backoff_delay(policy, 1, exc, lambda: 1.0) == pytest.approx(3.0)  # 2*1.5
        assert backoff_delay(policy, 10, exc, lambda: 1.0) == pytest.approx(45.0)  # 封顶 30*1.5

    def test_retry_after_hint_wins_when_larger(self):
        policy = RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0)
        exc = TransientError("x", status_code=429, retry_after_s=7.5)
        assert backoff_delay(policy, 1, exc, lambda: 0.0) == pytest.approx(7.5)


class TestQuotaGateProgressAge:
    async def test_passthrough_and_wrap(self):
        from polygateway.middleware.ratelimit import QuotaGate

        class _L:
            async def progress_age_s(self):
                return 12.5

        class _Broken:
            async def progress_age_s(self):
                raise OSError("down")

        assert await QuotaGate(_L(), scope="llm").progress_age_s() == 12.5
        with pytest.raises(GovernanceBackendError):
            await QuotaGate(_Broken(), scope="llm").progress_age_s()


class TestUnknownSourceIsAssemblyDefect:
    """未知源 = 限流后端的源名单与治理循环对不上,是装配缺陷不是后端故障。

    两个后端行为必须一致(Redis 版对应用例在 `test_redis_key_layout.py::
    TestConversions::test_unknown_source_rejected`);内存版此前无覆盖,
    该分支从未被测过(issue #7 §3.4)。
    """

    def test_memory_limiter_rejects_unknown_source(self):
        limiter = InMemoryLimiter(
            scope="llm", sources={"s1": make_source("s1")}, global_limits=_NO_GLOBAL
        )
        with pytest.raises(SourceNotConfiguredError) as ei:
            limiter._cfg("nope")
        # 关键: 若归入 scope 级家族,配置写错的任务会永远延期重投、永不进死信
        assert not isinstance(ei.value, GatewayUnavailableError)

    @pytest.mark.parametrize("method", ["try_acquire", "stats"])
    async def test_survives_the_quota_gate_wrapper(self, method):
        """必须穿透 QuotaGate,否则整个拆分在生产路径上等于没做。

        上面两条(以及 redis 版)打的都是私有 `_cfg`,绕过了包装器。而治理循环
        只经 QuotaGate 访问后端,包装器的 `except Exception` 会把装配缺陷重新
        包成 `GovernanceBackendError`——下游又拿到可重投异常,永远重投不告警。
        """
        src = make_source("s1")
        # 限流后端的源名单与治理循环拿到的源对不上 = 装配缺陷
        limiter = InMemoryLimiter(
            scope="llm", sources={"other": src}, global_limits=_NO_GLOBAL
        )
        gate = QuotaGate(limiter, scope="llm")
        with pytest.raises(SourceNotConfiguredError) as ei:
            await getattr(gate, method)(src)
        assert not isinstance(ei.value, GatewayUnavailableError)


class TestGateFailuresReachCallersAsScopeLevel:
    """闸门泄漏路径必须以 scope 级不可用的形态到达调用方(issue #7)。

    记账路径由 `_record_quietly` 降级为 warning,但闸门路径没有那层包裹,会一路
    抛给调用方。只写 `except GatewayUnavailableError` 的调用方此前接不住,后果
    是 Redis 抖一下就让积压任务烧掉业务失败预算进死信——而那是运维重启即可恢复
    的故障。全部五条为: `QuotaGate` 的 try_acquire / stats / progress_age_s,
    `BreakerGate` 的 try_enter / retry_after_s(判据是该调用点未被 `_record_quietly`
    包裹)。此处钉住其中三条代表路径,余两条由同一注入机制覆盖。
    """

    async def test_try_acquire_failure_is_scope_level(self):
        from polygateway.middleware.ratelimit import QuotaGate

        class _Broken:
            async def try_acquire(self, name, est):
                raise OSError("down")

        with pytest.raises(GatewayUnavailableError) as ei:
            await QuotaGate(_Broken(), scope="LLM").try_acquire(make_source("s1"))
        assert ei.value.scope == "llm"
        assert ei.value.reason == "governance_backend_down"
        assert ei.value.retry_after_s > 0  # 0 会让积压任务零延迟冲击已挂的后端

    async def test_try_enter_failure_is_scope_level(self):
        from polygateway.middleware.breaker import BreakerGate

        class _Broken:
            async def try_enter(self, name, owner):
                raise OSError("down")

        with pytest.raises(GatewayUnavailableError) as ei:
            await BreakerGate(_Broken(), scope="LLM").try_enter(make_source("s1"), "owner")
        assert ei.value.scope == "llm"
        assert ei.value.reason == "governance_backend_down"

    async def test_progress_age_failure_is_scope_level(self):
        from polygateway.middleware.ratelimit import QuotaGate

        class _Broken:
            async def progress_age_s(self):
                raise OSError("down")

        with pytest.raises(GatewayUnavailableError) as ei:
            await QuotaGate(_Broken(), scope="LLM").progress_age_s()
        assert ei.value.scope == "llm"
        assert ei.value.reason == "governance_backend_down"