diff --git a/locust/rpc/zmqrpc.py b/locust/rpc/zmqrpc.py index 164d3cb262..b891b3330d 100644 --- a/locust/rpc/zmqrpc.py +++ b/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] diff --git a/locust/util/exception_handler.py b/locust/util/exception_handler.py new file mode 100644 index 0000000000..ddb01877d8 --- /dev/null +++ b/locust/util/exception_handler.py @@ -0,0 +1,23 @@ +import time +import logging + +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 ) ) + time.sleep(delay) + return wrapper + return decorator \ No newline at end of file