perf: traceroute

apps/settings/tools/traceroute.py

@@ -1,12 +1,17 @@
 # -*- coding: utf-8 -*-
 #
 import asyncio
+import select
 import socket
 import struct
 import time
 
+from common.utils import get_logger
 from settings.utils import generate_ips
 
+logger = get_logger(__name__)
+ICMP_PROTO = getattr(socket, 'IPPROTO_ICMP', 1)
+
 
 def calculate_checksum(data):
     # 计算 ICMP 校验和
@@ -21,40 +26,45 @@ def calculate_checksum(data):
     return checksum
 
 
+def recv_with_timeout(sock_fd, timeout):
+    ready, _, _ = select.select([sock_fd], [], [], timeout)
+    if not ready:
+        raise TimeoutError(f'recv timeout after {timeout}s')
+    return sock_fd.recvfrom(1024)
+
+
 async def once_traceroute(target, display, max_hops=30, timeout=3):
-    loop = asyncio.get_event_loop()
-    icmp = socket.getprotobyname('icmp')
-    sock_fd = socket.socket(socket.AF_INET, socket.SOCK_RAW, icmp)
-    for ttl in range(1, max_hops+1):
-        # 设置 TTL
-        sock_fd.setsockopt(socket.IPPROTO_IP, socket.IP_TTL, ttl)
-        # 设置超时时间
-        sock_fd.settimeout(timeout)
-        # 发送 ICMP Echo 请求数据包
-        packet = struct.pack('!BBHHH', 8, 0, 0, 12345, ttl)
-        checksum = calculate_checksum(packet)
-        packet = struct.pack('!BBHHH', 8, 0, checksum, 12345, ttl)
-        start_time = time.time()
-        sock_fd.sendto(packet, (target, 0))
-        await asyncio.sleep(0.01)
-        try:
-            # 接收 ICMP Echo Reply 数据包
-            recv_packet, addr = await loop.sock_recvfrom(sock_fd, 1024)
-            end_time = time.time()
-            # 解析目标地址
-            dest_ip = addr[0]
-            # 获取跃点信息
-            hop_info = f'{ttl} {dest_ip} ({dest_ip})'
-            # 获取延迟时间信息
-            delay_info = f'{(end_time - start_time) * 1000:.3f} ms'
-            # 打印跃点信息和延迟时间
-            await display(f'{hop_info:<4} {delay_info}')
-            if dest_ip == target:
-                return
-        except socket.timeout:
-            # 发生超时，跳出循环
-            await display(f'{ttl} *')
-    sock_fd.close()
+    sock_fd = socket.socket(socket.AF_INET, socket.SOCK_RAW, ICMP_PROTO)
+    try:
+        for ttl in range(1, max_hops+1):
+            # 设置 TTL
+            sock_fd.setsockopt(socket.IPPROTO_IP, socket.IP_TTL, ttl)
+            # 发送 ICMP Echo 请求数据包
+            packet = struct.pack('!BBHHH', 8, 0, 0, 12345, ttl)
+            checksum = calculate_checksum(packet)
+            packet = struct.pack('!BBHHH', 8, 0, checksum, 12345, ttl)
+            start_time = time.monotonic()
+            sock_fd.sendto(packet, (target, 0))
+            await asyncio.sleep(0.01)
+            try:
+                # 使用 select 等待 raw socket 可读，避免与 loop.sock_recvfrom 混用
+                recv_packet, addr = await asyncio.to_thread(recv_with_timeout, sock_fd, timeout)
+                end_time = time.monotonic()
+                # 解析目标地址
+                dest_ip = addr[0]
+                # 获取跃点信息
+                hop_info = f'{ttl} {dest_ip} ({dest_ip})'
+                # 获取延迟时间信息
+                delay_info = f'{(end_time - start_time) * 1000:.3f} ms'
+                # 打印跃点信息和延迟时间
+                await display(f'{hop_info:<4} {delay_info}')
+                if dest_ip == target:
+                    return
+            except (TimeoutError, socket.timeout):
+                # 发生超时，跳出循环
+                await display(f'{ttl} *')
+    finally:
+        sock_fd.close()
 
 
 async def verbose_traceroute(dest_ips, timeout=10, display=None):
@@ -69,5 +79,6 @@ async def verbose_traceroute(dest_ips, timeout=10, display=None):
         try:
             await once_traceroute(dest_ip, display)
         except Exception as e:
+            logger.exception('Traceroute failed for %s', dest_ip)
             msg = f'Error: {e}'
         await display(msg)