fix(transport): ping 容错 TdxConnectionError,避免单台服务器拖垮测速

ping_host 仅 except OSError,但握手期 _recv_exact_sock 抛的 TdxConnectionError
继承自 TdxError(Exception) 而非 OSError,逃出捕获后经 ping_all 的 fut.result()
重新抛出,导致非交易时间服务器 accept 后立即 FIN 时整个 easy-tdx ping 崩溃。
- ping_host: except (OSError, TdxConnectionError),对齐 docstring 返回 None
- ping_all: fut.result() 加 try/except 防御层,单 host 失败只跳过不崩
- 新增 2 例回归测试
This commit is contained in:
Justin Gu
2026-06-13 21:09:55 +08:00
parent a6ed0eac16
commit 095c88f735
2 changed files with 52 additions and 3 deletions
+10 -2
View File
@@ -56,7 +56,10 @@ def ping_host(
if hdr.zipsize > 0:
_recv_exact_sock(sock, hdr.zipsize)
return time.monotonic() - t0
except OSError:
except (OSError, TdxConnectionError):
# OSError: 连接/超时层失败;TdxConnectionError: 握手期服务器关闭连接
# _recv_exact_sock 抛出,继承自 TdxError 而非 OSError)。
# 两者均属"服务器不可用",按 docstring 返回 None,不拖垮整个 ping_all。
return None
finally:
try:
@@ -85,7 +88,12 @@ def ping_all(
futures = {pool.submit(ping_host, h, port, timeout): h for h in hosts}
for fut in concurrent.futures.as_completed(futures):
host = futures[fut]
latency = fut.result()
try:
latency = fut.result()
except Exception:
# 防御层:即使 ping_host 因意外原因抛异常,也只跳过该 host,
# 不让单个服务器拖垮整个 ping_all / `easy-tdx ping` 命令。
continue
if latency is not None:
results.append((host, latency))
results.sort(key=lambda t: t[1])
+42 -1
View File
@@ -3,7 +3,7 @@
from unittest.mock import patch
from easy_tdx.exceptions import TdxConnectionError
from easy_tdx.transport.sync import TdxConnection
from easy_tdx.transport.sync import TdxConnection, ping_all, ping_host
class _FakeSocket:
@@ -18,6 +18,9 @@ class _FakeSocket:
def connect(self, address: tuple[str, int]) -> None:
self.connected_to = address
def sendall(self, data: bytes) -> None:
pass
def close(self) -> None:
self.closed = True
@@ -45,3 +48,41 @@ def test_sync_connection_closes_socket_when_setup_fails() -> None:
assert sock.connected_to == ("127.0.0.1", 7709)
assert sock.closed is True
assert conn._sock is None
def test_ping_host_returns_none_when_server_closes_during_handshake() -> None:
"""握手期服务器关闭连接(_recv_exact_sock 抛 TdxConnectionError)应返回 None。
回归:ping_host 旧版仅 except OSError,未捕获 TdxConnectionError(继承自
TdxError(Exception) 而非 OSError),导致单个服务器握手失败就让整个
ping_all / `easy-tdx ping` 命令崩溃。非交易时间服务器 accept 后立即 FIN
时必现。
"""
sock = _FakeSocket()
with (
patch("easy_tdx.transport.sync.socket.socket", return_value=sock),
patch(
"easy_tdx.transport.sync._recv_exact_sock",
side_effect=TdxConnectionError("连接被服务器关闭"),
),
):
result = ping_host("192.0.2.1", port=7709, timeout=0.2)
assert result is None
assert sock.closed is True
def test_ping_all_skips_handshake_failure_without_crash() -> None:
"""ping_all 中任一服务器握手失败(TdxConnectionError)应被跳过,不崩溃。"""
def fake_ping(host: str, port: int, timeout: float) -> float | None:
if host == "bad":
raise TdxConnectionError("连接被服务器关闭")
return 0.05
with patch("easy_tdx.transport.sync.ping_host", side_effect=fake_ping):
results = ping_all(hosts=["good", "bad"], port=7709, timeout=0.2)
hosts_ok = [h for h, _ in results]
assert "good" in hosts_ok
assert "bad" not in hosts_ok