add retry in zmqrpc

locust/rpc/zmqrpc.py

@@ -1,24 +1,28 @@
 import zmq.green as zmq
 
 from .protocol import Message
-
+from locust.util.exception_handler import retry
 
 class BaseSocket(object):
     def __init__(self, sock_type):
         context = zmq.Context()
         self.socket = context.socket(sock_type)
-
+
+    @retry()
     def send(self, msg):
         self.socket.send(msg.serialize())
 
+    @retry()
     def send_to_client(self, msg):
         self.socket.send_multipart([msg.node_id.encode(), msg.serialize()])
 
+    @retry()
     def recv(self):
         data = self.socket.recv()
         msg = Message.unserialize(data)
         return msg
 
+    @retry()
     def recv_from_client(self):
         data = self.socket.recv_multipart()
         addr = data[0]

locust/util/exception_handler.py

@@ -0,0 +1,23 @@
+import logging
+from time import sleep
+
+logger = logging.getLogger(__name__)
+
+def retry(delays=(1, 3, 5),
+          exception=Exception):
+    def decorator(function):
+        def wrapper(*args, **kwargs):
+            cnt = 0
+            for delay in delays + (None,):
+                try:
+                    return function(*args, **kwargs)
+                except exception:
+                    if delay is None:
+                        logger.info("Retry failed after %d times." % ( cnt ) )
+                        raise
+                    else:
+                        cnt += 1
+                        logger.info("Exception found on retry %d: -- retry after %ds" % ( cnt, delay ) )
+                        sleep(delay)
+        return wrapper
+    return decorator