Coverage for src\core\heartbeat.py: 55%

1"""

2Heartbeat Module - 心跳模块

4PAO系统的核心机制之一：

5- 设备定期发送心跳，证明自己还活着

6- 心跳管理器追踪所有设备的心跳状态

7- 超时未心跳的设备标记为离线

8- 支持心跳状态变化的回调通知

9"""

11import asyncio

12import logging

13import time

14from typing import Dict, Callable, Optional, List, Any

15from dataclasses import dataclass, field

16from datetime import datetime, timedelta

17from enum import Enum

19logger = logging.getLogger(__name__)

22class HeartbeatStatus(str, Enum):

23 """心跳状态"""

24 ALIVE = "alive" # 设备存活

25 WARNING = "warning" # 心跳警告（快超时）

26 DEAD = "dead" # 设备离线

29@dataclass

30class HeartbeatInfo:

31 """心跳信息"""

32 device_id: str

33 last_heartbeat: float # 时间戳

34 status: HeartbeatStatus = HeartbeatStatus.ALIVE

35 consecutive_failures: int = 0 # 连续失败次数

37 def time_since_heartbeat(self) -> float:

38 """距离上次心跳的秒数"""

39 return time.time() - self.last_heartbeat

42@dataclass

43class HeartbeatConfig:

44 """心跳配置"""

45 interval: int = 30 # 心跳间隔（秒）

46 timeout: int = 90 # 超时时间（秒），超过这个时间没心跳认为离线

47 warning_threshold: float = 0.7 # 警告阈值，timeout的70%时触发warning状态

48 check_interval: int = 10 # 检查间隔（秒）

51class HeartbeatCallback:

52 """心跳回调函数类型"""

53 # 回调签名: (device_id: str, old_status: HeartbeatStatus, new_status: HeartbeatStatus) -> None

54 pass

57class HeartbeatManager:

58 """

59 心跳管理器

61 核心功能：

62 1. 追踪所有设备的心跳状态

63 2. 定期检查超时设备

64 3. 状态变化时通知订阅者

65 4. 本设备也向其他设备发送心跳

66 """

68 def __init__(self, device_id: str, config: Optional[HeartbeatConfig] = None):

69 """

70 初始化心跳管理器

72 Args:

73 device_id: 本设备ID

74 config: 心跳配置

75 """

76 self.device_id = device_id

77 self.config = config or HeartbeatConfig()

79 # 心跳信息存储: device_id -> HeartbeatInfo

80 self._heartbeats: Dict[str, HeartbeatInfo] = {}

82 # 回调函数列表

83 self._callbacks: List[Callable] = []

85 # 任务

86 self._send_task: Optional[asyncio.Task] = None

87 self._check_task: Optional[asyncio.Task] = None

88 self._running = False

90 # 外部通信接口（发送心跳用）

91 self._sender: Optional[Any] = None

93 logger.info(f"心跳管理器初始化完成 [设备: {device_id}]")

95 def set_sender(self, sender: Any):

96 """

97 设置发送器（通常是Communication模块）

99 Args:

100 sender: 有 send_message 方法的对象

101 """

102 self._sender = sender

103

104 def register_callback(self, callback: Callable):

105 """

106 注册心跳状态变化回调

107

108 Args:

109 callback: 回调函数，签名: (device_id, old_status, new_status) -> None

110 """

111 self._callbacks.append(callback)

112 logger.debug(f"已注册心跳回调: {callback.__name__ if hasattr(callback, '__name__') else str(callback)}")

113

114 def _notify_status_change(self, device_id: str, old_status: HeartbeatStatus, new_status: HeartbeatStatus):

115 """通知所有回调设备状态变化"""

116 if old_status == new_status:

117 return

118

119 logger.info(f"设备状态变化 [{device_id}]: {old_status.value} -> {new_status.value}")

120

121 for callback in self._callbacks:

122 try:

123 callback(device_id, old_status, new_status)

124 except Exception as e:

125 logger.error(f"执行心跳回调失败: {e}")

126

127 async def start(self):

128 """启动心跳管理器"""

129 if self._running:

130 logger.warning("心跳管理器已在运行")

131 return

132

133 self._running = True

134

135 # 启动心跳发送任务

136 self._send_task = asyncio.create_task(self._heartbeat_sender())

137

138 # 启动心跳检查任务

139 self._check_task = asyncio.create_task(self._heartbeat_checker())

140

141 logger.info("心跳管理器已启动")

142

143 async def stop(self):

144 """停止心跳管理器"""

145 if not self._running:

146 return

147

148 self._running = False

149

150 # 取消任务

151 if self._send_task:

152 self._send_task.cancel()

153 try:

154 await self._send_task

155 except asyncio.CancelledError:

156 pass

157

158 if self._check_task:

159 self._check_task.cancel()

160 try:

161 await self._check_task

162 except asyncio.CancelledError:

163 pass

164

165 logger.info("心跳管理器已停止")

166

167 async def _heartbeat_sender(self):

168 """定期发送本设备心跳"""

169 while self._running:

170 try:

171 await asyncio.sleep(self.config.interval)

172

173 if self._sender:

174 # 通过通信模块发送心跳

175 success = await self._sender.send_message(

176 peer_id="broadcast", # 广播心跳

177 content={

178 "device_id": self.device_id,

179 "timestamp": time.time(),

180 "type": "heartbeat_request"

181 },

182 msg_type="heartbeat"

183 )

184

185 if success:

186 logger.debug(f"心跳已发送 [设备: {self.device_id}]")

187 else:

188 logger.warning(f"心跳发送失败 [设备: {self.device_id}]")

189

190 except asyncio.CancelledError:

191 break

192 except Exception as e:

193 logger.error(f"心跳发送异常: {e}")

194

195 async def _heartbeat_checker(self):

196 """定期检查所有设备的心跳状态"""

197 while self._running:

198 try:

199 await asyncio.sleep(self.config.check_interval)

200

201 now = time.time()

202 timeout_threshold = self.config.timeout

203 warning_threshold = self.config.timeout * self.config.warning_threshold

204

205 for device_id, info in list(self._heartbeats.items()):

206 elapsed = now - info.last_heartbeat

207

208 old_status = info.status

209

210 if elapsed >= timeout_threshold:

211 info.status = HeartbeatStatus.DEAD

212 elif elapsed >= warning_threshold:

213 info.status = HeartbeatStatus.WARNING

214

215 if old_status != info.status:

216 self._notify_status_change(device_id, old_status, info.status)

217

218 except asyncio.CancelledError:

219 break

220 except Exception as e:

221 logger.error(f"心跳检查异常: {e}")

222

223 def receive_heartbeat(self, device_id: str, timestamp: float):

224 """

225 收到其他设备的心跳

226

227 Args:

228 device_id: 发送心跳的设备ID

229 timestamp: 心跳时间戳

230 """

231 if device_id not in self._heartbeats:

232 # 新设备

233 self._heartbeats[device_id] = HeartbeatInfo(

234 device_id=device_id,

235 last_heartbeat=timestamp,

236 status=HeartbeatStatus.ALIVE

237 )

238 logger.info(f"新设备心跳注册 [{device_id}]")

239 else:

240 info = self._heartbeats[device_id]

241 old_status = info.status

242

243 info.last_heartbeat = timestamp

244 info.status = HeartbeatStatus.ALIVE

245 info.consecutive_failures = 0

246

247 if old_status != HeartbeatStatus.ALIVE:

248 self._notify_status_change(device_id, old_status, HeartbeatStatus.ALIVE)

249

250 logger.debug(f"收到心跳 [{device_id}] elapsed={time.time() - timestamp:.1f}s")

251

252 def get_status(self, device_id: str) -> Optional[HeartbeatStatus]:

253 """

254 获取设备心跳状态

255

256 Args:

257 device_id: 设备ID

258

259 Returns:

260 HeartbeatStatus 或 None（设备不存在）

261 """

262 info = self._heartbeats.get(device_id)

263 return info.status if info else None

264

265 def get_all_status(self) -> Dict[str, HeartbeatStatus]:

266 """

267 获取所有设备的心跳状态

268

269 Returns:

270 device_id -> HeartbeatStatus 的字典

271 """

272 return {

273 device_id: info.status

274 for device_id, info in self._heartbeats.items()

275 }

276

277 def get_device_info(self, device_id: str) -> Optional[HeartbeatInfo]:

278 """

279 获取设备详细信息

280

281 Args:

282 device_id: 设备ID

283

284 Returns:

285 HeartbeatInfo 或 None

286 """

287 return self._heartbeats.get(device_id)

288

289 def get_online_devices(self) -> List[str]:

290 """

291 获取所有在线设备ID

292

293 Returns:

294 在线设备ID列表

295 """

296 return [

297 device_id

298 for device_id, info in self._heartbeats.items()

299 if info.status != HeartbeatStatus.DEAD

300 ]

301

302 def get_alive_devices(self) -> List[str]:

303 """

304 获取所有存活设备ID（完全正常）

305

306 Returns:

307 存活设备ID列表

308 """

309 return [

310 device_id

311 for device_id, info in self._heartbeats.items()

312 if info.status == HeartbeatStatus.ALIVE

313 ]

314

315 def remove_device(self, device_id: str):

316 """

317 移除设备心跳追踪

318

319 Args:

320 device_id: 设备ID

321 """

322 if device_id in self._heartbeats:

323 del self._heartbeats[device_id]

324 logger.info(f"移除设备心跳追踪 [{device_id}]")

325

326 def get_summary(self) -> Dict[str, Any]:

327 """

328 获取心跳管理器摘要

329

330 Returns:

331 包含统计信息的字典

332 """

333 alive = warning = dead = 0

334

335 for info in self._heartbeats.values():

336 if info.status == HeartbeatStatus.ALIVE:

337 alive += 1

338 elif info.status == HeartbeatStatus.WARNING:

339 warning += 1

340 else:

341 dead += 1

342

343 return {

344 "device_id": self.device_id,

345 "total_devices": len(self._heartbeats),

346 "alive": alive,

347 "warning": warning,

348 "dead": dead,

349 "config": {

350 "interval": self.config.interval,

351 "timeout": self.config.timeout,

352 "check_interval": self.config.check_interval

353 }

354 }

355

356

357class HeartbeatProtocol:

358 """

359 心跳协议处理器

360

361 处理心跳相关的协议消息

362 """

363

364 @staticmethod

365 def create_heartbeat_message(sender_id: str) -> Dict:

366 """

367 创建心跳消息

368

369 Args:

370 sender_id: 发送者设备ID

371

372 Returns:

373 心跳消息字典

374 """

375 return {

376 "type": "heartbeat",

377 "sender_id": sender_id,

378 "timestamp": time.time()

379 }

380

381 @staticmethod

382 def create_heartbeat_response(sender_id: str, target_id: str) -> Dict:

383 """

384 创建心跳响应消息

385

386 Args:

387 sender_id: 发送者设备ID

388 target_id: 目标设备ID

389

390 Returns:

391 心跳响应消息字典

392 """

393 return {

394 "type": "heartbeat_response",

395 "sender_id": sender_id,

396 "target_id": target_id,

397 "timestamp": time.time()

398 }

399

400 @staticmethod

401 def is_heartbeat_message(message: Dict) -> bool:

402 """

403 判断是否为心跳消息

404

405 Args:

406 message: 消息字典

407

408 Returns:

409 bool

410 """

411 msg_type = message.get("type", "")

412 return msg_type in ("heartbeat", "heartbeat_response")

413

414 @staticmethod

415 def parse_heartbeat(message: Dict) -> Optional[Dict]:

416 """

417 解析心跳消息

418

419 Args:

420 message: 消息字典

421

422 Returns:

423 解析后的数据或None

424 """

425 if not HeartbeatProtocol.is_heartbeat_message(message):

426 return None

427

428 return {

429 "type": message.get("type"),

430 "sender_id": message.get("sender_id"),

431 "timestamp": message.get("timestamp", 0)

432 }