From d5b310c646806bf3f501bf0f8700f31fc09ca1af Mon Sep 17 00:00:00 2001
From: delulu <deluzhu@yahoo.com>
Date: Mon, 18 Feb 2019 16:07:45 +0800
Subject: [PATCH] fix inconsistent status when slave status is changed

---
 locust/main.py    |  2 ++
 locust/runners.py | 11 ++++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/locust/main.py b/locust/main.py
index f9a02c91d5..fe3186d861 100644
--- a/locust/main.py
+++ b/locust/main.py
@@ -132,6 +132,7 @@ def parse_options():
         '--heartbeat-liveness',
         action='store',
         type='int',
+        dest='heartbeat_liveness',
         default=3,
         help="set number of seconds before failed heartbeat from slave"
     )
@@ -140,6 +141,7 @@ def parse_options():
         '--heartbeat-interval',
         action='store',
         type='int',
+        dest='heartbeat_interval',
         default=1,
         help="set number of seconds delay between slave heartbeats to master"
     )
diff --git a/locust/runners.py b/locust/runners.py
index f23d62fb9e..a444d722c7 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -283,6 +283,7 @@ def start_hatching(self, locust_count, hatch_rate):
             return
 
         self.num_clients = locust_count
+        self.hatch_rate = hatch_rate
         slave_num_clients = locust_count // (num_slaves or 1)
         slave_hatch_rate = float(hatch_rate) / (num_slaves or 1)
         remaining = locust_count % num_slaves
@@ -328,6 +329,7 @@ def heartbeat_worker(self):
                 if client.heartbeat < 0 and client.state != STATE_MISSING:
                     logger.info('Slave %s failed to send heartbeat, setting state to missing.' % str(client.id))
                     client.state = STATE_MISSING
+                    client.user_count = 0
                 else:
                     client.heartbeat -= 1
 
@@ -338,7 +340,10 @@ def client_listener(self):
             if msg.type == "client_ready":
                 id = msg.node_id
                 self.clients[id] = SlaveNode(id, heartbeat_liveness=self.heartbeat_liveness)
-                logger.info("Client %r reported as ready. Currently %i clients ready to swarm." % (id, len(self.clients.ready)))
+                logger.info("Client %r reported as ready. Currently %i clients ready to swarm." % (id, len(self.clients.ready + self.clients.running + self.clients.hatching)))
+                # balance the load distribution when new client joins
+                if self.state == STATE_RUNNING or self.state == STATE_HATCHING:
+                    self.start_hatching(self.user_count, self.hatch_rate)
                 ## emit a warning if the slave's clock seem to be out of sync with our clock
                 #if abs(time() - msg.data["time"]) > 5.0:
                 #    warnings.warn("The slave node's clock seem to be out of sync. For the statistics to be correct the different locust servers need to have synchronized clocks.")
@@ -366,8 +371,8 @@ def client_listener(self):
             elif msg.type == "exception":
                 self.log_exception(msg.node_id, msg.data["msg"], msg.data["traceback"])
 
-            if not self.state == STATE_INIT and all(map(lambda x: x.state == STATE_INIT, self.clients.all)):
-                    self.state = STATE_STOPPED
+            if not self.state == STATE_INIT and all(map(lambda x: x.state != STATE_RUNNING and x.state != STATE_HATCHING, self.clients.all)):
+                self.state = STATE_STOPPED
 
     @property
     def slave_count(self):