Update user classes across workers / processes

locustio · Apr 4, 2024 · f57e004 · f57e004
1 parent a214363
commit f57e004
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git 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
@@ -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]