Repository: locustio/locust

Issue:
cannot specify task of UserClass when locust run on distributed mode
### Prerequisites

- [X] I am using [the latest version of Locust](https://github.com/locustio/locust/releases/)
- [X] I am reporting a bug, not asking a question

### Description

if launch locust with `--processes` argument, choosing task will not working. 

![image](https://github.com/locustio/locust/assets/3401642/2067a1b1-26d5-40b1-9f65-37b729096d19)


### Command line

locust --processes 2 --class-picker -f locustfiles

### Locustfile contents

```python3
...
```


### Python version

3.12

### Locust version

2.24

### Operating system

macOS or centOS


Discussion:
huh. good find. unfortunately it will probably be hard to support, so I'll probably "fix" it by throwing an error (at least for now)
I guess the reason it's hard to support is that in distributed mode, there's no shared data strategy. The single-process mode uses memory, while distributed mode would require introducing Redis or other memory database.






Actually. When I think about it, it _should_ work. Not sure when I'll have time to look at it though...
It would be the same for distributed mode, we need to communicate between the workers and the master node so that the workers know that the list of tasks has changed. @cyberw Perhaps in the meantime we should throw an error or note something in the documentation that these arguments don't work together yet
We shouldnt really have options that only work with non-distributed so we need to try and solve it (or remove task selection entirely). But erroring out when combined with --processes or --master is better than the current state I guess :)

Resolving patch:
diff --git a/locust/env.py b/locust/env.py
--- a/locust/env.py
+++ b/locust/env.py
@@ -198,6 +198,9 @@ def create_web_ui(
         return self.web_ui
 
     def update_user_class(self, user_settings):
+        if isinstance(self.runner, MasterRunner):
+            self.runner.send_message("update_user_class", user_settings)
+
         user_class_name = user_settings.get("user_class_name")
         user_class = self.available_user_classes[user_class_name]
         user_tasks = self.available_user_tasks[user_class_name]
diff --git a/locust/runners.py b/locust/runners.py
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -1395,6 +1395,8 @@ def worker(self) -> NoReturn:
                 self.reset_connection()
             elif msg.type == "heartbeat":
                 self.last_heartbeat_timestamp = time.time()
+            elif msg.type == "update_user_class":
+                self.environment.update_user_class(msg.data)
             elif msg.type in self.custom_messages:
                 logger.debug("Received %s message from master" % msg.type)
                 listener, concurrent = self.custom_messages[msg.type]


Test patch:
diff --git a/locust/test/test_env.py b/locust/test/test_env.py
--- a/locust/test/test_env.py
+++ b/locust/test/test_env.py
@@ -241,3 +241,43 @@ def my_task(self):
             environment.available_user_classes["User1"].json(),
             {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
         )
+
+    def test_distributed_update_user_class(self):
+        class MyUser1(User):
+            @task
+            def my_task(self):
+                pass
+
+            @task
+            def my_task_2(self):
+                pass
+
+        class MyUser2(User):
+            @task
+            def my_task(self):
+                pass
+
+        master_env = Environment(
+            user_classes=[MyUser1, MyUser2],
+            available_user_classes={"User1": MyUser1, "User2": MyUser2},
+            available_user_tasks={"User1": MyUser1.tasks, "User2": MyUser2.tasks},
+        )
+        master = master_env.create_master_runner("*", 0)
+
+        worker_env = Environment(
+            user_classes=[MyUser1, MyUser2],
+            available_user_classes={"User1": MyUser1, "User2": MyUser2},
+            available_user_tasks={"User1": MyUser1.tasks, "User2": MyUser2.tasks},
+        )
+        worker = worker_env.create_worker_runner("127.0.0.1", master.server.port)
+
+        master_env.update_user_class({"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]})
+
+        self.assertEqual(
+            master_env.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )
+        self.assertEqual(
+            worker_env.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )
diff --git a/locust/test/test_web.py b/locust/test/test_web.py
--- a/locust/test/test_web.py
+++ b/locust/test/test_web.py
@@ -1351,3 +1351,32 @@ def test_html_stats_report(self):
         self.assertTrue(d("#root"))
         self.assertIn('"locustfile": "locust.py"', str(d))
         self.assertIn('"host": "http://localhost"', str(d))
+
+    def test_update_user(self):
+        class MyUser1(User):
+            @task
+            def my_task(self):
+                pass
+
+            @task
+            def my_task_2(self):
+                pass
+
+        class MyUser2(User):
+            @task
+            def my_task(self):
+                pass
+
+        self.environment.user_classes = [MyUser1, MyUser2]
+        self.environment.available_user_classes = {"User1": MyUser1, "User2": MyUser2}
+        self.environment.available_user_tasks = {"User1": MyUser1.tasks, "User2": MyUser2.tasks}
+
+        requests.post(
+            "http://127.0.0.1:%i/user" % self.web_port,
+            json={"user_class_name": "User1", "host": "http://localhost", "tasks": ["my_task_2"]},
+        )
+
+        self.assertEqual(
+            self.environment.available_user_classes["User1"].json(),
+            {"host": "http://localhost", "tasks": ["my_task_2"], "fixed_count": 0, "weight": 1},
+        )