Merge pull request #1101 from albertowar/tcp_keep_alive

Adding TCP Keep Alive to guarantee master-slave communication after i…
locustio · Oct 1, 2019 · c95dda3 · c95dda3
2 parents 30f7928 + 3abafce
commit c95dda3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py
@@ -7,6 +7,9 @@ class BaseSocket(object):
     def __init__(self, sock_type):
         context = zmq.Context()
         self.socket = context.socket(sock_type)
+
+        self.socket.setsockopt(zmq.TCP_KEEPALIVE, 1)
+        self.socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 30)
 
     @retry()
     def send(self, msg):

diff --git a/locust/test/test_zmqrpc.py b/locust/test/test_zmqrpc.py
@@ -14,6 +14,12 @@ def tearDown(self):
         self.server.socket.close()
         self.client.socket.close()
 
+    def test_constructor(self):
+        self.assertEqual(self.server.socket.getsockopt(zmq.TCP_KEEPALIVE), 1)
+        self.assertEqual(self.server.socket.getsockopt(zmq.TCP_KEEPALIVE_IDLE), 30)
+        self.assertEqual(self.client.socket.getsockopt(zmq.TCP_KEEPALIVE), 1)
+        self.assertEqual(self.client.socket.getsockopt(zmq.TCP_KEEPALIVE_IDLE), 30)
+
     def test_client_send(self):
         self.client.send(Message('test', 'message', 'identity'))
         addr, msg = self.server.recv_from_client()