Only raise GreenletExit in TaskSet.run(), because that is the only pl…

…ace it is really needed.
locustio · Oct 26, 2019 · 96d6759 · 96d6759
1 parent 7cbdb69
commit 96d6759
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/locust/core.py b/locust/core.py
@@ -361,6 +361,9 @@ def run(self, *args, **kwargs):
 
         while (True):
             try:
+                if self.locust._state == LOCUST_STATE_STOPPING:
+                    raise GreenletExit()
+
                 if self.locust.stop_timeout is not None and time() - self._time_start > self.locust.stop_timeout:
                     return
 
@@ -370,8 +373,6 @@ def run(self, *args, **kwargs):
                 try:
                     self.execute_next_task()
                 except RescheduleTaskImmediately:
-                    if self.locust._state == LOCUST_STATE_STOPPING:
-                        raise GreenletExit()
                     pass
                 except RescheduleTask:
                     self.wait()
@@ -435,8 +436,6 @@ def get_wait_secs(self):
         return millis / 1000.0
 
     def wait(self):
-        if self.locust._state == LOCUST_STATE_STOPPING:
-            raise GreenletExit()
         self.locust._state = LOCUST_STATE_WAITING
         self._sleep(self.get_wait_secs())
         self.locust._state = LOCUST_STATE_RUNNING