mirror of
https://ghfast.top/https://github.com/aeroxw/easy_tdx_max.git
synced 2026-09-12 15:44:18 +08:00
feat: AsyncTdxClient 增加心跳保活机制
1. 在 AsyncTdxClient 中增加后台心跳任务,默认每 60 秒发送一次请求。 2. 支持在连接成功后自动启动,并在连接关闭时自动停止心跳任务。 3. 增加单元测试 tests/unit/test_heartbeat.py 验证心跳循环与清理逻辑。
This commit is contained in:
+39
-1
@@ -293,13 +293,16 @@ class AsyncTdxClient:
|
||||
port: int = _DEFAULT_PORT,
|
||||
timeout: float = 15.0,
|
||||
auto_reconnect: bool = True,
|
||||
heartbeat_interval: float = 60.0,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._timeout = timeout
|
||||
self._auto_reconnect = auto_reconnect
|
||||
self._heartbeat_interval = heartbeat_interval
|
||||
self._conn = AsyncTdxConnection(host, port, timeout)
|
||||
self._execute_lock = asyncio.Lock()
|
||||
self._heartbeat_task: asyncio.Task[None] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_best_host(
|
||||
@@ -309,11 +312,12 @@ class AsyncTdxClient:
|
||||
timeout: float = 15.0,
|
||||
ping_timeout: float = 5.0,
|
||||
auto_reconnect: bool = True,
|
||||
heartbeat_interval: float = 60.0,
|
||||
) -> "AsyncTdxClient":
|
||||
"""测量 hosts 中所有服务器延迟,选最低延迟的建立连接。"""
|
||||
ranked = ping_all(hosts, port, ping_timeout)
|
||||
best = ranked[0][0] if ranked else hosts[0]
|
||||
return cls(best, port, timeout, auto_reconnect)
|
||||
return cls(best, port, timeout, auto_reconnect, heartbeat_interval)
|
||||
|
||||
@staticmethod
|
||||
def ping_all(
|
||||
@@ -326,8 +330,10 @@ class AsyncTdxClient:
|
||||
|
||||
async def connect(self) -> None:
|
||||
await self._conn.connect()
|
||||
self._start_heartbeat()
|
||||
|
||||
async def close(self) -> None:
|
||||
await self._stop_heartbeat()
|
||||
await self._conn.close()
|
||||
|
||||
async def __aenter__(self) -> "AsyncTdxClient":
|
||||
@@ -342,6 +348,38 @@ class AsyncTdxClient:
|
||||
) -> None:
|
||||
await self.close()
|
||||
|
||||
def _start_heartbeat(self) -> None:
|
||||
"""启动后台心跳任务。"""
|
||||
if self._heartbeat_interval <= 0:
|
||||
return
|
||||
if self._heartbeat_task is not None:
|
||||
self._heartbeat_task.cancel()
|
||||
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
|
||||
|
||||
async def _stop_heartbeat(self) -> None:
|
||||
"""停止并清理心跳任务。"""
|
||||
if self._heartbeat_task:
|
||||
self._heartbeat_task.cancel()
|
||||
try:
|
||||
await self._heartbeat_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._heartbeat_task = None
|
||||
|
||||
async def _heartbeat_loop(self) -> None:
|
||||
"""心跳循环:定期发送轻量级请求保活。"""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(self._heartbeat_interval)
|
||||
# 使用 get_security_count 作为心跳包
|
||||
await self.get_security_count(Market.SH)
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception:
|
||||
# 心跳失败通常意味着连接已断开
|
||||
# 下一次正常的业务请求或下一次心跳会通过 _execute 触发重连
|
||||
pass
|
||||
|
||||
async def _execute(self, cmd: "BaseCommand[_T]") -> _T:
|
||||
"""执行命令;断线时尝试重连一次再重试(若 auto_reconnect=True)。"""
|
||||
async with self._execute_lock:
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
"""心跳机制单元测试。"""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
from xmtdx import AsyncTdxClient, Market
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_sends_periodically():
|
||||
# 模拟连接和执行
|
||||
with patch("xmtdx.client.AsyncTdxConnection") as mock_conn_cls:
|
||||
mock_conn = mock_conn_cls.return_value
|
||||
mock_conn.connect = AsyncMock()
|
||||
mock_conn.close = AsyncMock()
|
||||
|
||||
# 记录调用次数
|
||||
call_count = 0
|
||||
async def mock_execute(cmd):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return 5 # 模拟 get_security_count 返回值
|
||||
|
||||
mock_conn.execute.side_effect = mock_execute
|
||||
|
||||
# 设置非常短的心跳间隔以便测试
|
||||
client = AsyncTdxClient("127.0.0.1", heartbeat_interval=0.1)
|
||||
await client.connect()
|
||||
|
||||
# 等待几次心跳周期
|
||||
await asyncio.sleep(0.35)
|
||||
|
||||
await client.close()
|
||||
|
||||
# 0.35s 应该触发约 3 次心跳 (0.1, 0.2, 0.3)
|
||||
assert call_count >= 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_stops_on_close():
|
||||
with patch("xmtdx.client.AsyncTdxConnection") as mock_conn_cls:
|
||||
mock_conn = mock_conn_cls.return_value
|
||||
mock_conn.connect = AsyncMock()
|
||||
mock_conn.close = AsyncMock()
|
||||
mock_conn.execute = AsyncMock(return_value=5)
|
||||
|
||||
client = AsyncTdxClient("127.0.0.1", heartbeat_interval=0.01)
|
||||
await client.connect()
|
||||
assert client._heartbeat_task is not None
|
||||
|
||||
task = client._heartbeat_task
|
||||
await client.close()
|
||||
|
||||
assert client._heartbeat_task is None
|
||||
assert task.done() or task.cancelled()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 手动跑一下
|
||||
async def run():
|
||||
await test_heartbeat_sends_periodically()
|
||||
await test_heartbeat_stops_on_close()
|
||||
print("Heartbeat tests passed!")
|
||||
|
||||
asyncio.run(run())
|
||||
Reference in New Issue
Block a user