From 1f040afabf05a20afc2ec0172f1693da8d7947c4 Mon Sep 17 00:00:00 2001 From: Justin Gu <97915@qq.com> Date: Mon, 13 Jul 2026 15:46:02 +0800 Subject: [PATCH] =?UTF-8?q?feat(reconnect):=20=E5=BC=95=E5=85=A5=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E5=99=A8=E5=81=A5=E5=BA=B7=E5=88=86=E5=BC=95=E6=93=8E?= =?UTF-8?q?=20+=20K=E7=BA=BF=E7=A9=BA=E6=95=B0=E6=8D=AE=E6=95=85=E9=9A=9C?= =?UTF-8?q?=E8=BD=AC=E7=A7=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 彻底解决通达信服务器"跳来跳去"问题： 1. 新增 _health.py 健康分引擎：失败乘性降权(×0.5)、连续失败3次进 120s 冷却、成功加性恢复(+0.2)。rank_by_health 按 latency/score (有效延迟)重排，冷却中的剔除。全健康时恒等映射，对既有测试零影响。 2. get_index_bars/get_security_bars 空数据时自动逐台换台（此前直接 返回空 DataFrame，是日志"指数K线响应在第1/800条处被截断"后用户拿 不到数据的根因）。复用泛化后的 _find_host_returning_data。 3. select_best_host_*/find_working_host_* 应用 rank_by_health 重排； 空数据验证失败/异常时调 record_failure，命中调 record_success。 4. 8 个 _execute(A股/MAC/EX/MAC-EX × sync/async)统一注入健康分记录： 成功 record_success、连接失败 record_failure。 5. security_bars 截断日志区分"首条即空(服务器无数据)"与"末尾截断"。 测试：26 个新增（15 health + 7 failover + 4 ex-client 健康分追踪）， 全量 reconnect/failover/decode 回归通过，ruff/mypy 通过。 --- src/easy_tdx/_health.py | 163 ++++++++++++++++++++ src/easy_tdx/_reconnect.py | 18 +++ src/easy_tdx/client.py | 197 +++++++++++++------------ src/easy_tdx/commands/security_bars.py | 47 ++++-- src/easy_tdx/ex/client.py | 40 ++++- src/easy_tdx/ex/mac_client.py | 37 ++++- src/easy_tdx/mac/client.py | 42 +++++- tests/unit/test_ex_reconnect.py | 125 ++++++++++++++++ tests/unit/test_failover.py | 91 ++++++++++-- tests/unit/test_health.py | 149 +++++++++++++++++++ 10 files changed, 770 insertions(+), 139 deletions(-) create mode 100644 src/easy_tdx/_health.py create mode 100644 tests/unit/test_health.py diff --git a/src/easy_tdx/_health.py b/src/easy_tdx/_health.py new file mode 100644 index 0000000..0d5f2fa --- /dev/null +++ b/src/easy_tdx/_health.py @@ -0,0 +1,163 @@ +"""服务器健康分（health score）引擎。 + +为通达信候选主机维护一份**进程级**的健康记录：每次失败降权、连续失败 +触发冷却、成功缓慢恢复。``select_best_host_*`` / ``find_working_host_*`` +据此重排候选列表，让"延迟低但数据不全 / 频繁断连"的服务器自动靠后， +避免被低延迟反复选中又反复触发空数据故障转移（日志里"服务器跳来跳去" +的根因之一）。 + +设计要点： + 1. 模块级单例 + ``threading.Lock``。8 个 client（A股/MAC/EX/MAC-EX × + sync/async）共享同一份记录——一台服务器的好坏不分协议。 + 2. score ∈ ``(0, 1.0]``，初始 1.0。``record_failure`` 乘性衰减（×0.5）， + 连续失败 ≥ ``_COOLDOWN_FAIL_THRESHOLD`` 台进入 ``_COOLDOWN_SEC`` 秒 + 冷却期；``record_success`` 加性恢复（+0.2，上限 1.0）并重置计数。 + 3. ``rank_by_health`` 把 ping 结果按 ``latency / score`` 重排（score 越低 + 惩罚越大），冷却中的主机直接剔除——既非永久黑名单（恢复后可回归）， + 也非纯延迟（数据不全的低延迟服务器会被压下去）。 + 4. 全健康时 ``rank_by_health`` 近似恒等映射，对既有 failover 单测零影响。 +""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass, field + +# --------------------------------------------------------------------------- +# 可调常数 +# --------------------------------------------------------------------------- + +# 失败一次的乘性衰减因子。score *= _FAILURE_DECAY（0.5 → 失败 3 次后 score ≈ 0.125）。 +_FAILURE_DECAY: float = 0.5 + +# 成功一次的加性恢复量。score += _SUCCESS_RECOVER（上限 1.0）。 +_SUCCESS_RECOVER: float = 0.2 + +# 进入冷却所需的连续失败次数。达到即认为该主机"持续不可用"，冷却期内剔除。 +_COOLDOWN_FAIL_THRESHOLD: int = 3 + +# 冷却时长（秒）。冷却期内 ``is_in_cooldown`` 返回 True，``rank_by_health`` 剔除该主机。 +_COOLDOWN_SEC: float = 120.0 + +# score 下限：保留一个极小正值而非归零，确保恢复路径可达且不会除零。 +_SCORE_FLOOR: float = 1e-3 + + +# --------------------------------------------------------------------------- +# 状态 +# --------------------------------------------------------------------------- + + +@dataclass +class _HostHealth: + """单台主机的健康记录。""" + + score: float = 1.0 + consecutive_failures: int = 0 + cooldown_until: float = 0.0 # monotonic 时间戳；0 表示未进入冷却 + + +@dataclass +class _HealthBook: + """所有主机的健康记录簿（模块单例）。""" + + hosts: dict[str, _HostHealth] = field(default_factory=dict) + lock: threading.Lock = field(default_factory=threading.Lock) + + def _get(self, host: str) -> _HostHealth: + hh = self.hosts.get(host) + if hh is None: + hh = _HostHealth() + self.hosts[host] = hh + return hh + + +_BOOK = _HealthBook() + + +# --------------------------------------------------------------------------- +# 写入：失败 / 成功 / 重置 +# --------------------------------------------------------------------------- + + +def record_failure(host: str) -> float: + """记录一次失败：score 乘性衰减，连续失败达阈值则进入冷却。 + + Returns: + 衰减后的 score（便于调用方记录日志）。 + """ + now = time.monotonic() + with _BOOK.lock: + hh = _BOOK._get(host) + hh.score = max(_SCORE_FLOOR, hh.score * _FAILURE_DECAY) + hh.consecutive_failures += 1 + if hh.consecutive_failures >= _COOLDOWN_FAIL_THRESHOLD: + hh.cooldown_until = now + _COOLDOWN_SEC + return hh.score + + +def record_success(host: str) -> None: + """记录一次成功：score 加性恢复（上限 1.0），重置连续失败计数与冷却。""" + with _BOOK.lock: + hh = _BOOK._get(host) + hh.score = min(1.0, hh.score + _SUCCESS_RECOVER) + hh.consecutive_failures = 0 + hh.cooldown_until = 0.0 + + +def reset_health() -> None: + """清空全部健康记录。主要供测试隔离使用。""" + with _BOOK.lock: + _BOOK.hosts.clear() + + +# --------------------------------------------------------------------------- +# 读取：冷却判定 / 健康分 / 排序 +# --------------------------------------------------------------------------- + + +def is_in_cooldown(host: str) -> bool: + """该主机是否处于冷却期（True=应剔除，暂不选用）。""" + now = time.monotonic() + with _BOOK.lock: + hh = _BOOK.hosts.get(host) + if hh is None: + return False + return hh.cooldown_until > now + + +def get_score(host: str) -> float: + """返回该主机当前 score（无记录则 1.0）。""" + with _BOOK.lock: + hh = _BOOK.hosts.get(host) + return hh.score if hh is not None else 1.0 + + +def rank_by_health( + ranked_hosts: list[tuple[str, float]], +) -> list[tuple[str, float]]: + """按健康分重排 ping 结果。 + + 输入是 ``ping_all`` 返回的 ``[(host, latency), ...]``（已按延迟升序）。 + 本函数： + 1. 剔除处于冷却期的主机； + 2. 对剩余主机按 ``latency / score``（有效延迟）升序排序——score 越低 + 惩罚越大，但低延迟仍占优，两者平滑权衡。 + + 全健康时（所有 score=1.0、无冷却），输出与输入排序一致（恒等映射）， + 故对既有 ``test_failover.py`` 的 mock 测试零影响。 + """ + now = time.monotonic() + with _BOOK.lock: + kept: list[tuple[str, float, float]] = [] + for host, latency in ranked_hosts: + hh = _BOOK.hosts.get(host) + if hh is not None and hh.cooldown_until > now: + # 冷却中：剔除 + continue + score = hh.score if hh is not None else 1.0 + kept.append((host, latency, score)) + # 有效延迟 = latency / score；score 越小有效延迟越大，越靠后 + kept.sort(key=lambda t: t[1] / t[2]) + return [(h, lat) for h, lat, _ in kept] diff --git a/src/easy_tdx/_reconnect.py b/src/easy_tdx/_reconnect.py index 730ba16..910cec8 100644 --- a/src/easy_tdx/_reconnect.py +++ b/src/easy_tdx/_reconnect.py @@ -24,6 +24,7 @@ import logging import time from collections.abc import Awaitable, Callable +from ._health import rank_by_health, record_failure, record_success from .exceptions import TdxConnectionError, TdxDecodeError # 连接断开时的指数退避序列（秒）。每次重连失败后按此序列 sleep 再重试， @@ -116,6 +117,9 @@ def select_best_host_sync( # 无论测速是否拿到结果，都视为"完成一次测速"，开启节流窗口， # 避免失败时被高频重试反复触发。 _mark_failover_done() + # 健康分重排：冷却中的剔除，剩余按 latency/score（有效延迟）排序。 + # 全健康时近似恒等映射，对既有 mock 测试零影响。 + ranked = rank_by_health(ranked) # 跳过当前（已判定不可用）主机，取延迟最低的另一台 for host, _latency in ranked: if host != current_host: @@ -149,6 +153,8 @@ async def select_best_host_async( ranked = await asyncio.to_thread(ping_fn, hosts, port, ping_timeout) finally: _mark_failover_done() + # 健康分重排（与 sync 版对称）：冷却剔除 + 有效延迟排序。 + ranked = rank_by_health(ranked) for host, _latency in ranked: if host != current_host: save_fn(host) @@ -195,6 +201,8 @@ def find_working_host_sync( """ log = logging.getLogger(__name__) tried = 0 + # 健康分预过滤：剔除冷却中的主机，避免对"持续不可用"的服务器反复实测。 + ranked_hosts = rank_by_health(ranked_hosts) for host, _latency in ranked_hosts: if host == current_host: continue @@ -204,6 +212,7 @@ def find_working_host_sync( try: if try_fn(host): save_fn(host) + record_success(host) log.info( "空数据故障转移：从 %s 切换到 %s（第 %d 台候选可用）", current_host, @@ -211,9 +220,12 @@ def find_working_host_sync( tried, ) return host + # 连通但返回空数据：记一次失败降权，下次优先级下降 + record_failure(host) except Exception: # 验证单台主机时的任何异常（连接失败、解析错误等）都只跳过该台， # 继续尝试下一台，不让单台拖垮整个轮询。 + record_failure(host) log.debug("空数据故障转移：%s 验证失败，尝试下一台", host, exc_info=True) return None @@ -231,6 +243,8 @@ async def find_working_host_async( """ log = logging.getLogger(__name__) tried = 0 + # 健康分预过滤（与 sync 版对称）：剔除冷却中的主机。 + ranked_hosts = rank_by_health(ranked_hosts) for host, _latency in ranked_hosts: if host == current_host: continue @@ -240,6 +254,7 @@ async def find_working_host_async( try: if await try_fn(host): save_fn(host) + record_success(host) log.info( "空数据故障转移：从 %s 切换到 %s（第 %d 台候选可用）", current_host, @@ -247,7 +262,10 @@ async def find_working_host_async( tried, ) return host + # 连通但返回空数据：记一次失败降权 + record_failure(host) except Exception: + record_failure(host) log.debug("空数据故障转移：%s 验证失败，尝试下一台", host, exc_info=True) return None diff --git a/src/easy_tdx/client.py b/src/easy_tdx/client.py index dbcb8e6..9263e6c 100644 --- a/src/easy_tdx/client.py +++ b/src/easy_tdx/client.py @@ -22,6 +22,7 @@ from ._df import ( _merge_txn_datetime, _to_df, ) +from ._health import record_failure, record_success from ._reconnect import ( _RETRY_DELAYS, AsyncHeartbeatMixin, @@ -344,20 +345,29 @@ class TdxClient: 2. 跨主机故障转移——同主机重试仍失败时，重新测速选延迟最低的另 一台服务器再试一轮。服务器连不上时用户无需手动 ``ping``。 ``auto_reconnect=False`` 时两阶段都不触发，直接抛出原异常。 + + 健康分联动：成功路径记 ``record_success``（恢复 score），连接失败 + 记 ``record_failure``（降权），让频繁断连的服务器在后续故障转移中 + 自动靠后。 """ try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + # 连接失败：当前主机降权 + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: time.sleep(delay) self._reconnect() try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移——重新测速切到另一台服务器再试一次 new_host = select_best_host_sync( get_known_hosts(), @@ -370,10 +380,16 @@ class TdxClient: if new_host is not None: self._reconnect(new_host) try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 市场信息 @@ -508,10 +524,12 @@ class TdxClient: """ cmd = GetSecurityBarsCmd(market, code, category, start, count) bars = self._execute(cmd) - # 空数据故障转移：部分服务器对所有证券返回空 body（不报错），延迟最低的不 - # 一定有数据，故空时按延迟逐台实测找首台返回数据的 host。 + # 空数据故障转移：服务器连通但返回空/截断（部分服务器对所有证券返回空 body + # 且不报错），按延迟顺序逐台实测找首台有效数据的服务器。与 get_market_stat + # 同源逻辑。注意：真·无历史数据（如新股）所有服务器都返回空，此时换台仍为空， + # 直接返回空 DataFrame 而非 raise——避免把"该股票本就没数据"误报为故障。 if not bars and self._auto_reconnect: - bars = self._find_host_returning_bars(cmd) + bars = self._find_host_returning_data(cmd) df = _to_df(bars) delta = _category_to_minutes(int(category)) is_intraday = delta is not None @@ -541,10 +559,11 @@ class TdxClient: """ cmd = GetIndexBarsCmd(market, code, category, start, count) bars = self._execute(cmd) - # 空数据故障转移：指数/板块指数（880xxx 等）并非所有服务器都提供，延迟最低 - # 的不一定返回数据，故空时按延迟逐台实测找首台返回数据的 host。 + # 空数据故障转移：指数/板块指数（880xxx 等）并非所有服务器都提供，服务端 + # 截断返回 0 条是已知现象（日志"指数K线响应在第1/800条处被截断"）。 + # 按延迟顺序逐台实测换台，避免上层拿到空数据。全失败则返回空 DataFrame。 if not bars and self._auto_reconnect: - bars = self._find_host_returning_bars(cmd) + bars = self._find_host_returning_data(cmd) df = _to_df(bars) delta = _category_to_minutes(int(category)) is_intraday = delta is not None @@ -769,56 +788,44 @@ class TdxClient: ) ) + def _find_host_returning_data(self, cmd: "BaseCommand[_T]") -> _T: + """空数据故障转移：测速后按延迟顺序逐台实测，返回首台有效数据的结果。 + + 泛化版：支持任意返回 ``list``/序列的命令（quotes、K 线 bars 等）。 + 统计指数（880005 等）、指数 K 线等并非所有服务器都提供，延迟最低的不 + 一定返回数据，故需逐台实际查询。最多尝试 ``_WORKING_HOST_MAX_ATTEMPTS`` + 台（见 ``_reconnect``）。找到后 client 停在该 host；全失败返回空值 + （由 ``bool()`` 判空）。 + + Note: + 命令返回值必须可被 ``bool()`` 判空（list / DataFrame 均满足）。 + """ + bad_host = self._host + ranked = ping_all(get_known_hosts(), self._port, 5.0) + + def _try(host: str) -> bool: + # 切换到候选 host 并实测；非空即视为该 host 可用 + self._reconnect(host) + return bool(self._execute(cmd)) + + new_host = find_working_host_sync(ranked, _try, save_best_host, bad_host) + if new_host is None: + # 全部候选都不可用，回退到原 host（保持状态可预测） + if self._host != bad_host: + self._reconnect(bad_host) + return [] # type: ignore[return-value] + # _try 已把 client 切到 new_host 并执行过 cmd，重新取一次拿结果 + return self._execute(cmd) + def _find_host_returning_quotes( self, cmd: "BaseCommand[list[SecurityQuote]]" ) -> list[SecurityQuote]: - """空数据故障转移：测速后按延迟顺序逐台实测，返回首台有效数据的 quotes。 + """空数据故障转移（quotes 专用薄封装）。 - 专供 ``get_market_stat`` 使用——统计指数并非所有服务器都提供，延迟最低 - 的不一定返回数据，故需逐台实际查询。最多尝试 ``_WORKING_HOST_MAX_ATTEMPTS`` - 台（见 ``_reconnect``）。找到后 client 停在该 host；全失败返回空。 + 保留独立方法以兼容既有 ``get_market_stat`` 调用与外部测试；实际委托 + 给泛化版 :meth:`_find_host_returning_data`。 """ - bad_host = self._host - ranked = ping_all(get_known_hosts(), self._port, 5.0) - - def _try(host: str) -> bool: - # 切换到候选 host 并实测；非空即视为该 host 可用 - self._reconnect(host) - return bool(self._execute(cmd)) - - new_host = find_working_host_sync(ranked, _try, save_best_host, bad_host) - if new_host is None: - # 全部候选都不可用，回退到原 host（保持状态可预测） - if self._host != bad_host: - self._reconnect(bad_host) - return [] - # _try 已把 client 切到 new_host 并执行过 cmd，重新取一次拿结果 - return self._execute(cmd) - - def _find_host_returning_bars(self, cmd: "BaseCommand[list[SecurityBar]]") -> list[SecurityBar]: - """空数据故障转移（K 线类）：与 :meth:`_find_host_returning_quotes` 同模式。 - - 指数/板块指数（880xxx 等）并非所有服务器都提供，延迟最低的不一定返回 - 数据（实测约 1/8 的服务器对所有指数返回空 body 且不报错）。故 K 线查询 - 空数据时按延迟顺序逐台实测，返回首台返回非空 bars 的结果。最多尝试 - ``_WORKING_HOST_MAX_ATTEMPTS`` 台；全失败回退原 host 并返回空。 - """ - bad_host = self._host - ranked = ping_all(get_known_hosts(), self._port, 5.0) - - def _try(host: str) -> bool: - # 切换到候选 host 并实测；非空即视为该 host 可用 - self._reconnect(host) - return bool(self._execute(cmd)) - - new_host = find_working_host_sync(ranked, _try, save_best_host, bad_host) - if new_host is None: - # 全部候选都不可用，回退到原 host（保持状态可预测） - if self._host != bad_host: - self._reconnect(bad_host) - return [] - # _try 已把 client 切到 new_host 并执行过 cmd，重新取一次拿结果 - return self._execute(cmd) + return self._find_host_returning_data(cmd) def _collect_transaction_records( self, @@ -1043,18 +1050,23 @@ class AsyncTdxClient(AsyncHeartbeatMixin): """ async with self._execute_lock: try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + # 连接失败：当前主机降权 + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: await asyncio.sleep(delay) await self._areconnect() try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移 new_host = await select_best_host_async( get_known_hosts(), @@ -1067,10 +1079,16 @@ class AsyncTdxClient(AsyncHeartbeatMixin): if new_host is not None: await self._areconnect(new_host) try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result async def get_security_count(self, market: Market) -> int: return await self._execute(GetSecurityCountCmd(market)) @@ -1180,8 +1198,9 @@ class AsyncTdxClient(AsyncHeartbeatMixin): """获取 K 线数据。``bar_time`` 见同步版 :meth:`get_security_bars`。""" cmd = GetSecurityBarsCmd(market, code, category, start, count) bars = await self._execute(cmd) + # 空数据故障转移（与 sync 版对称）：服务器连通但返回空/截断时逐台换台。 if not bars and self._auto_reconnect: - bars = await self._find_host_returning_bars(cmd) + bars = await self._find_host_returning_data(cmd) df = _to_df(bars) delta = _category_to_minutes(int(category)) is_intraday = delta is not None @@ -1207,8 +1226,9 @@ class AsyncTdxClient(AsyncHeartbeatMixin): """获取指数 K 线数据。``bar_time`` 见同步版 :meth:`get_index_bars`。""" cmd = GetIndexBarsCmd(market, code, category, start, count) bars = await self._execute(cmd) + # 空数据故障转移（与 sync 版对称）：指数 K 线截断/空时逐台换台。 if not bars and self._auto_reconnect: - bars = await self._find_host_returning_bars(cmd) + bars = await self._find_host_returning_data(cmd) df = _to_df(bars) delta = _category_to_minutes(int(category)) is_intraday = delta is not None @@ -1399,45 +1419,32 @@ class AsyncTdxClient(AsyncHeartbeatMixin): ) ) + async def _find_host_returning_data(self, cmd: "BaseCommand[_T]") -> _T: + """空数据故障转移（async）：与 sync ``_find_host_returning_data`` 对称。 + + 泛化版，支持任意可 ``bool()`` 判空的命令返回值（quotes / K 线 bars）。 + """ + bad_host = self._host + ranked = await asyncio.to_thread(ping_all, get_known_hosts(), self._port, 5.0) + + async def _try(host: str) -> bool: + await self._areconnect(host) + # mypy 对 async 闭包内泛型参数的推断会宽化为 BaseCommand[object] + # （sync 同模式可正确推断），此处为已知 mypy 限制，非真实类型错误。 + return bool(await self._execute(cmd)) # type: ignore[arg-type] + + new_host = await find_working_host_async(ranked, _try, save_best_host, bad_host) + if new_host is None: + if self._host != bad_host: + await self._areconnect(bad_host) + return [] # type: ignore[return-value] + return await self._execute(cmd) + async def _find_host_returning_quotes( self, cmd: "BaseCommand[list[SecurityQuote]]" ) -> list[SecurityQuote]: - """空数据故障转移（async）：与 sync ``_find_host_returning_quotes`` 对称。""" - bad_host = self._host - ranked = await asyncio.to_thread(ping_all, get_known_hosts(), self._port, 5.0) - - async def _try(host: str) -> bool: - await self._areconnect(host) - # mypy 对 async 闭包内泛型参数的推断会宽化为 BaseCommand[object] - # （sync 同模式可正确推断），此处为已知 mypy 限制，非真实类型错误。 - return bool(await self._execute(cmd)) # type: ignore[arg-type] - - new_host = await find_working_host_async(ranked, _try, save_best_host, bad_host) - if new_host is None: - if self._host != bad_host: - await self._areconnect(bad_host) - return [] - return await self._execute(cmd) - - async def _find_host_returning_bars( - self, cmd: "BaseCommand[list[SecurityBar]]" - ) -> list[SecurityBar]: - """空数据故障转移（K 线类，async）：与 sync ``_find_host_returning_bars`` 对称。""" - bad_host = self._host - ranked = await asyncio.to_thread(ping_all, get_known_hosts(), self._port, 5.0) - - async def _try(host: str) -> bool: - await self._areconnect(host) - # mypy 对 async 闭包内泛型参数的推断会宽化为 BaseCommand[object] - # （sync 同模式可正确推断），此处为已知 mypy 限制，非真实类型错误。 - return bool(await self._execute(cmd)) # type: ignore[arg-type] - - new_host = await find_working_host_async(ranked, _try, save_best_host, bad_host) - if new_host is None: - if self._host != bad_host: - await self._areconnect(bad_host) - return [] - return await self._execute(cmd) + """空数据故障转移（async quotes 专用薄封装）。""" + return await self._find_host_returning_data(cmd) async def _collect_transaction_records( self, diff --git a/src/easy_tdx/commands/security_bars.py b/src/easy_tdx/commands/security_bars.py index 7868baa..5a19675 100644 --- a/src/easy_tdx/commands/security_bars.py +++ b/src/easy_tdx/commands/security_bars.py @@ -87,13 +87,24 @@ class GetSecurityBarsCmd(BaseCommand[list[SecurityBar]]): # 注意：即使 bars 为空（第 1 条就崩）也 return 而非 raise —— # 服务器返回 0 条数据但 ret_count 撒谎是已知现象，返回空列表 # 让调用方分页重试比直接 500 更友好。 - _log.warning( - "K线响应在第 %d/%d 条处被截断（%s），已丢弃末尾残缺记录，返回前 %d 条", - i + 1, - ret_count, - e, - len(bars), - ) + if i == 0 and not bars: + # 第 1 条即崩且无任何已解析记录：典型"服务器空响应" + # （ret_count 撒谎）。用更明确的措辞，便于上层故障转移逻辑 + # 与人工排查识别"这是该服务器没数据，该换台"。 + _log.warning( + "K线响应为空（声称 %d 条但首条即解析失败：%s），" + "该服务器可能未提供此标的，返回空列表", + ret_count, + e, + ) + else: + _log.warning( + "K线响应在第 %d/%d 条处被截断（%s），已丢弃末尾残缺记录，返回前 %d 条", + i + 1, + ret_count, + e, + len(bars), + ) return bars # 差分还原（与 pytdx 完全一致） @@ -153,13 +164,21 @@ class GetIndexBarsCmd(GetSecurityBarsCmd): # 指数记录额外 4 字节：上涨家数 + 下跌家数（各 uint16 LE） pos += 4 except TdxDecodeError as e: - _log.warning( - "指数K线响应在第 %d/%d 条处被截断（%s），已丢弃末尾残缺记录，返回前 %d 条", - i + 1, - ret_count, - e, - len(bars), - ) + if i == 0 and not bars: + _log.warning( + "指数K线响应为空（声称 %d 条但首条即解析失败：%s），" + "该服务器可能未提供此指数，返回空列表", + ret_count, + e, + ) + else: + _log.warning( + "指数K线响应在第 %d/%d 条处被截断（%s），已丢弃末尾残缺记录，返回前 %d 条", + i + 1, + ret_count, + e, + len(bars), + ) return bars # 差分还原（与 pytdx 完全一致） diff --git a/src/easy_tdx/ex/client.py b/src/easy_tdx/ex/client.py index 0081bdf..d81bcb9 100644 --- a/src/easy_tdx/ex/client.py +++ b/src/easy_tdx/ex/client.py @@ -9,6 +9,7 @@ from types import TracebackType from typing import TypeVar from .._df import _apply_bar_time_align_bars, _category_to_minutes +from .._health import record_failure, record_success from .._reconnect import ( _RETRY_DELAYS, AsyncHeartbeatMixin, @@ -143,20 +144,26 @@ class ExTdxClient: 两阶段韧性与 A 股/MAC 统一（审计 #2）：先同主机重试 4 次，再跨主机 故障转移（测速切到另一台扩展行情服务器）。 + + 健康分联动与 A 股 client 一致（成功记 success、连接失败记 failure）。 """ try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: time.sleep(delay) self._reconnect() try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移 new_host = select_best_host_sync( get_ex_hosts(), @@ -169,10 +176,16 @@ class ExTdxClient: if new_host is not None: self._reconnect(new_host) try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 市场信息 @@ -400,21 +413,28 @@ class AsyncExTdxClient(AsyncHeartbeatMixin): self._start_heartbeat() async def _execute(self, cmd: "BaseCommand[_T]") -> _T: - """执行命令；断线时指数退避重试，同主机耗尽则跨主机故障转移。""" + """执行命令；断线时指数退避重试，同主机耗尽则跨主机故障转移。 + + 健康分联动与 sync 版对称。 + """ async with self._execute_lock: try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: await asyncio.sleep(delay) await self._areconnect() try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移 new_host = await select_best_host_async( get_ex_hosts(), @@ -427,10 +447,16 @@ class AsyncExTdxClient(AsyncHeartbeatMixin): if new_host is not None: await self._areconnect(new_host) try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 市场信息 diff --git a/src/easy_tdx/ex/mac_client.py b/src/easy_tdx/ex/mac_client.py index aa96f2e..8b9b266 100644 --- a/src/easy_tdx/ex/mac_client.py +++ b/src/easy_tdx/ex/mac_client.py @@ -15,6 +15,7 @@ from typing import Any, TypeVar import pandas as pd from .._df import _to_df +from .._health import record_failure, record_success from .._reconnect import ( _RETRY_DELAYS, AsyncHeartbeatMixin, @@ -159,12 +160,15 @@ class MacExClient: ``TdxConnectionError`` 与业务请求一样计入退避重试；``TdxCommandError`` （登录被拒等确定性失败）不重试，直接抛出。跨主机故障转移阶段同样遵循 ``connect + login`` 纳入重试的语义。 + + 健康分联动与 A 股 client 一致（成功记 success、连接失败记 failure）。 """ try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: time.sleep(delay) @@ -176,9 +180,12 @@ class MacExClient: try: self._conn.connect() self._login() - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移——测速切到另一台 MAC 扩展行情服务器 new_host = select_best_host_sync( get_mac_ex_hosts(), @@ -197,10 +204,16 @@ class MacExClient: try: self._conn.connect() self._login() - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 商品列表 @@ -617,13 +630,16 @@ class AsyncMacExClient(AsyncHeartbeatMixin): 每次重连后必须重新 ``_login()``（MAC 协议扩展行情特有）。登录握手期的 ``TdxConnectionError`` 与业务请求一样计入退避重试；``TdxCommandError`` （登录被拒等确定性失败）不重试，直接抛出。 + + 健康分联动与 sync 版对称。 """ async with self._execute_lock: try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: await asyncio.sleep(delay) @@ -635,9 +651,12 @@ class AsyncMacExClient(AsyncHeartbeatMixin): try: await self._conn.connect() await self._login() - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移 new_host = await select_best_host_async( get_mac_ex_hosts(), @@ -656,10 +675,16 @@ class AsyncMacExClient(AsyncHeartbeatMixin): try: await self._conn.connect() await self._login() - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 商品列表 diff --git a/src/easy_tdx/mac/client.py b/src/easy_tdx/mac/client.py index f8d3ae2..a99701d 100644 --- a/src/easy_tdx/mac/client.py +++ b/src/easy_tdx/mac/client.py @@ -13,6 +13,7 @@ from typing import Any, TypeVar import pandas as pd from .._df import _apply_bar_time_align_df, _period_to_minutes, _to_df +from .._health import record_failure, record_success from .._reconnect import ( _RETRY_DELAYS, AsyncHeartbeatMixin, @@ -268,20 +269,28 @@ class MacClient: 两阶段韧性与 TdxClient 对称：先同主机重试（``_RETRY_DELAYS``）， 再跨主机故障转移（重新测速切到另一台 MAC 服务器）。 + + 健康分联动：成功路径记 ``record_success``，连接失败记 ``record_failure``， + 与 A 股 client 一致（健康分全局共享，但 MAC 服务器 IP 与 A 股不重叠， + 不会互相干扰）。 """ try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: time.sleep(delay) self._reconnect() try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移——测速切到另一台 MAC 服务器再试一次。 # save_best_mac_host（而非 save_best_host）：MAC 服务器写入独立的 # best_mac_host 配置项，不污染标准 best_host（v1.19.4 修复的回归）。 @@ -296,10 +305,16 @@ class MacClient: if new_host is not None: self._reconnect(new_host) try: - return self._conn.execute(cmd) + result = self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 报价 @@ -1309,21 +1324,28 @@ class AsyncMacClient(AsyncHeartbeatMixin): # ------------------------------------------------------------------ # async def _execute(self, cmd: BaseCommand[_T]) -> _T: - """执行命令；断线时指数退避重试，同主机耗尽则跨主机故障转移。""" + """执行命令；断线时指数退避重试，同主机耗尽则跨主机故障转移。 + + 健康分联动与 sync 版对称（见 :meth:`MacClient._execute`）。 + """ async with self._execute_lock: try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) except TdxConnectionError: if not self._auto_reconnect: raise + record_failure(self._host) last_exc: TdxConnectionError | None = None for delay in _RETRY_DELAYS: await asyncio.sleep(delay) await self._areconnect() try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) # 第二阶段：跨主机故障转移 # save_best_mac_host：写入独立配置项，不污染标准 best_host # （v1.19.4 修复的回归：MAC failover 不可用 save_best_host） @@ -1338,10 +1360,16 @@ class AsyncMacClient(AsyncHeartbeatMixin): if new_host is not None: await self._areconnect(new_host) try: - return await self._conn.execute(cmd) + result = await self._conn.execute(cmd) + record_success(self._host) + return result except TdxConnectionError as e: last_exc = e + record_failure(self._host) raise last_exc # type: ignore[misc] + else: + record_success(self._host) + return result # ------------------------------------------------------------------ # # 报价 diff --git a/tests/unit/test_ex_reconnect.py b/tests/unit/test_ex_reconnect.py index 46c356b..da745b0 100644 --- a/tests/unit/test_ex_reconnect.py +++ b/tests/unit/test_ex_reconnect.py @@ -191,3 +191,128 @@ class TestMacExLoginRetriedOnConnectionError: client._execute(GetExMarketsCmd()) # 关键：_login 异常被纳入重试，4 次都跑了（而非第 1 次就逃逸） assert mock_sleep.call_count == len(_RETRY_DELAYS) + + +# --------------------------------------------------------------------------- # +# MAC/EX client 健康分联动（防 pattern-fix 回归） +# --------------------------------------------------------------------------- # + + +class TestExClientHealthTracking: + """锁定修复：MAC/EX 的 _execute 必须像 A 股 client 一样写健康分。 + + 此前只有 A 股 TdxClient/AsyncTdxClient 注入了 record_failure/record_success， + MAC/EX 的 6 个 _execute 漏改（它们的服务器 IP 与 A 股不重叠，失败时不会被 + 降权，功能残缺）。本测试防止再次漏改。 + """ + + def test_ex_client_failure_records_health(self) -> None: + """ExTdxClient 连接失败时应调 record_failure 降权当前 host。""" + from easy_tdx._health import reset_health + + reset_health() + try: + with ( + patch("easy_tdx.ex.client.ExTdxConnection") as mock_conn_cls, + patch("easy_tdx.ex.client.time.sleep"), + patch("easy_tdx.ex.client.select_best_host_sync", return_value=None), + ): + mock_conn = MagicMock() + mock_conn.execute.side_effect = TdxConnectionError("down") + mock_conn_cls.return_value = mock_conn + + client = ExTdxClient("ex-bad", auto_reconnect=True) + with pytest.raises(TdxConnectionError): + client._execute(GetExMarketsCmd()) + + from easy_tdx._health import get_score + + # ex-bad 经历首次 + 4 次重试共 5 次失败，score 应远低于 1.0 + assert get_score("ex-bad") < 1.0 + finally: + reset_health() + + def test_ex_client_success_records_health(self) -> None: + """ExTdxClient 首次成功应调 record_success（score 保持 1.0）。""" + from easy_tdx._health import get_score, reset_health + + reset_health() + try: + with patch("easy_tdx.ex.client.ExTdxConnection") as mock_conn_cls: + mock_conn = MagicMock() + mock_conn.execute.return_value = ["market"] + mock_conn_cls.return_value = mock_conn + + client = ExTdxClient("ex-good", auto_reconnect=True) + result = client._execute(GetExMarketsCmd()) + assert result == ["market"] + assert get_score("ex-good") == 1.0 + finally: + reset_health() + + def test_mac_ex_client_failure_records_health(self) -> None: + """MacExClient（含 _login 重连路径）连接失败也应降权当前 host。""" + from easy_tdx._health import reset_health + + reset_health() + try: + with ( + patch("easy_tdx.ex.mac_client.ExTdxConnection") as mock_conn_cls, + patch("easy_tdx.ex.mac_client.time.sleep"), + patch("easy_tdx.ex.mac_client.select_best_host_sync", return_value=None), + ): + mock_conn = MagicMock() + mock_conn.execute.side_effect = TdxConnectionError("down") + mock_conn_cls.return_value = mock_conn + + client = MacExClient("macex-bad", auto_reconnect=True) + with pytest.raises(TdxConnectionError): + client._execute(GetExMarketsCmd()) + + from easy_tdx._health import get_score + + assert get_score("macex-bad") < 1.0 + finally: + reset_health() + + def test_async_ex_client_failure_records_health(self) -> None: + """AsyncExTdxClient 连接失败时也应降权。""" + from easy_tdx._health import reset_health + + reset_health() + try: + + async def main() -> None: + with patch("easy_tdx.ex.client.AsyncExTdxConnection") as mock_conn_cls: + mock_conn = MagicMock() + + async def _execute(cmd: object) -> list[str]: + raise TdxConnectionError("down") + + async def _noop() -> None: + return None + + mock_conn.execute = _execute + mock_conn.close = _noop + mock_conn.connect = _noop + mock_conn_cls.return_value = mock_conn + + client = AsyncExTdxClient("aex-bad", auto_reconnect=True, heartbeat_interval=0) + with ( + patch("easy_tdx.ex.client.asyncio.sleep"), + patch( + "easy_tdx.ex.client.select_best_host_async", + new_callable=AsyncMock, + return_value=None, + ), + ): + with pytest.raises(TdxConnectionError): + await client._execute(GetExMarketsCmd()) + + asyncio.run(main()) + + from easy_tdx._health import get_score + + assert get_score("aex-bad") < 1.0 + finally: + reset_health() diff --git a/tests/unit/test_failover.py b/tests/unit/test_failover.py index 0b38af3..c756530 100644 --- a/tests/unit/test_failover.py +++ b/tests/unit/test_failover.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch import pytest +from easy_tdx._health import _FAILURE_DECAY, reset_health from easy_tdx._reconnect import ( _FAILOVER_PING_THROTTLE_SEC, _WORKING_HOST_MAX_ATTEMPTS, @@ -27,6 +28,23 @@ from easy_tdx.exceptions import TdxConnectionError from easy_tdx.models.bar import SecurityBar from easy_tdx.models.enums import KlineCategory, Market + +@pytest.fixture(autouse=True) +def _isolate_health_and_throttle(): + """每个测试前后重置健康分 + 节流时间戳，避免跨测试污染。 + + failover 与空数据转移现在会写健康分（record_failure/success）， + 若不重置，一个测试里降权的 host 会影响后续测试的 rank_by_health 排序。 + """ + import easy_tdx._reconnect as r + + r._last_failover_ts = 0.0 + reset_health() + yield + reset_health() + r._last_failover_ts = 0.0 + + # --------------------------------------------------------------------------- # # select_best_host_sync 单元逻辑 # --------------------------------------------------------------------------- # @@ -400,19 +418,72 @@ class TestMarketStatEmptyFailover: # --------------------------------------------------------------------------- # -# get_index_bars / get_security_bars 空数据故障转移 -# （与 TestMarketStatEmptyFailover 对称：指数/板块指数 880xxx 并非所有服务器都提供） +# 健康分联动：select_best_host / find_working_host 感知健康分 # --------------------------------------------------------------------------- # -class TestIndexBarsEmptyFailover: +class TestHealthAwareFailover: + """验证故障转移会读取/写入健康分：坏主机被降权后排序靠后。""" + + def test_select_best_host_skips_cooldown_host(self) -> None: + """冷却中的主机即使延迟最低，也不会被 select_best_host 选中。""" + from easy_tdx._health import record_failure + + # host-fast 连续失败进入冷却 + for _ in range(3): + record_failure("host-fast") + + ping_fn = MagicMock(return_value=[("host-fast", 0.01), ("host-slow", 0.10)]) + save_fn = MagicMock() + result = select_best_host_sync( + ["host-fast", "host-slow", "cur"], ping_fn, save_fn, 7709, 1.0, "cur" + ) + # host-fast 在冷却中被剔除，应选 host-slow + assert result == "host-slow" + save_fn.assert_called_once_with("host-slow") + + def test_find_working_host_records_failure_on_empty(self) -> None: + """候选返回空数据时记一次 failure（降权），下次轮询优先级下降。""" + from easy_tdx._health import get_score + + ranked = [("empty-host", 0.01), ("good-host", 0.02)] + try_fn = MagicMock(side_effect=[False, True]) # empty 空，good 非空 + save_fn = MagicMock() + + result = find_working_host_sync(ranked, try_fn, save_fn, "cur") + assert result == "good-host" + # empty-host 被记一次失败，score < 1.0 + assert get_score("empty-host") < 1.0 + # good-host 被记成功，score = 1.0 + assert get_score("good-host") == 1.0 + + def test_find_working_host_records_success_on_hit(self) -> None: + """命中的主机 score 恢复到 1.0。""" + from easy_tdx._health import get_score, record_failure + + # 先把 good-host 降权 + record_failure("good-host") + assert get_score("good-host") < 1.0 + + ranked = [("good-host", 0.01)] + try_fn = MagicMock(return_value=True) + save_fn = MagicMock() + + find_working_host_sync(ranked, try_fn, save_fn, "cur") + # 命中后 score 恢复（+0.2，但初始降权后 0.5+0.2=0.7，未到 1.0； + # 关键是比失败前上升了） + assert get_score("good-host") > _FAILURE_DECAY + + +# --------------------------------------------------------------------------- # +# get_index_bars / get_security_bars 空数据故障转移 +# （指数/板块指数 880xxx 并非所有服务器都提供，空时逐台实测切 host） +# --------------------------------------------------------------------------- # + + +class TestBarsEmptyFailover: """K 线空数据故障转移——验证 get_index_bars/get_security_bars 空时逐台实测切 host。""" - def setup_method(self) -> None: - import easy_tdx._reconnect as r - - r._last_failover_ts = 0.0 - def _make_bar(self) -> SecurityBar: """构造一根字段合法的日 K，让 get_index_bars 下游处理走通。""" return SecurityBar( @@ -478,7 +549,7 @@ class TestIndexBarsEmptyFailover: with ( patch.object(client, "_execute", return_value=[bar]) as mock_exec, - patch.object(client, "_find_host_returning_bars") as mock_failover, + patch.object(client, "_find_host_returning_data") as mock_failover, ): df = client.get_index_bars(Market.SH, "880008", KlineCategory.DAY, 0, 10) @@ -492,7 +563,7 @@ class TestIndexBarsEmptyFailover: with ( patch.object(client, "_execute", return_value=[]) as mock_exec, - patch.object(client, "_find_host_returning_bars") as mock_failover, + patch.object(client, "_find_host_returning_data") as mock_failover, ): df = client.get_index_bars(Market.SH, "880008", KlineCategory.DAY, 0, 10) diff --git a/tests/unit/test_health.py b/tests/unit/test_health.py new file mode 100644 index 0000000..68d0980 --- /dev/null +++ b/tests/unit/test_health.py @@ -0,0 +1,149 @@ +"""服务器健康分（health score）引擎单元测试。 + +覆盖： + - record_failure 乘性衰减 + 连续失败触发冷却 + - record_success 加性恢复 + 重置计数与冷却 + - is_in_cooldown / get_score 读取语义 + - rank_by_health：冷却剔除 + 有效延迟（latency/score）排序 + - 全健康时 rank_by_health 近似恒等映射（向后兼容保证） +""" + +from __future__ import annotations + +import time + +import pytest + +from easy_tdx._health import ( + _COOLDOWN_FAIL_THRESHOLD, + _FAILURE_DECAY, + _SUCCESS_RECOVER, + get_score, + is_in_cooldown, + rank_by_health, + record_failure, + record_success, + reset_health, +) + + +@pytest.fixture(autouse=True) +def _isolate_health(): + """每个测试前后清空健康记录，避免跨测试污染。""" + reset_health() + yield + reset_health() + + +# --------------------------------------------------------------------------- # +# record_failure / record_success +# --------------------------------------------------------------------------- # + + +class TestRecordFailure: + def test_single_failure_decays_score(self) -> None: + s = record_failure("h1") + assert s == pytest.approx(_FAILURE_DECAY) + assert get_score("h1") == pytest.approx(_FAILURE_DECAY) + + def test_repeated_failure_decays_multiplicatively(self) -> None: + record_failure("h1") + record_failure("h1") + s = record_failure("h1") + assert s == pytest.approx(_FAILURE_DECAY**3) + + def test_consecutive_failures_below_threshold_no_cooldown(self) -> None: + for _ in range(_COOLDOWN_FAIL_THRESHOLD - 1): + record_failure("h1") + assert not is_in_cooldown("h1") + + def test_consecutive_failures_at_threshold_enters_cooldown(self) -> None: + for _ in range(_COOLDOWN_FAIL_THRESHOLD): + record_failure("h1") + assert is_in_cooldown("h1") + + def test_score_never_drops_below_floor(self) -> None: + for _ in range(100): + record_failure("h1") + assert get_score("h1") > 0 + + +class TestRecordSuccess: + def test_success_recovers_score_additively(self) -> None: + record_failure("h1") # score = 0.5 + record_success("h1") + assert get_score("h1") == pytest.approx(_FAILURE_DECAY + _SUCCESS_RECOVER) + + def test_success_caps_at_one(self) -> None: + record_success("h1") + record_success("h1") + assert get_score("h1") == pytest.approx(1.0) + + def test_success_resets_consecutive_failures_and_cooldown(self) -> None: + for _ in range(_COOLDOWN_FAIL_THRESHOLD): + record_failure("h1") + assert is_in_cooldown("h1") + record_success("h1") + assert not is_in_cooldown("h1") + # 再次失败一次不应立即进冷却（计数已重置） + record_failure("h1") + assert not is_in_cooldown("h1") + + +# --------------------------------------------------------------------------- # +# is_in_cooldown +# --------------------------------------------------------------------------- # + + +class TestIsInCooldown: + def test_unknown_host_not_in_cooldown(self) -> None: + assert not is_in_cooldown("never-seen") + + def test_cooldown_expires(self) -> None: + # 手动模拟过期：记录到阈值进入冷却后，快进时间戳 + for _ in range(_COOLDOWN_FAIL_THRESHOLD): + record_failure("h1") + assert is_in_cooldown("h1") + # 直接篡改内部状态模拟冷却过期（避免真睡 120s） + from easy_tdx._health import _BOOK + + with _BOOK.lock: + _BOOK.hosts["h1"].cooldown_until = time.monotonic() - 1.0 + assert not is_in_cooldown("h1") + + +# --------------------------------------------------------------------------- # +# rank_by_health +# --------------------------------------------------------------------------- # + + +class TestRankByHealth: + def test_identity_when_all_healthy(self) -> None: + """全健康时输出排序与输入一致（向后兼容关键保证）。""" + ranked = [("a", 0.01), ("b", 0.05), ("c", 0.10)] + assert rank_by_health(ranked) == ranked + + def test_filters_out_cooldown_hosts(self) -> None: + for _ in range(_COOLDOWN_FAIL_THRESHOLD): + record_failure("bad") + ranked = [("bad", 0.01), ("good", 0.02)] + result = rank_by_health(ranked) + assert "bad" not in [h for h, _ in result] + assert result == [("good", 0.02)] + + def test_low_score_host_pushed_back(self) -> None: + # b 延迟最低但 score 被打到很低，使其有效延迟（latency/score）反超 a。 + # 2 次失败 → b score = 0.25；有效延迟 = 0.03/0.25 = 0.12。 + # a 全健康，有效延迟 = 0.05/1.0 = 0.05 < 0.12 → a 应排前。 + for _ in range(2): + record_failure("b") # 不进冷却（阈值 3），但 score 衰减到 0.25 + ranked = [("b", 0.03), ("a", 0.05)] + result = rank_by_health(ranked) + assert result[0][0] == "a" + + def test_empty_input(self) -> None: + assert rank_by_health([]) == [] + + def test_preserves_latency_order_among_equal_scores(self) -> None: + ranked = [("a", 0.01), ("b", 0.02), ("c", 0.03)] + assert rank_by_health(ranked) == ranked