Merge pull request #1710 from eNote-GmbH/feature/worker-health

Fix automatic distributed test shutdown
locustio · Feb 26, 2021 · e17032b · e17032b
2 parents 4480674 + bfa2084
commit e17032b
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 10 deletions.
diff --git a/locust/runners.py b/locust/runners.py
@@ -608,7 +608,7 @@ def heartbeat_worker(self):
                     logger.info("Worker %s failed to send heartbeat, setting state to missing." % str(client.id))
                     client.state = STATE_MISSING
                     client.user_count = 0
-                    if self.worker_count - len(self.clients.missing) <= 0:
+                    if self.worker_count <= 0:
                         logger.info("The last worker went missing, stopping test.")
                         self.stop()
                         self.check_stopped()
@@ -654,7 +654,15 @@ def client_listener(self):
                 if msg.node_id in self.clients:
                     c = self.clients[msg.node_id]
                     c.heartbeat = HEARTBEAT_LIVENESS
-                    c.state = msg.data["state"]
+                    client_state = msg.data["state"]
+                    if c.state == STATE_MISSING:
+                        logger.info(
+                            "Worker %s self-healed with heartbeat, setting state to %s." % (str(c.id), client_state)
+                        )
+                        user_count = msg.data.get("count")
+                        if user_count:
+                            c.user_count = user_count
+                    c.state = client_state
                     c.cpu_usage = msg.data["current_cpu_usage"]
                     if not c.cpu_warning_emitted and c.cpu_usage > 90:
                         self.worker_cpu_warning_emitted = True  # used to fail the test in the end
@@ -751,7 +759,11 @@ def heartbeat(self):
                 self.client.send(
                     Message(
                         "heartbeat",
-                        {"state": self.worker_state, "current_cpu_usage": self.current_cpu_usage},
+                        {
+                            "state": self.worker_state,
+                            "current_cpu_usage": self.current_cpu_usage,
+                            "count": self.user_count,
+                        },
                         self.client_id,
                     )
                 )

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
@@ -22,6 +22,7 @@
     STATE_SPAWNING,
     STATE_RUNNING,
     STATE_MISSING,
+    STATE_STOPPING,
     STATE_STOPPED,
 )
 from locust.stats import RequestStats
@@ -754,20 +755,45 @@ def test_last_worker_missing_stops_test(self):
             master = self.get_runner()
             server.mocked_send(Message("client_ready", None, "fake_client1"))
             server.mocked_send(Message("client_ready", None, "fake_client2"))
+            server.mocked_send(Message("client_ready", None, "fake_client3"))
 
-            master.start(1, 2)
+            master.start(3, 3)
             server.mocked_send(Message("spawning", None, "fake_client1"))
             server.mocked_send(Message("spawning", None, "fake_client2"))
+            server.mocked_send(Message("spawning", None, "fake_client3"))
+
+            sleep(0.2)
+            server.mocked_send(
+                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client1")
+            )
+            server.mocked_send(
+                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client2")
+            )
+            server.mocked_send(
+                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client3")
+            )
 
-            sleep(0.3)
-            server.mocked_send(Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50}, "fake_client1"))
+            sleep(0.2)
+            self.assertEqual(0, len(master.clients.missing))
+            self.assertEqual(3, master.worker_count)
+            self.assertNotIn(
+                master.state, [STATE_STOPPED, STATE_STOPPING], "Not all workers went missing but test stopped anyway."
+            )
 
-            sleep(0.3)
-            self.assertEqual(1, len(master.clients.missing))
-            self.assertNotEqual(STATE_STOPPED, master.state, "Not all workers went missing but test stopped anyway.")
+            server.mocked_send(
+                Message("heartbeat", {"state": STATE_RUNNING, "current_cpu_usage": 50, "count": 1}, "fake_client1")
+            )
 
-            sleep(0.3)
+            sleep(0.4)
             self.assertEqual(2, len(master.clients.missing))
+            self.assertEqual(1, master.worker_count)
+            self.assertNotIn(
+                master.state, [STATE_STOPPED, STATE_STOPPING], "Not all workers went missing but test stopped anyway."
+            )
+
+            sleep(0.2)
+            self.assertEqual(3, len(master.clients.missing))
+            self.assertEqual(0, master.worker_count)
             self.assertEqual(STATE_STOPPED, master.state, "All workers went missing but test didn't stop.")
 
     def test_master_total_stats(self):