diff --git a/locust/runners.py b/locust/runners.py index eb1b5726cc..699e0103c6 100644 --- 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 index 4c7690edd0..25aeab6492 100644 --- 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):