Horizontal42/netsleuth
1
0
Fork 0
mirror of https://github.com/Horizontal42/netsleuth.git synced 2026-11-10 23:52:04 +03:00
Code Issues Projects Releases Packages Wiki Activity Actions

⚡ fix: refactor ping_samples_win to be asynchronous to avoid thread pool blocking

Browse source 
By converting `ping_samples_win` from a synchronous function that blocks with `time.sleep` into an async function using `await asyncio.sleep()`, we prevent thread-pool exhaustion during heavy concurrent ICMP sweeps. Only the strictly blocking Windows ICMP API call (`echo_once`) is delegated to a background thread.

Co-authored-by: Horizontal42 <75383519+Horizontal42@users.noreply.github.com>
This commit is contained in:
google-labs-jules[bot] 2026-09-04 07:23:48 +00:00
parent 8c99ac78b7
commit dc07077abd
3 changed files with 71 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
src/netsleuth/probes/icmp_win.py
View file

@ -48,9 +48,12 @@ def parse_echo_reply(buffer: bytes, pointer_size: int = 8) -> IcmpReply:
ttl = buffer[ttl_offset]
if status in (IP_REQ_TIMED_OUT,) or (a, b, c, d) == (0, 0, 0, 0):
return IcmpReply(address=None, status=status, rtt_ms=None, ttl=ttl)
return IcmpReply(address=f"{a}.{b}.{c}.{d}", status=status, rtt_ms=float(rtt), ttl=ttl)
return IcmpReply(
address=f"{a}.{b}.{c}.{d}", status=status, rtt_ms=float(rtt), ttl=ttl
)
import asyncio
import ctypes
import platform
import socket
@ -139,7 +142,11 @@ def echo_once(
iphlpapi, handle = _handle()
try:
options = _IpOptionInformation(
Ttl=ttl, Tos=0, Flags=IP_FLAG_DF if df else 0, OptionsSize=0, OptionsData=None
Ttl=ttl,
Tos=0,
Flags=IP_FLAG_DF if df else 0,
OptionsSize=0,
OptionsData=None,
)
buffer = ctypes.create_string_buffer(_REPLY_BUFFER_SIZE)
address = struct.unpack("<I", socket.inet_aton(dest))[0]
@ -180,7 +187,7 @@ def echo_once(
iphlpapi.IcmpCloseHandle(handle)
def ping_samples_win(
async def ping_samples_win(
host: str, count: int, interval: float, timeout: float, source_ip: str | None = None
) -> list[float | None]:
dest = socket.gethostbyname(host)
@ -188,8 +195,10 @@ def ping_samples_win(
samples: list[float | None] = []
for index in range(count):
if index:
time.sleep(interval)
reply = echo_once(dest, ttl=128, timeout_ms=timeout_ms, source_ip=source_ip)
await asyncio.sleep(interval)
reply = await asyncio.to_thread(
echo_once, dest, ttl=128, timeout_ms=timeout_ms, source_ip=source_ip
)
samples.append(reply.rtt_ms if classify_status(reply.status) == "ok" else None)
return samples
@ -207,12 +216,16 @@ def trace_hops_win(
# into 30 minutes of blocking IcmpSendEcho2 calls.
address = socket.gethostbyname(dest)
hops: list[tuple[int, IcmpReply]] = []
per_hop_timeout_ms = max(1, min(timeout_ms // max(max_hops, 1), _MAX_HOP_TIMEOUT_MS))
per_hop_timeout_ms = max(
1, min(timeout_ms // max(max_hops, 1), _MAX_HOP_TIMEOUT_MS)
)
deadline = time.monotonic() + timeout_ms / 1000.0
for ttl in range(1, max_hops + 1):
if time.monotonic() >= deadline:
break
reply = echo_once(address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip)
reply = echo_once(
address, ttl=ttl, timeout_ms=per_hop_timeout_ms, source_ip=source_ip
)
hops.append((ttl, reply))
if classify_status(reply.status) == "ok":
break

36
src/netsleuth/probes/latency.py
View file

@ -47,7 +47,7 @@ async def tcp_connect_rtt(
_reader, writer = await asyncio.wait_for(
asyncio.open_connection(host, port, local_addr=local_addr), timeout=timeout
)
except (TimeoutError, OSError):
except TimeoutError, OSError:
return None
writer.close()
try:
@ -62,22 +62,32 @@ async def _resolve(host: str) -> str | None:
loop = asyncio.get_running_loop()
res = await loop.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
return res[0][4][0]
except (OSError, IndexError):
except OSError, IndexError:
return None
async def _icmp_samples(
host: str, count: int, interval: float, timeout: float, backend: str, source_ip: str | None = None
host: str,
count: int,
interval: float,
timeout: float,
backend: str,
source_ip: str | None = None,
) -> list[float | None]:
if backend == "icmp_win":
from netsleuth.probes.icmp_win import ping_samples_win
return await asyncio.to_thread(ping_samples_win, host, count, interval, timeout, source_ip)
return await ping_samples_win(host, count, interval, timeout, source_ip)
from icmplib import async_ping
privileged = backend == "icmp_raw"
host_result = await async_ping(
host, count=count, interval=interval, timeout=timeout, privileged=privileged, source=source_ip
host,
count=count,
interval=interval,
timeout=timeout,
privileged=privileged,
source=source_ip,
)
samples: list[float | None] = list(host_result.rtts)
samples.extend([None] * (count - len(samples)))
@ -96,11 +106,15 @@ async def ping_host(
) -> PingResult:
if semaphore is not None:
async with semaphore:
return await ping_host(host, label, count, interval, timeout, backend, source_ip)
return await ping_host(
host, label, count, interval, timeout, backend, source_ip
)
resolved = await _resolve(host)
if backend in ("icmp_win", "icmp_dgram", "icmp_raw"):
try:
samples = await _icmp_samples(host, count, interval, timeout, backend, source_ip)
samples = await _icmp_samples(
host, count, interval, timeout, backend, source_ip
)
return summarize_ping(label, host, resolved, backend, samples)
except Exception:
backend = "tcp"
@ -108,7 +122,9 @@ async def ping_host(
for index in range(count):
if index:
await asyncio.sleep(interval)
samples.append(await tcp_connect_rtt(host, timeout=timeout, source_ip=source_ip))
samples.append(
await tcp_connect_rtt(host, timeout=timeout, source_ip=source_ip)
)
return summarize_ping(label, host, resolved, "tcp", samples)
@ -125,7 +141,9 @@ async def ping_fanout(
return list(
await asyncio.gather(
*(
ping_host(host, label, count, interval, timeout, backend, source_ip, semaphore)
ping_host(
host, label, count, interval, timeout, backend, source_ip, semaphore
)
for label, host in hosts
)
)

40
tests/test_latency.py
View file

@ -6,7 +6,12 @@ from unittest.mock import AsyncMock, patch, MagicMock
import pytest
from netsleuth.models import PingResult
from netsleuth.probes.latency import summarize_ping, tcp_connect_rtt, _resolve, _icmp_samples
from netsleuth.probes.latency import (
summarize_ping,
tcp_connect_rtt,
_resolve,
_icmp_samples,
)
def test_summarize_ping_calculates_correct_stats():
@ -68,6 +73,7 @@ async def test_tcp_connect_rtt_oserror(mock_open_conn):
rtt = await tcp_connect_rtt("host.com")
assert rtt is None
@pytest.mark.asyncio
async def test_resolve_success():
# Mock loop.getaddrinfo via asyncio.get_running_loop
@ -75,9 +81,7 @@ async def test_resolve_success():
with patch.object(loop, "getaddrinfo", new_callable=AsyncMock) as mock_getaddrinfo:
# getaddrinfo returns list of tuples: (family, type, proto, canonname, sockaddr)
# sockaddr is (address, port) for IPv4
mock_getaddrinfo.return_value = [
(2, 1, 6, '', ('1.2.3.4', 0))
]
mock_getaddrinfo.return_value = [(2, 1, 6, "", ("1.2.3.4", 0))]
ip = await _resolve("host.com")
assert ip == "1.2.3.4"
@ -92,17 +96,12 @@ async def test_resolve_failure():
@pytest.mark.asyncio
@patch("netsleuth.probes.latency.asyncio.to_thread", new_callable=AsyncMock)
async def test_icmp_samples_win(mock_to_thread):
mock_to_thread.return_value = [10.0, 15.0]
@patch("netsleuth.probes.icmp_win.ping_samples_win", new_callable=AsyncMock)
async def test_icmp_samples_win(mock_ping_samples_win):
mock_ping_samples_win.return_value = [10.0, 15.0]
samples = await _icmp_samples("host.com", 2, 0.5, 2.0, "icmp_win", source_ip=None)
assert samples == [10.0, 15.0]
mock_to_thread.assert_called_once()
# The first argument should be the ping_samples_win function
func_arg = mock_to_thread.call_args[0][0]
assert func_arg.__name__ == "ping_samples_win"
# The remaining arguments
assert mock_to_thread.call_args[0][1:] == ("host.com", 2, 0.5, 2.0, None)
mock_ping_samples_win.assert_called_once_with("host.com", 2, 0.5, 2.0, None)
@pytest.mark.asyncio
@ -133,7 +132,12 @@ async def test_icmp_samples_unix_dgram_padded():
# Should be padded with Nones to length 3
assert samples == [10.0, None, None]
mock_async_ping.assert_called_once_with(
"host.com", count=3, interval=0.5, timeout=2.0, privileged=False, source=None
"host.com",
count=3,
interval=0.5,
timeout=2.0,
privileged=False,
source=None,
)
@ -160,7 +164,9 @@ async def test_ping_host_icmp_success(mock_icmp_samples, mock_resolve):
@patch("netsleuth.probes.latency._resolve", new_callable=AsyncMock)
@patch("netsleuth.probes.latency._icmp_samples", new_callable=AsyncMock)
@patch("netsleuth.probes.latency.tcp_connect_rtt", new_callable=AsyncMock)
async def test_ping_host_icmp_exception_fallback(mock_tcp_connect, mock_icmp_samples, mock_resolve):
async def test_ping_host_icmp_exception_fallback(
mock_tcp_connect, mock_icmp_samples, mock_resolve
):
mock_resolve.return_value = "1.2.3.4"
mock_icmp_samples.side_effect = Exception("ICMP failed")
mock_tcp_connect.side_effect = [15.0, 16.0]
@ -190,7 +196,9 @@ async def test_ping_host_tcp_direct(mock_tcp_connect, mock_resolve):
async def test_ping_host_semaphore():
semaphore = asyncio.Semaphore(1)
with patch("netsleuth.probes.latency.ping_host", new_callable=AsyncMock) as mock_ping_host:
with patch(
"netsleuth.probes.latency.ping_host", new_callable=AsyncMock
) as mock_ping_host:
mock_ping_host.return_value = MagicMock()
# When called with a semaphore, ping_host calls itself without the semaphore inside the block