Repository: locustio/locust

Issue:
MasterRunner target_user_count no longer set for test_start event listeners
### Describe the bug
I was following the `custom_messages.py` example and noticed that `environment.runner.target_user_count` is no longer resolving to the number of users being created for the test. This functionality appears to have been broken between tags `2.0.0b1` and `2.0.0b2`.

### Expected behavior
The MasterRunner's `target_user_count` should resolve to the number of users specified in the WebUI for the test.
The `custom_messages.py` example for a test with 10 users and 2 workers should log a message from each worker indicating that they received 5 users.
```
master_1  | [2021-09-14 02:48:02,550] d4b6940eebd0/INFO/locust.runners: Sending spawn jobs of 10 users at 1.00 spawn rate to 2 ready clients
master_1  | Thanks for the 5 users!
master_1  | Thanks for the 5 users!
worker_1  | User4
worker_2  | User9
worker_1  | User3
...
```

### Actual behavior
The MasterRunner's `target_user_count` is 0, causing numerous exceptions in the `custom_messages.py` example.
```
master_1  | [2021-09-14 03:01:33,449] 0134f51ea2b2/INFO/locust.runners: Sending spawn jobs of 10 users at 1.00 spawn rate to 2 ready clients
master_1  | Thanks for the 0 users!
master_1  | Thanks for the 0 users!
worker_2  | Traceback (most recent call last):
worker_2  |   File "src/gevent/greenlet.py", line 906, in gevent._gevent_cgreenlet.Greenlet.run
worker_2  |   File "/usr/local/lib/python3.8/site-packages/locust/runners.py", line 1138, in <lambda>
worker_2  |     self.spawning_greenlet = self.greenlet.spawn(lambda: self.start_worker(job["user_classes_count"]))
worker_2  |   File "/usr/local/lib/python3.8/site-packages/locust/runners.py", line 1075, in start_worker
worker_2  |     self.spawn_users(user_classes_spawn_count)
worker_2  |   File "/usr/local/lib/python3.8/site-packages/locust/runners.py", line 221, in spawn_users
worker_2  |     new_users += spawn(user_class, spawn_count)
worker_2  |   File "/usr/local/lib/python3.8/site-packages/locust/runners.py", line 210, in spawn
worker_2  |     new_user = self.user_classes_by_name[user_class](self.environment)
worker_2  |   File "/mnt/locust/custom_messages.py", line 55, in __init__
worker_2  |     self.username = usernames.pop()
worker_2  | IndexError: pop from empty list
worker_2  | 2021-09-14T03:01:33Z <Greenlet at 0x7fc070a7ae10: <lambda>> failed with IndexError
...
```

### Steps to reproduce
1. Create `docker-compose.yml`:
```yaml
version: '3'

services:
  master:
    image: locustio/locust:2.2.1
    # image: locustio/locust:2.0.0b1 ## Last known working version
    ports:
     - "8089:8089"
    volumes:
      - ./:/mnt/locust
    command: -f /mnt/locust/custom_messages.py --master -H http://master:8089
  
  worker:
    image: locustio/locust:2.2.1
    # image: locustio/locust:2.0.0b1 ## Last known working version
    volumes:
      - ./:/mnt/locust
    command: -f /mnt/locust/custom_messages.py --worker --master-host master
    deploy:
      mode: replicated
      replicas: 2
```
2. Copy the existing [examples/custom_messages.py](https://raw.githubusercontent.com/locustio/locust/master/examples/custom_messages.py) into the same directory as the `docker-compose.yml`
3. Run `docker compose up`
4. Visit http://localhost:8089/ and start a test with 10 users
5. (Optional) Modify the `docker-compose.yml` to use a different image tag listed below and repeat steps 3 and 4.

### Environment

- OS: `Debian GNU/Linux 11 (bullseye)` from `locustio/locust` docker image.
- Python version: `3.8.11` and `3.8.12` from the various `locustio/locust` docker images.
- Locust version: `2.0.0b2`, `2.0.0b3`, `2.1.0`, `2.2.0`, `2.2.1`
- Locust command line that you ran: See `docker-compose.yml` above.
- Locust file contents (anonymized if necessary): [examples/custom_messages.py](https://raw.githubusercontent.com/locustio/locust/master/examples/custom_messages.py)


Discussion:
Thanks for the nice and detailed description. I think this is either @mboutet 's area (the only relevant PR between b1 and b2 was made by him) or @nathan-beam 's as he wrote the custom messages feature.
I see where the issue is coming from. The relevant changes are [here](https://github.com/locustio/locust/compare/2.0.0b1...2.0.0b2#diff-1c6279e1bae3e0b2b378639df87fc31fdb5da6e806b8d4e8cbc929bb04e78c81L645-R730).

Before `2.0.0b2`, the "weighing" of users was done before the dispatch to workers took place. However, in https://github.com/locustio/locust/pull/1809, the weighing is done on-the-fly during dispatch using a smooth weighted round robin generator. As seen in the above linked diff, `target_user_classes_count` is only set after the ramp-down/up is finished. I agree this is not ideal because it makes sense to know what is the target before actually reaching this target. 

Switching to this approach of weighing on-the-fly was one of the key point in addressing the performance issues while addressing the other issues in https://github.com/locustio/locust/pull/1621.

Right now, the `target_user_count` is a `property` computed from `target_user_classes_count`. So we could add:
```python
self.target_user_classes_count = self._users_dispatcher._distribute_users(user_count)[0]
```
at line:
https://github.com/locustio/locust/blob/1c65699f26fd5f4368841bbaea83256300e9594c/locust/runners.py#L690
The downside is that calling `_distribute_users` for a large user count (I'd say more than 30k-50k users) can be computationally expensive, but I think for most use cases, this should not be a problem. 

Another solution could be to no longer compute `target_user_count` from `target_user_classes_count`. But I don't feel that this solution solves the problem for real because the day someone needs `target_user_classes_count` in a custom message instead of `target_user_count`, we'll have the same problem.

Hmm... I think your second proposed solution, while a little more limited in what flexibility we give the user, is better. Doing all that work just to get the target user count (which already exists) just doesnt feel right. 

I think it is reasonable to assume that 99% of tests will not require exact knowledge of the count of each user class. If that is necessary, then maybe it is something we could allow the locustfile to calculate only if it is needed.
If there's another mechanism for retrieving the target user count that was entered in the WebUI or passed in via the `-u/--users` CLI arg in a `test_start` event listener, that would similarly solve my issue.
I'm not too savvy with Python, though I did try `inspect.getmembers(...)` on various objects that were available to the function, but could not find anything that contained the value.
As a workaround, if you specify the number of users on command line you can get it from `environment.parsed_options.users`. I'm not sure, but it could maybe be argued that the web ui should overwrite that value when the test is started from there.
Hello, I'm having the exact same issue, even after updating to 2.2.3. 
What would be an alternative to use instead of `target_user_count`? I'm using `environment.parsed_options.num_users` for the time being, but it doesn't update if you change the user count in the web UI.

Resolving patch:
diff --git a/locust/runners.py b/locust/runners.py
--- a/locust/runners.py
+++ b/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 = {}


Test patch:
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
--- a/locust/test/test_runners.py
+++ b/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):