Fix #1552 by using separate greenlets to call .stop(), to allow them …

…to run in parallel.
locustio · Sep 15, 2020 · f6db44e · f6db44e
1 parent 46fe8c4
commit f6db44e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/locust/runners.py b/locust/runners.py
@@ -9,6 +9,7 @@
 from time import time
 
 import gevent
+import greenlet
 import psutil
 from gevent.pool import Group
 
@@ -18,6 +19,8 @@
 from .stats import RequestStats, setup_distributed_stats_event_listeners
 
 from .exception import RPCError
+from .user.task import LOCUST_STATE_STOPPING
+
 
 logger = logging.getLogger(__name__)
 
@@ -225,25 +228,28 @@ def stop_users(self, user_count, stop_rate=None):
             sleep_time = 1.0 / stop_rate
             logger.info("Stopping %i users at rate of %g users/s" % (user_count, stop_rate))
 
-        if self.environment.stop_timeout:
-            stop_group = Group()
+        async_calls_to_stop = Group()
+        stop_group = Group()
 
         while True:
             user_to_stop: User = to_stop.pop(random.randint(0, len(to_stop) - 1))
             logger.debug("Stopping %s" % user_to_stop._greenlet.name)
-            if self.environment.stop_timeout:
-                if not user_to_stop.stop(force=False):
-                    # User.stop() returns False if the greenlet was not stopped, so we'll need
-                    # to add it's greenlet to our stopping Group so we can wait for it to finish it's task
-                    stop_group.add(user_to_stop._greenlet)
+            if user_to_stop._greenlet is greenlet.getcurrent():
+                # User called runner.quit(), so dont block waiting for killing to finish"
+                user_to_stop._group.killone(user_to_stop._greenlet, block=False)
+            elif self.environment.stop_timeout:
+                async_calls_to_stop.add(gevent.spawn_later(0, User.stop, user_to_stop, force=False))
+                stop_group.add(user_to_stop._greenlet)
             else:
-                user_to_stop.stop(force=True)
+                async_calls_to_stop.add(gevent.spawn_later(0, User.stop, user_to_stop, force=True))
             if to_stop:
                 gevent.sleep(sleep_time)
             else:
                 break
 
-        if self.environment.stop_timeout and not stop_group.join(timeout=self.environment.stop_timeout):
+        async_calls_to_stop.join()
+
+        if not stop_group.join(timeout=self.environment.stop_timeout):
             logger.info(
                 "Not all users finished their tasks & terminated in %s seconds. Stopping them..."
                 % self.environment.stop_timeout

diff --git a/locust/user/users.py b/locust/user/users.py
@@ -185,11 +185,7 @@ def stop(self, force=False):
                       methods are called. If force is True the greenlet will be killed immediately.
         :returns: True if the greenlet was killed immediately, otherwise False
         """
-        if self._greenlet is greenlet.getcurrent():
-            # the user is stopping itself (from within a task), so blocking would deadlock
-            self._group.killone(self._greenlet, block=False)
-            return True
-        elif force or self._state == LOCUST_STATE_WAITING:
+        if force or self._state == LOCUST_STATE_WAITING:
             self._group.killone(self._greenlet)
             return True
         elif self._state == LOCUST_STATE_RUNNING: