Merge pull request #1127 from cyberw/fix-stop-timeout

Fix and add test for when locusts fail to exit at end of iteration during stop timeout.
locustio · Oct 31, 2019 · dd549f6 · dd549f6
2 parents cbf2834 + b5c808b
commit dd549f6
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 5 deletions.
diff --git a/locust/core.py b/locust/core.py
@@ -358,15 +358,18 @@ def run(self, *args, **kwargs):
 
         while (True):
             try:
-                if self.locust._state == LOCUST_STATE_STOPPING:
-                    raise GreenletExit()
-
                 if not self._task_queue:
                     self.schedule_task(self.get_next_task())
 
                 try:
+                    if self.locust._state == LOCUST_STATE_STOPPING:
+                        raise GreenletExit()
                     self.execute_next_task()
+                    if self.locust._state == LOCUST_STATE_STOPPING:
+                        raise GreenletExit()
                 except RescheduleTaskImmediately:
+                    if self.locust._state == LOCUST_STATE_STOPPING:
+                        raise GreenletExit()
                     pass
                 except RescheduleTask:
                     self.wait()

diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
@@ -502,6 +502,8 @@ def my_task(self):
 
         class MyTestLocust(Locust):
             task_set = MyTaskSet
+            min_wait = 0
+            max_wait = 0
 
         options = mocked_options()
         runner = LocalLocustRunner([MyTestLocust], options)
@@ -517,13 +519,49 @@ class MyTestLocust(Locust):
         runner.quit()
         self.assertEqual("second", MyTaskSet.state)
 
-        options.stop_timeout = short_time * 2 # allow task iteration to complete, with some margin
+        options.stop_timeout = short_time * 3 # allow task iteration to complete, with some margin
         runner = LocalLocustRunner([MyTestLocust], options)
         runner.start_hatching(1, 1)
         gevent.sleep(short_time)
-        runner.quit()
+        timeout = gevent.Timeout(short_time * 2)
+        timeout.start()
+        try:
+            runner.quit()
+            runner.greenlet.join()
+        except gevent.Timeout:
+            self.fail("Got Timeout exception. Some locusts must have kept runnining after iteration finish")
+        finally:
+            timeout.cancel()
         self.assertEqual("third", MyTaskSet.state)
 
+    def test_stop_timeout_during_on_start(self):
+        short_time = 0.05
+        class MyTaskSet(TaskSet):
+            finished_on_start = False
+            my_task_run = False
+            def on_start(self):
+                gevent.sleep(short_time)
+                MyTaskSet.finished_on_start = True
+
+            @task
+            def my_task(self):
+                MyTaskSet.my_task_run = True
+
+        class MyTestLocust(Locust):
+            task_set = MyTaskSet
+            min_wait = 0
+            max_wait = 0
+
+        options = mocked_options()
+        options.stop_timeout = short_time
+        runner = LocalLocustRunner([MyTestLocust], options)
+        runner.start_hatching(1, 1)
+        gevent.sleep(short_time / 2)
+        runner.quit()
+
+        self.assertTrue(MyTaskSet.finished_on_start)
+        self.assertFalse(MyTaskSet.my_task_run)
+
     def test_stop_timeout_exit_during_wait(self):
         short_time = 0.05
         class MyTaskSet(TaskSet):