Move User selection responsibility from worker to master in order to fix unbalanced distribution of users and uneven ramp-up

locust/distribution.py

@@ -0,0 +1,115 @@
+import math
+from itertools import combinations_with_replacement
+from operator import attrgetter
+from typing import Dict, List, Type
+
+from locust import User
+
+
+def weight_users(user_classes: List[Type[User]], user_count: int) -> Dict[str, int]:
+    """
+    Compute the desired state of users using the weight of each user class.
+
+    :param user_classes: the list of user class
+    :param user_count: total number of users
+    :return: the set of users to run
+    """
+    assert user_count >= 0
+
+    if len(user_classes) == 0:
+        return {}
+
+    user_classes = sorted(user_classes, key=attrgetter("__name__"))
+
+    user_classes_count = {user_class.__name__: 0 for user_class in user_classes}
+
+    # If the number of users is less than the number of user classes, at most one user of each user class
+    # is chosen. User classes with higher weight are chosen first.
+    if user_count <= len(user_classes):
+        user_classes_count.update(
+            {
+                user_class.__name__: 1
+                for user_class in sorted(user_classes, key=attrgetter("weight"), reverse=True)[:user_count]
+            }
+        )
+        return user_classes_count
+
+    # If the number of users is greater than or equal to the number of user classes, at least one user of each
+    # user class will be chosen. The greater number of users is, the better the actual distribution
+    # of users will match the desired one (as dictated by the weight attributes).
+    weights = list(map(attrgetter("weight"), user_classes))
+    relative_weights = [weight / sum(weights) for weight in weights]
+    user_classes_count = {
+        user_class.__name__: round(relative_weight * user_count) or 1
+        for user_class, relative_weight in zip(user_classes, relative_weights)
+    }
+
+    if sum(user_classes_count.values()) == user_count:
+        return user_classes_count
+
+    else:
+        user_classes_count = _find_ideal_users_to_add_or_remove(
+            user_classes, user_count - sum(user_classes_count.values()), user_classes_count
+        )
+        assert sum(user_classes_count.values()) == user_count
+        return user_classes_count
+
+
+def _find_ideal_users_to_add_or_remove(
+    user_classes: List[Type[User]], user_count_to_add_or_remove: int, user_classes_count: Dict[str, int]
+) -> Dict[str, int]:
+    sign = -1 if user_count_to_add_or_remove < 0 else 1
+
+    user_count_to_add_or_remove = abs(user_count_to_add_or_remove)
+
+    assert user_count_to_add_or_remove <= len(user_classes), user_count_to_add_or_remove
+
+    # Formula for combination with replacement
+    # (https://www.tutorialspoint.com/statistics/combination_with_replacement.htm)
+    number_of_combinations = math.factorial(len(user_classes) + user_count_to_add_or_remove - 1) / (
+        math.factorial(user_count_to_add_or_remove) * math.factorial(len(user_classes) - 1)
+    )
+
+    # If the number of combinations with replacement is above this threshold, we simply add/remove
+    # users for the first "number_of_users_to_add_or_remove" users. Otherwise, computing the best
+    # distribution is too expensive in terms of computation.
+    max_number_of_combinations_threshold = 1000
+
+    if number_of_combinations <= max_number_of_combinations_threshold:
+        user_classes_count_candidates: Dict[float, Dict[str, int]] = {}
+        for user_classes_combination in combinations_with_replacement(user_classes, user_count_to_add_or_remove):
+            # Copy in order to not mutate `user_classes_count` for the parent scope
+            user_classes_count_candidate = user_classes_count.copy()
+            for user_class in user_classes_combination:
+                user_classes_count_candidate[user_class.__name__] += sign
+            distance = distance_from_desired_distribution(user_classes, user_classes_count_candidate)
+            if distance not in user_classes_count_candidates:
+                user_classes_count_candidates[distance] = user_classes_count_candidate.copy()
+
+        return user_classes_count_candidates[min(user_classes_count_candidates.keys())]
+
+    else:
+        # Copy in order to not mutate `user_classes_count` for the parent scope
+        user_classes_count_candidate = user_classes_count.copy()
+        for user_class in user_classes[:user_count_to_add_or_remove]:
+            user_classes_count_candidate[user_class.__name__] += sign
+        return user_classes_count_candidate
+
+
+def distance_from_desired_distribution(user_classes: List[Type[User]], user_classes_count: Dict[str, int]) -> float:
+    actual_ratio_of_user_class = {
+        user_class: user_class_count / sum(user_classes_count.values())
+        for user_class, user_class_count in user_classes_count.items()
+    }
+
+    expected_ratio_of_user_class = {
+        user_class.__name__: user_class.weight / sum(map(attrgetter("weight"), user_classes))
+        for user_class in user_classes
+    }
+
+    differences = [
+        actual_ratio_of_user_class[user_class] - expected_ratio
+        for user_class, expected_ratio in expected_ratio_of_user_class.items()
+    ]
+
+    return math.sqrt(math.fsum(map(lambda x: x ** 2, differences)))

locust/env.py

@@ -1,3 +1,11 @@
+from operator import methodcaller
+from typing import (
+    Dict,
+    List,
+    Type,
+    TypeVar,
+)
+
 from .event import Events
 from .exception import RunnerAlreadyExistsError
 from .stats import RequestStats
@@ -6,7 +14,9 @@
 from .user import User
 from .user.task import filter_tasks_by_tags
 from .shape import LoadTestShape
-from typing import List
+
+
+RunnerType = TypeVar("RunnerType", bound=Runner)
 
 
 class Environment:
@@ -16,7 +26,7 @@ class Environment:
     See :ref:`events` for available events.
     """
 
-    user_classes: List[User] = []
+    user_classes: List[Type[User]] = []
     """User classes that the runner will run"""
 
     shape_class: LoadTestShape = None
@@ -95,24 +105,37 @@ def __init__(
 
         self._filter_tasks_by_tags()
 
-    def _create_runner(self, runner_class, *args, **kwargs):
+        # Validate there's no class with the same name but in different modules
+        if len(set(user_class.__name__ for user_class in self.user_classes)) != len(self.user_classes):
+            raise ValueError(
+                "The following user classes have the same class name: {}".format(
+                    ", ".join(map(methodcaller("fullname"), self.user_classes))
+                )
+            )
+
+    def _create_runner(
+        self,
+        runner_class: Type[RunnerType],
+        *args,
+        **kwargs,
+    ) -> RunnerType:
         if self.runner is not None:
             raise RunnerAlreadyExistsError("Environment.runner already exists (%s)" % self.runner)
-        self.runner = runner_class(self, *args, **kwargs)
+        self.runner: RunnerType = runner_class(self, *args, **kwargs)
 
         # Attach the runner to the shape class so that the shape class can access user count state
         if self.shape_class:
             self.shape_class.runner = self.runner
 
         return self.runner
 
-    def create_local_runner(self):
+    def create_local_runner(self) -> LocalRunner:
         """
         Create a :class:`LocalRunner <locust.runners.LocalRunner>` instance for this Environment
         """
         return self._create_runner(LocalRunner)
 
-    def create_master_runner(self, master_bind_host="*", master_bind_port=5557):
+    def create_master_runner(self, master_bind_host="*", master_bind_port=5557) -> MasterRunner:
         """
         Create a :class:`MasterRunner <locust.runners.MasterRunner>` instance for this Environment
 
@@ -126,7 +149,7 @@ def create_master_runner(self, master_bind_host="*", master_bind_port=5557):
             master_bind_port=master_bind_port,
         )
 
-    def create_worker_runner(self, master_host, master_port):
+    def create_worker_runner(self, master_host, master_port) -> WorkerRunner:
         """
         Create a :class:`WorkerRunner <locust.runners.WorkerRunner>` instance for this Environment
 
@@ -191,3 +214,7 @@ def _filter_tasks_by_tags(self):
 
         for user_class in self.user_classes:
             filter_tasks_by_tags(user_class, self.tags, self.exclude_tags)
+
+    @property
+    def user_classes_by_name(self) -> Dict[str, Type[User]]:
+        return {u.__name__: u for u in self.user_classes}

locust/html.py

@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 from jinja2 import Environment, FileSystemLoader
 import os
 import pathlib

locust/input_events.py

@@ -1,3 +1,5 @@
+from typing import Dict
+
 import gevent
 import logging
 import os
@@ -86,7 +88,7 @@ def get_poller():
         return UnixKeyPoller()
 
 
-def input_listener(key_to_func_map):
+def input_listener(key_to_func_map: Dict[str, callable]):
     def input_listener_func():
         try:
             poller = get_poller()

locust/main.py

@@ -306,6 +306,7 @@ def main():
         web_ui.start()
         main_greenlet = web_ui.greenlet
 
+    headless_master_greenlet = None
     if options.headless:
         # headless mode
         if options.master:
@@ -316,6 +317,9 @@ def main():
                     len(runner.clients.ready),
                     options.expect_workers,
                 )
+                # TODO: Handle KeyboardInterrupt and send quit signal to workers that are started.
+                #       Right now, if the user sends a ctrl+c, the master will not gracefully
+                #       shutdown resulting in all the already started workers to stay active.
                 time.sleep(1)
         if not options.worker:
             # apply headless mode defaults
@@ -328,7 +332,8 @@ def main():
             if environment.shape_class:
                 environment.runner.start_shape()
             else:
-                runner.start(options.num_users, options.spawn_rate)
+                headless_master_greenlet = gevent.spawn(runner.start, options.num_users, options.spawn_rate)
+                headless_master_greenlet.link_exception(greenlet_exception_handler)
 
     def spawn_run_time_limit_greenlet():
         def timelimit_stop():
@@ -351,19 +356,19 @@ def timelimit_stop():
         input_listener_greenlet = gevent.spawn(
             input_listener(
                 {
-                    "w": lambda: runner.spawn_users(1, 100)
+                    "w": lambda: runner.start(runner.user_count + 1, 100)
                     if runner.state != "spawning"
                     else logging.warning("Already spawning users, can't spawn more right now"),
-                    "W": lambda: runner.spawn_users(10, 100)
+                    "W": lambda: runner.start(runner.user_count + 10, 100)
                     if runner.state != "spawning"
                     else logging.warning("Already spawning users, can't spawn more right now"),
-                    "s": lambda: runner.stop_users(1)
+                    "s": lambda: runner.start(max(0, runner.user_count - 1), 100)
                     if runner.state != "spawning"
                     else logging.warning("Spawning users, can't stop right now"),
-                    "S": lambda: runner.stop_users(10)
+                    "S": lambda: runner.start(max(0, runner.user_count - 10), 100)
                     if runner.state != "spawning"
                     else logging.warning("Spawning users, can't stop right now"),
-                }
+                },
             )
         )
         input_listener_greenlet.link_exception(greenlet_exception_handler)
@@ -403,6 +408,8 @@ def shutdown():
         logger.info("Shutting down (exit code %s), bye." % code)
         if stats_printer_greenlet is not None:
             stats_printer_greenlet.kill(block=False)
+        if headless_master_greenlet is not None:
+            headless_master_greenlet.kill(block=False)
         logger.info("Cleaning up runner...")
         if runner is not None:
             runner.quit()