Merge pull request #1324 from Trouv/worker-quit-ui-fix

Worker quitting then stopping via web UI bug fix

locust/runners.py

@@ -340,6 +340,10 @@ def hatching(self):
             @property
             def running(self):
                 return self.get_by_state(STATE_RUNNING)
+
+            @property
+            def missing(self):
+                return self.get_by_state(STATE_MISSING)
 
         self.clients = WorkerNodesDict()
         self.server = rpc.Server(master_bind_host, master_bind_port)
@@ -411,10 +415,11 @@ def start(self, locust_count, hatch_rate):
         self.state = STATE_HATCHING
 
     def stop(self):
-        self.state = STATE_STOPPING
-        for client in self.clients.all:
-            self.server.send_to_client(Message("stop", None, client.id))
-        self.environment.events.test_stop.fire(environment=self.environment)
+        if self.state not in [STATE_INIT, STATE_STOPPED, STATE_STOPPING]:
+            self.state = STATE_STOPPING
+            for client in self.clients.all:
+                self.server.send_to_client(Message("stop", None, client.id))
+            self.environment.events.test_stop.fire(environment=self.environment)
 
     def quit(self):
         if self.state not in [STATE_INIT, STATE_STOPPED, STATE_STOPPING]:
@@ -425,18 +430,28 @@ def quit(self):
             self.server.send_to_client(Message("quit", None, client.id))
         gevent.sleep(0.5) # wait for final stats report from all workers
         self.greenlet.kill(block=True)
+
+    def check_stopped(self):
+        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
+
 
     def heartbeat_worker(self):
         while True:
             gevent.sleep(HEARTBEAT_INTERVAL)
             if self.connection_broken:
                 self.reset_connection()
                 continue
+
             for client in self.clients.all:
                 if client.heartbeat < 0 and client.state != STATE_MISSING:
                     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:
+                        logger.info("The last worker went missing, stopping test.")
+                        self.stop()
+                        self.check_stopped()
                 else:
                     client.heartbeat -= 1
 
@@ -496,11 +511,14 @@ def client_listener(self):
                 if msg.node_id in self.clients:
                     del self.clients[msg.node_id]
                     logger.info("Client %r quit. Currently %i clients connected." % (msg.node_id, len(self.clients.ready)))
+                    if self.worker_count - len(self.clients.missing) <= 0:
+                        logger.info("The last worker quit, stopping test.")
+                        self.stop()
             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_RUNNING and x.state != STATE_HATCHING, self.clients.all)):
-                self.state = STATE_STOPPED
+            self.check_stopped()
+
 
     @property
     def worker_count(self):

locust/static/locust.js

@@ -4,14 +4,18 @@ $(window).ready(function() {
     }
 });
 
-$("#box_stop a.stop-button").click(function(event) {
-    event.preventDefault();
-    $.get($(this).attr("href"));
-    $("body").attr("class", "stopped");
+function appearStopped() {
     $(".box_stop").hide();
     $("a.new_test").show();
     $("a.edit_test").hide();
     $(".user_count").hide();
+}
+
+$("#box_stop a.stop-button").click(function(event) {
+    event.preventDefault();
+    $.get($(this).attr("href"));
+    $("body").attr("class", "stopped");
+    appearStopped()
 });
 
 $("#box_stop a.reset-button").click(function(event) {
@@ -173,6 +177,8 @@ function updateStats() {
             rpsChart.addValue([total.current_rps, total.current_fail_per_sec]);
             responseTimeChart.addValue([report.current_response_time_percentile_50, report.current_response_time_percentile_95]);
             usersChart.addValue([report.user_count]);
+        } else {
+            appearStopped();
         }
 
         setTimeout(updateStats, 2000);
@@ -187,4 +193,4 @@ function updateExceptions() {
         setTimeout(updateExceptions, 5000);
     });
 }
-updateExceptions();
+updateExceptions();

locust/static/style.css

@@ -139,8 +139,8 @@ a:hover {
 }
 .hatching .boxes .box_running {display: block;}
 .running .boxes .box_running {display: block;}
-.stopped .boxes .box_running {display: block;}
-.stopped .boxes .box_stop {display: none;}
+.stopped .boxes .box_running, .stopping .boxes .box_running {display: block;}
+.stopped .boxes .box_stop, .stopping .boxes .box_stop {display: none;}
 
 .container {
     max-width: 1800px;
@@ -206,7 +206,7 @@ a:hover {
 }
 
 
-.stopped .start {
+.stopped .start, .stopping .start {
     display: none;
     border-radius: 5px;
     -moz-border-radius: 5px;
@@ -216,7 +216,7 @@ a:hover {
     box-shadow: 0 0 60px rgba(0,0,0,0.3);
 }
 
-.stopped .edit {display: none;}
+.stopped .edit, .stopping .edit {display: none;}
 .running .edit, .hatching .edit {
     display: none;
     border-radius: 5px;
@@ -232,25 +232,25 @@ a:hover {
 .ready .start {display: block;}
 
 .running .status, .hatching .status {display: block;}
-.stopped .status {display: block;}
+.stopped .status, .stopping .status {display: block;}
 .ready .status {display: none;}
 
-.stopped .boxes .edit_test, .ready .boxes .edit_test {display: none;}
-.stopped .boxes .user_count, .ready .boxes .user_count {display: none;}
+.stopped .boxes .edit_test, .stopping .boxes .edit_test, .ready .boxes .edit_test {display: none;}
+.stopped .boxes .user_count, .stopping .boxes .user_count, .ready .boxes .user_count {display: none;}
 
 .running a.new_test, .ready a.new_test, .hatching a.new_test {display: none;}
 .running a.new_test {display: none;}
-.stopped a.new_test {display: block;}
+.stopped a.new_test, .stopping a.new_test {display: block;}
 
 .start a.close_link, .edit a.close_link{
     position: absolute;
     right: 10px;
     top: 10px;
 }
-.stopped .start a.close_link {display: inline;}
+.stopped .start a.close_link, .stopping .start a.close_link {display: inline;}
 .running .start a.close_link, .ready .start a.close_link, .hatching .start a.close_link  {display: none;}
 
-.stopped .edit a.close_link, .ready .edit a.close_link {display: none;}
+.stopped .edit a.close_link, .stopping .edit a.close_link, .ready .edit a.close_link {display: none;}
 .running .edit a.close_link, .hatching .edit a.close_link {display: inline;}
 
 .stats_label {
@@ -461,7 +461,7 @@ ul.tabs li a.current:after {
 }
 
 .running .hostname, .hatching .hostname {display: block;}
-.stopped .hostname {display: block;}
+.stopped .hostname, .stopping .hostname {display: block;}
 .ready .hostname {display: none;}
 
 .footer {

locust/test/test_runners.py

@@ -12,7 +12,7 @@
 from locust.exception import LocustError, RPCError, StopLocust
 from locust.rpc import Message
 from locust.runners import LocustRunner, LocalLocustRunner, MasterLocustRunner, WorkerNode, \
-     WorkerLocustRunner, STATE_INIT, STATE_HATCHING, STATE_RUNNING, STATE_MISSING
+     WorkerLocustRunner, STATE_INIT, STATE_HATCHING, STATE_RUNNING, STATE_MISSING, STATE_STOPPED
 from locust.stats import RequestStats
 from locust.test.testcases import LocustTestCase
 
@@ -452,6 +452,48 @@ def test_master_marks_downed_workers_as_missing(self):
             # print(master.clients['fake_client'].__dict__)
             assert master.clients['fake_client'].state == STATE_MISSING
 
+    def test_last_worker_quitting_stops_test(self):
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner()
+            server.mocked_send(Message("client_ready", None, "fake_client1"))
+            server.mocked_send(Message("client_ready", None, "fake_client2"))
+
+            master.start(1, 2)
+            server.mocked_send(Message("hatching", None, "fake_client1"))
+            server.mocked_send(Message("hatching", None, "fake_client2"))
+
+            server.mocked_send(Message("quit", None, "fake_client1"))
+            sleep(0)
+            self.assertEqual(1, len(master.clients.all))
+            self.assertNotEqual(STATE_STOPPED, master.state, "Not all workers quit but test stopped anyway.")
+
+            server.mocked_send(Message("quit", None, "fake_client2"))
+            sleep(0)
+            self.assertEqual(0, len(master.clients.all))
+            self.assertEqual(STATE_STOPPED, master.state, "All workers quit but test didn't stop.")
+
+    @mock.patch("locust.runners.HEARTBEAT_INTERVAL", new=0.1)
+    def test_last_worker_missing_stops_test(self):
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            master = self.get_runner()
+            server.mocked_send(Message("client_ready", None, "fake_client1"))
+            server.mocked_send(Message("client_ready", None, "fake_client2"))
+
+            master.start(1, 2)
+            server.mocked_send(Message("hatching", None, "fake_client1"))
+            server.mocked_send(Message("hatching", None, "fake_client2"))
+
+            sleep(0.3)
+            server.mocked_send(Message("heartbeat", {'state': STATE_RUNNING, 'current_cpu_usage': 50}, "fake_client1"))
+
+            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.")
+
+            sleep(0.3)
+            self.assertEqual(2, len(master.clients.missing))
+            self.assertEqual(STATE_STOPPED, master.state, "All workers went missing but test didn't stop.")
+
     def test_master_total_stats(self):
         with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
             master = self.get_runner()