From 2cffdb65f20f27fade42c00dfdabd3dae9051b25 Mon Sep 17 00:00:00 2001 From: Marco Mazzucco Date: Wed, 2 Oct 2024 23:11:29 +0200 Subject: [PATCH 1/5] Decide if ipv6 can work --- locust/rpc/zmqrpc.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index bb344cf9dd..4a2a1eee9a 100644 --- a/locust/rpc/zmqrpc.py +++ b/locust/rpc/zmqrpc.py @@ -4,19 +4,19 @@ import msgpack.exceptions as msgerr import zmq.error as zmqerr import zmq.green as zmq -from urllib3.util.connection import HAS_IPV6 +from socket import getaddrinfo, IPPROTO_TCP, has_dualstack_ipv6 from .protocol import Message class BaseSocket: - def __init__(self, sock_type): + def __init__(self, sock_type, ipv4_only): context = zmq.Context() self.socket = context.socket(sock_type) self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1) self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30) - if HAS_IPV6: + if has_dualstack_ipv6() and not ipv4_only: self.socket.setsockopt(zmq.IPV6, 1) @retry() @@ -60,10 +60,15 @@ def recv_from_client(self): def close(self, linger=None): self.socket.close(linger=linger) + def ipv4_only(self, host, port) -> bool: + if str(getaddrinfo(host, port, proto=IPPROTO_TCP)).find("Family.AF_INET6") == -1: + return True + return False + class Server(BaseSocket): def __init__(self, host, port): - BaseSocket.__init__(self, zmq.ROUTER) + BaseSocket.__init__(self, zmq.ROUTER, self.ipv4_only(host, port)) if port == 0: self.port = self.socket.bind_to_random_port(f"tcp://{host}") else: @@ -76,6 +81,6 @@ def __init__(self, host, port): class Client(BaseSocket): def __init__(self, host, port, identity): - BaseSocket.__init__(self, zmq.DEALER) + BaseSocket.__init__(self, zmq.DEALER, self.ipv4_only(host, port)) self.socket.setsockopt(zmq.IDENTITY, identity.encode()) self.socket.connect("tcp://%s:%i" % (host, port)) From d478fedad455fd1f2ec7b49e7446fe0326da6e81 Mon Sep 17 00:00:00 2001 From: Marco Mazzucco Date: Fri, 4 Oct 2024 08:54:14 +0200 Subject: [PATCH 2/5] ruff fix --- locust/rpc/zmqrpc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index 4a2a1eee9a..f2e3b07f3f 100644 --- a/locust/rpc/zmqrpc.py +++ b/locust/rpc/zmqrpc.py @@ -1,10 +1,11 @@ from locust.exception import RPCError, RPCReceiveError, RPCSendError from locust.util.exception_handler import retry +from socket import IPPROTO_TCP, getaddrinfo, has_dualstack_ipv6 + import msgpack.exceptions as msgerr import zmq.error as zmqerr import zmq.green as zmq -from socket import getaddrinfo, IPPROTO_TCP, has_dualstack_ipv6 from .protocol import Message From f37caa835f0295e2ecaf8491aa64c0af3251cef3 Mon Sep 17 00:00:00 2001 From: Marco Mazzucco Date: Thu, 10 Oct 2024 23:16:56 +0200 Subject: [PATCH 3/5] got back to urllib3 for IPV6 check --- locust/rpc/zmqrpc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index f2e3b07f3f..cf2d21e432 100644 --- a/locust/rpc/zmqrpc.py +++ b/locust/rpc/zmqrpc.py @@ -1,11 +1,12 @@ from locust.exception import RPCError, RPCReceiveError, RPCSendError from locust.util.exception_handler import retry -from socket import IPPROTO_TCP, getaddrinfo, has_dualstack_ipv6 +from socket import IPPROTO_TCP, getaddrinfo import msgpack.exceptions as msgerr import zmq.error as zmqerr import zmq.green as zmq +from urllib3.util.connection import HAS_IPV6 from .protocol import Message @@ -17,7 +18,7 @@ def __init__(self, sock_type, ipv4_only): self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1) self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30) - if has_dualstack_ipv6() and not ipv4_only: + if HAS_IPV6 and not ipv4_only: self.socket.setsockopt(zmq.IPV6, 1) @retry() From b5c36f1ba265a748d517c934957340a0e93efaca Mon Sep 17 00:00:00 2001 From: Marco Mazzucco Date: Fri, 11 Oct 2024 16:29:04 +0200 Subject: [PATCH 4/5] Catch if resolution doesn't work --- locust/rpc/zmqrpc.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index cf2d21e432..a2ff1f9e04 100644 --- a/locust/rpc/zmqrpc.py +++ b/locust/rpc/zmqrpc.py @@ -1,12 +1,11 @@ from locust.exception import RPCError, RPCReceiveError, RPCSendError from locust.util.exception_handler import retry -from socket import IPPROTO_TCP, getaddrinfo +import socket as csocket import msgpack.exceptions as msgerr import zmq.error as zmqerr import zmq.green as zmq -from urllib3.util.connection import HAS_IPV6 from .protocol import Message @@ -18,7 +17,7 @@ def __init__(self, sock_type, ipv4_only): self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1) self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30) - if HAS_IPV6 and not ipv4_only: + if csocket.has_dualstack_ipv6() and not ipv4_only: self.socket.setsockopt(zmq.IPV6, 1) @retry() @@ -63,8 +62,12 @@ def close(self, linger=None): self.socket.close(linger=linger) def ipv4_only(self, host, port) -> bool: - if str(getaddrinfo(host, port, proto=IPPROTO_TCP)).find("Family.AF_INET6") == -1: - return True + try: + if str(csocket.getaddrinfo(host, port, proto=csocket.IPPROTO_TCP)).find("Family.AF_INET6") == -1: + return True + except csocket.gaierror as e: + print(f"Error resolving address: {e}") + return False return False From 39bc8ee82df1082a5a152e9bb671d621dcb347b1 Mon Sep 17 00:00:00 2001 From: Marco Mazzucco Date: Fri, 11 Oct 2024 22:31:13 +0200 Subject: [PATCH 5/5] adjusted import syntax --- locust/rpc/zmqrpc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index a2ff1f9e04..d5d5586129 100644 --- a/locust/rpc/zmqrpc.py +++ b/locust/rpc/zmqrpc.py @@ -2,6 +2,7 @@ from locust.util.exception_handler import retry import socket as csocket +from socket import gaierror, has_dualstack_ipv6 import msgpack.exceptions as msgerr import zmq.error as zmqerr @@ -17,7 +18,7 @@ def __init__(self, sock_type, ipv4_only): self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1) self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30) - if csocket.has_dualstack_ipv6() and not ipv4_only: + if has_dualstack_ipv6() and not ipv4_only: self.socket.setsockopt(zmq.IPV6, 1) @retry() @@ -65,7 +66,7 @@ def ipv4_only(self, host, port) -> bool: try: if str(csocket.getaddrinfo(host, port, proto=csocket.IPPROTO_TCP)).find("Family.AF_INET6") == -1: return True - except csocket.gaierror as e: + except gaierror as e: print(f"Error resolving address: {e}") return False return False