Coverage for src\p2p_network.py: 0%

1""" 

2P2P Network Module - P2P网络优化模块 

3 

4提供连接池管理、断线重连、消息队列等功能 

5""" 

6 

7import asyncio 

8import logging 

9import time 

10from typing import Dict, Optional, Set 

11from dataclasses import dataclass, field 

12from enum import Enum 

13from collections import deque 

14import uuid 

15 

16logger = logging.getLogger(__name__) 

17 

18 

19class ConnectionState(Enum): 

20 """连接状态""" 

21 DISCONNECTED = "disconnected" 

22 CONNECTING = "connecting" 

23 CONNECTED = "connected" 

24 RECONNECTING = "reconnecting" 

25 FAILED = "failed" 

26 

27 

28@dataclass 

29class ConnectionInfo: 

30 """连接信息""" 

31 peer_id: str 

32 peer_name: str 

33 state: ConnectionState = ConnectionState.DISCONNECTED 

34 connected_at: float = 0 

35 last_heartbeat: float = 0 

36 reconnect_attempts: int = 0 

37 max_reconnect_attempts: int = 5 

38 latency_ms: float = 0 

39 

40 

41@dataclass 

42class Message: 

43 """消息结构""" 

44 msg_id: str 

45 msg_type: str 

46 payload: dict 

47 created_at: float = field(default_factory=time.time) 

48 retry_count: int = 0 

49 max_retries: int = 3 

50 

51 

52class ConnectionPool: 

53 """连接池管理器""" 

54 

55 def __init__(self, max_connections: int = 10): 

56 self.max_connections = max_connections 

57 self.connections: Dict[str, ConnectionInfo] = {} 

58 self.connection_lock = asyncio.Lock() 

59 self._callbacks: Dict[str, callable] = {} 

60 

61 async def add_connection(self, peer_id: str, peer_name: str) -> ConnectionInfo: 

62 """添加新连接""" 

63 async with self.connection_lock: 

64 if peer_id in self.connections: 

65 return self.connections[peer_id] 

66 

67 if len(self.connections) >= self.max_connections: 

68 # 移除最老的连接 

69 oldest = min( 

70 self.connections.values(), 

71 key=lambda x: x.connected_at 

72 ) 

73 await self.remove_connection(oldest.peer_id) 

74 

75 conn_info = ConnectionInfo( 

76 peer_id=peer_id, 

77 peer_name=peer_name, 

78 state=ConnectionState.CONNECTING, 

79 connected_at=time.time() 

80 ) 

81 self.connections[peer_id] = conn_info 

82 logger.info(f"添加连接到池: {peer_id}") 

83 return conn_info 

84 

85 async def update_state(self, peer_id: str, state: ConnectionState) -> None: 

86 """更新连接状态""" 

87 async with self.connection_lock: 

88 if peer_id in self.connections: 

89 self.connections[peer_id].state = state 

90 if state == ConnectionState.CONNECTED: 

91 self.connections[peer_id].connected_at = time.time() 

92 self.connections[peer_id].reconnect_attempts = 0 

93 

94 async def remove_connection(self, peer_id: str) -> None: 

95 """移除连接""" 

96 async with self.connection_lock: 

97 if peer_id in self.connections: 

98 del self.connections[peer_id] 

99 logger.info(f"从池中移除连接: {peer_id}") 

100 

101 async def get_connection(self, peer_id: str) -> Optional[ConnectionInfo]: 

102 """获取连接信息""" 

103 return self.connections.get(peer_id) 

104 

105 async def get_all_connections(self) -> Dict[str, ConnectionInfo]: 

106 """获取所有连接""" 

107 return self.connections.copy() 

108 

109 def register_callback(self, event: str, callback: callable) -> None: 

110 """注册状态变化回调""" 

111 self._callbacks[event] = callback 

112 

113 async def trigger_callback(self, event: str, *args, **kwargs) -> None: 

114 """触发回调""" 

115 if event in self._callbacks: 

116 try: 

117 await self._callbacks[event](*args, **kwargs) 

118 except Exception as e: 

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

120 

121 

122class ReconnectionManager: 

123 """断线重连管理器""" 

124 

125 def __init__(self, base_delay: float = 1.0, max_delay: float = 60.0): 

126 self.base_delay = base_delay 

127 self.max_delay = max_delay 

128 self._reconnect_tasks: Dict[str, asyncio.Task] = {} 

129 self._running = False 

130 

131 def calculate_delay(self, attempt: int) -> float: 

132 """计算重连延迟（指数退避）""" 

133 delay = min(self.base_delay * (2 ** attempt), self.max_delay) 

134 # 添加随机抖动 

135 import random 

136 return delay * (0.5 + random.random() * 0.5) 

137 

138 async def start_reconnect( 

139 self, 

140 peer_id: str, 

141 connect_func: callable, 

142 max_attempts: int = 5 

143 ) -> bool: 

144 """启动重连""" 

145 if peer_id in self._reconnect_tasks: 

146 logger.warning(f"重连任务已存在: {peer_id}") 

147 return False 

148 

149 self._running = True 

150 task = asyncio.create_task( 

151 self._reconnect_loop(peer_id, connect_func, max_attempts) 

152 ) 

153 self._reconnect_tasks[peer_id] = task 

154 return True 

155 

156 async def stop_reconnect(self, peer_id: str) -> None: 

157 """停止重连""" 

158 if peer_id in self._reconnect_tasks: 

159 self._reconnect_tasks[peer_id].cancel() 

160 del self._reconnect_tasks[peer_id] 

161 logger.info(f"停止重连任务: {peer_id}") 

162 

163 async def _reconnect_loop( 

164 self, 

165 peer_id: str, 

166 connect_func: callable, 

167 max_attempts: int 

168 ) -> None: 

169 """重连循环""" 

170 attempt = 0 

171 while attempt < max_attempts and self._running: 

172 delay = self.calculate_delay(attempt) 

173 logger.info(f"等待 {delay:.1f}秒后重连 {peer_id} (尝试 {attempt + 1}/{max_attempts})") 

174 

175 await asyncio.sleep(delay) 

176 

177 try: 

178 success = await connect_func(peer_id) 

179 if success: 

180 logger.info(f"重连成功: {peer_id}") 

181 return 

182 except Exception as e: 

183 logger.error(f"重连失败: {peer_id}, {e}") 

184 

185 attempt += 1 

186 

187 logger.warning(f"重连次数耗尽: {peer_id}") 

188 

189 def stop_all(self) -> None: 

190 """停止所有重连任务""" 

191 self._running = False 

192 for task in self._reconnect_tasks.values(): 

193 task.cancel() 

194 self._reconnect_tasks.clear() 

195 

196 

197class MessageQueue: 

198 """消息队列""" 

199 

200 def __init__(self, max_size: int = 1000): 

201 self.max_size = max_size 

202 self._queue: deque = deque(maxlen=max_size) 

203 self._pending: Dict[str, Message] = {} 

204 self._lock = asyncio.Lock() 

205 self._callbacks: Dict[str, callable] = {} 

206 

207 async def enqueue(self, msg_type: str, payload: dict) -> str: 

208 """入队消息""" 

209 async with self._lock: 

210 msg_id = str(uuid.uuid4()) 

211 msg = Message(msg_id=msg_id, msg_type=msg_type, payload=payload) 

212 self._queue.append(msg) 

213 self._pending[msg_id] = msg 

214 logger.debug(f"消息入队: {msg_id}, 类型: {msg_type}") 

215 return msg_id 

216 

217 async def dequeue(self) -> Optional[Message]: 

218 """出队消息""" 

219 async with self._lock: 

220 if self._queue: 

221 msg = self._queue.popleft() 

222 return msg 

223 return None 

224 

225 async def get_pending(self, msg_id: str) -> Optional[Message]: 

226 """获取待确认消息""" 

227 return self._pending.get(msg_id) 

228 

229 async def confirm(self, msg_id: str) -> bool: 

230 """确认消息已送达""" 

231 async with self._lock: 

232 if msg_id in self._pending: 

233 del self._pending[msg_id] 

234 logger.debug(f"消息已确认: {msg_id}") 

235 return True 

236 return False 

237 

238 async def retry(self, msg_id: str) -> bool: 

239 """重试消息""" 

240 async with self._lock: 

241 if msg_id in self._pending: 

242 msg = self._pending[msg_id] 

243 if msg.retry_count < msg.max_retries: 

244 msg.retry_count += 1 

245 self._queue.append(msg) 

246 logger.info(f"消息重试: {msg_id}, 次数: {msg.retry_count}") 

247 return True 

248 else: 

249 del self._pending[msg_id] 

250 logger.warning(f"消息重试次数耗尽: {msg_id}") 

251 return False 

252 

253 async def size(self) -> int: 

254 """获取队列大小""" 

255 return len(self._queue) 

256 

257 async def pending_count(self) -> int: 

258 """获取待确认消息数""" 

259 return len(self._pending) 

260 

261 

262class P2PNetworkManager: 

263 """P2P网络管理器（整合模块）""" 

264 

265 def __init__(self): 

266 self.connection_pool = ConnectionPool() 

267 self.reconnect_manager = ReconnectionManager() 

268 self.message_queue = MessageQueue() 

269 self._running = False 

270 

271 async def start(self) -> None: 

272 """启动P2P网络""" 

273 self._running = True 

274 logger.info("P2P网络管理器已启动") 

275 

276 async def stop(self) -> None: 

277 """停止P2P网络""" 

278 self._running = False 

279 self.reconnect_manager.stop_all() 

280 logger.info("P2P网络管理器已停止") 

281 

282 async def connect(self, peer_id: str, peer_name: str) -> bool: 

283 """连接到对等节点""" 

284 conn_info = await self.connection_pool.add_connection(peer_id, peer_name) 

285 

286 try: 

287 # 模拟连接建立 

288 await asyncio.sleep(0.1) 

289 await self.connection_pool.update_state(peer_id, ConnectionState.CONNECTED) 

290 conn_info.last_heartbeat = time.time() 

291 logger.info(f"已连接到: {peer_id}") 

292 return True 

293 except Exception as e: 

294 logger.error(f"连接失败: {peer_id}, {e}") 

295 await self.connection_pool.update_state(peer_id, ConnectionState.FAILED) 

296 return False 

297 

298 async def disconnect(self, peer_id: str) -> None: 

299 """断开连接""" 

300 self.reconnect_manager.stop_reconnect(peer_id) 

301 await self.connection_pool.remove_connection(peer_id) 

302 logger.info(f"已断开: {peer_id}") 

303 

304 async def send_message(self, peer_id: str, msg_type: str, payload: dict) -> Optional[str]: 

305 """发送消息""" 

306 conn_info = await self.connection_pool.get_connection(peer_id) 

307 if not conn_info or conn_info.state != ConnectionState.CONNECTED: 

308 # 消息入队等待发送 

309 return await self.message_queue.enqueue(msg_type, payload) 

310 

311 msg_id = await self.message_queue.enqueue(msg_type, payload) 

312 # 模拟发送 

313 await asyncio.sleep(0.01) 

314 await self.message_queue.confirm(msg_id) 

315 return msg_id 

316 

317 async def get_status(self) -> dict: 

318 """获取网络状态""" 

319 connections = await self.connection_pool.get_all_connections() 

320 return { 

321 "running": self._running, 

322 "connections": { 

323 peer_id: { 

324 "state": conn.state.value, 

325 "latency_ms": conn.latency_ms, 

326 "last_heartbeat": conn.last_heartbeat 

327 } 

328 for peer_id, conn in connections.items() 

329 }, 

330 "queue_size": await self.message_queue.size(), 

331 "pending_messages": await self.message_queue.pending_count() 

332 }