Merge pull request #2923 from nc-marco/feature/ipv4_only_routing

Decide if ipv6 can work
locustio · Oct 11, 2024 · 4d12906 · 4d12906
2 parents 1954b26 + 39bc8ee
commit 4d12906
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py
@@ -1,22 +1,24 @@
 from locust.exception import RPCError, RPCReceiveError, RPCSendError
 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
 import zmq.green as zmq
-from urllib3.util.connection import HAS_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 +62,19 @@ def recv_from_client(self):
     def close(self, linger=None):
         self.socket.close(linger=linger)
 
+    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 gaierror as e:
+            print(f"Error resolving address: {e}")
+            return False
+        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 +87,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))