Ensure target_user_count is set before ramping-up or down

locust/runners.py

@@ -89,7 +89,12 @@ def __init__(self, environment):
         self.worker_cpu_warning_emitted = False
         self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
         self.exceptions = {}
+        # Because of the way the ramp-up/ramp-down is implemented, target_user_classes_count
+        # is only updated at the end of the ramp-up/ramp-down.
+        # See https://github.com/locustio/locust/issues/1883#issuecomment-919239824 for context.
         self.target_user_classes_count: Dict[str, int] = {}
+        # target_user_count is set before the ramp-up/ramp-down occurs.
+        self.target_user_count: int = 0
         self.custom_messages = {}
 
         # Only when running in standalone mode (non-distributed)
@@ -324,6 +329,8 @@ def start(self, user_count: int, spawn_rate: float, wait: bool = False):
 
         self._users_dispatcher.new_dispatch(user_count, spawn_rate)
 
+        self.target_user_count = user_count
+
         try:
             for dispatched_users in self._users_dispatcher:
                 user_classes_spawn_count = {}
@@ -356,6 +363,8 @@ def start(self, user_count: int, spawn_rate: float, wait: bool = False):
 
         logger.info("All users spawned: %s" % _format_user_classes_count_for_log(self.user_classes_count))
 
+        self.target_user_classes_count = self.user_classes_count
+
         self.environment.events.spawning_complete.fire(user_count=sum(self.target_user_classes_count.values()))
 
     def start_shape(self):
@@ -448,10 +457,6 @@ def log_exception(self, node_id, msg, formatted_tb):
         row["nodes"].add(node_id)
         self.exceptions[key] = row
 
-    @property
-    def target_user_count(self) -> int:
-        return sum(self.target_user_classes_count.values())
-
     def register_message(self, msg_type, listener):
         """
         Register a listener for a custom message from another node
@@ -695,6 +700,8 @@ def start(self, user_count: int, spawn_rate: float, **kwargs) -> None:
 
         self._users_dispatcher.new_dispatch(target_user_count=user_count, spawn_rate=spawn_rate)
 
+        self.target_user_count = user_count
+
         try:
             for dispatched_users in self._users_dispatcher:
                 dispatch_greenlets = Group()
@@ -1056,6 +1063,7 @@ def start_worker(self, user_classes_count: Dict[str, int], **kwargs):
         :param user_classes_count: Users to run
         """
         self.target_user_classes_count = user_classes_count
+        self.target_user_count = sum(user_classes_count.values())
         if self.worker_state != STATE_RUNNING and self.worker_state != STATE_SPAWNING:
             self.stats.clear_all()
             self.exceptions = {}

locust/test/test_runners.py

@@ -608,6 +608,37 @@ def my_task(self):
         local_runner.stop()
         web_ui.stop()
 
+    def test_target_user_count_is_set_before_ramp_up(self):
+        """Test for https://github.com/locustio/locust/issues/1883"""
+
+        class MyUser1(User):
+            wait_time = constant(0)
+
+            @task
+            def my_task(self):
+                pass
+
+        environment = Environment(user_classes=[MyUser1])
+        runner = LocalRunner(environment)
+
+        runner.start(user_count=3, spawn_rate=1, wait=False)
+
+        gevent.sleep(1)
+
+        self.assertEqual(runner.target_user_count, 3)
+        self.assertEqual(runner.user_count, 1)
+        # However, target_user_classes_count is only updated at the end of the ramp-up/ramp-down
+        # due to the way it is implemented.
+        self.assertDictEqual({}, runner.target_user_classes_count)
+
+        runner.spawning_greenlet.join()
+
+        self.assertEqual(runner.target_user_count, 3)
+        self.assertEqual(runner.user_count, 3)
+        self.assertDictEqual({"MyUser1": 3}, runner.target_user_classes_count)
+
+        runner.quit()
+
 
 class TestMasterWorkerRunners(LocustTestCase):
     def test_distributed_integration_run(self):
@@ -1548,6 +1579,51 @@ def my_task(self):
             master.stop()
             web_ui.stop()
 
+    def test_target_user_count_is_set_before_ramp_up(self):
+        """Test for https://github.com/locustio/locust/issues/1883"""
+
+        class MyUser1(User):
+            wait_time = constant(0)
+
+            @task
+            def my_task(self):
+                pass
+
+        with mock.patch("locust.runners.WORKER_REPORT_INTERVAL", new=0.3):
+            # start a Master runner
+            master_env = Environment(user_classes=[MyUser1])
+            master = master_env.create_master_runner("*", 0)
+
+            sleep(0)
+
+            # start 1 worker runner
+            worker_env = Environment(user_classes=[MyUser1])
+            worker = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+
+            # give worker time to connect
+            sleep(0.1)
+
+            gevent.spawn(master.start, 3, spawn_rate=1)
+
+            sleep(1)
+
+            self.assertEqual(master.target_user_count, 3)
+            self.assertEqual(master.user_count, 1)
+            # However, target_user_classes_count is only updated at the end of the ramp-up/ramp-down
+            # due to the way it is implemented.
+            self.assertDictEqual({}, master.target_user_classes_count)
+
+            sleep(2)
+
+            self.assertEqual(master.target_user_count, 3)
+            self.assertEqual(master.user_count, 3)
+            self.assertDictEqual({"MyUser1": 3}, master.target_user_classes_count)
+
+            master.quit()
+
+            # make sure users are killed
+            self.assertEqual(0, worker.user_count)
+
 
 class TestMasterRunner(LocustTestCase):
     def setUp(self):