Adding ability to generate any custom load shape with LoadTestShape class

docs/generating-custom-load-shape.rst

@@ -0,0 +1,38 @@
+.. _generating-custom-load-shape:
+
+=================================
+Generating a custom load shape using a `LoadTestShape` class
+=================================
+
+Sometimes a completely custom shaped load test is required that cannot be achieved by simply setting or changing the user count and hatch rate. For example, generating a spike during a test or ramping up and down at custom times. In these cases using the `LoadTestShape` class can give complete control over the user count and hatch rate at all times.
+
+How does a `LoadTestShape` class work?
+---------------------------------------------
+
+Define your class inheriting the `LoadTestShape` class in your locust file. If this type of class is found then it will be automatically used by Locust. Add a `tick()` method to return either a tuple containing the user count and hatch rate or `None` to stop the load test. Locust will call the `tick()` method every second and update the load test with new settings or stop the test.
+
+Examples
+---------------------------------------------
+
+There are also some [examples on github](https://github.com/locustio/locust/tree/master/examples/custom_shape) including:
+
+- Generating a double wave shape
+- Time based stages like K6
+- Step load pattern like Visual Studio
+
+Here is a simple example that will increase user count in blocks of 100 then stop the load test after 10 minutes:
+
+```python
+class MyCustomShape(LoadTestShape):
+    time_limit = 600
+    hatch_rate = 20
+
+    def tick(self):
+        run_time = self.get_run_time()
+
+        if run_time < self.time_limit:
+            user_count = round(run_time, -2)
+            return (user_count, hatch_rate)
+
+        return None
+```

docs/index.rst

@@ -52,6 +52,7 @@ Other functionalities
 .. toctree ::
     :maxdepth: 2
 
+    generating-custom-load-shape
     running-locust-in-step-load-mode
     retrieving-stats
     testing-other-systems

docs/running-locust-distributed.rst

@@ -100,6 +100,12 @@ Running Locust distributed without the web UI
 See :ref:`running-locust-distributed-without-web-ui`
 
 
+Generating a custom load shape using a `LoadTestShape` class
+=============================================
+
+See :ref:`generating-custom-load-shape`
+
+
 Running Locust distributed in Step Load mode
 =============================================
 

examples/custom_shape/double_wave.py

@@ -0,0 +1,43 @@
+import math
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class DoubleWave(LoadTestShape):
+    """
+    A shape to immitate some specific user behaviour. In this example, midday
+    and evening meal times.
+
+    Settings:
+        min_users -- minimum users
+        peak_one_users -- users in first peak
+        peak_two_users -- users in second peak
+        time_limit -- total length of test
+    """
+
+    min_users = 20
+    peak_one_users = 60
+    peak_two_users = 40
+    time_limit = 600
+
+    def tick(self):
+        run_time = round(self.get_run_time())
+
+        if run_time < self.time_limit:
+            user_count = ((self.peak_one_users - self.min_users) * math.e ** -((run_time/(self.time_limit/10*2/3))-5) ** 2
+                        + (self.peak_two_users - self.min_users) * math.e ** - ((run_time/(self.time_limit/10*2/3))-10) ** 2
+                        + self.min_users)
+            return (round(user_count), round(user_count))
+        else:
+            return None

examples/custom_shape/stages.py

@@ -0,0 +1,49 @@
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class StagesShape(LoadTestShape):
+    """
+    A simply load test shape class that has different user and hatch_rate at
+    different stages.
+
+    Keyword arguments:
+
+        stages -- A list of dicts, each representing a stage with the following keys:
+            duration -- When this many seconds pass the test is advanced to the next stage
+            users -- Total user count
+            hatch_rate -- Hatch rate
+            stop -- A boolean that can stop that test at a specific stage
+
+        stop_at_end -- Can be set to stop once all stages have run.
+    """
+
+    stages = [
+        {'duration': 60, 'users': 10, 'hatch_rate': 10},
+        {'duration': 100, 'users': 50, 'hatch_rate': 10},
+        {'duration': 180, 'users': 100, 'hatch_rate': 10},
+        {'duration': 220, 'users': 30, 'hatch_rate': 10},
+        {'duration': 230, 'users': 10, 'hatch_rate': 10},
+        {'duration': 240, 'users': 1, 'hatch_rate': 1},
+    ]
+
+    def tick(self):
+        run_time = self.get_run_time()
+
+        for stage in self.stages:
+            if run_time < stage["duration"]:
+                tick_data = (stage["users"], stage["hatch_rate"])
+                return tick_data
+
+        return None

examples/custom_shape/step_load.py

@@ -0,0 +1,43 @@
+import math
+from locust import HttpUser, TaskSet, task, constant
+from locust import LoadTestShape
+
+
+class UserTasks(TaskSet):
+    @task
+    def get_root(self):
+        self.client.get("/")
+
+
+class WebsiteUser(HttpUser):
+    wait_time = constant(0.5)
+    tasks = [UserTasks]
+
+
+class StepLoadShape(LoadTestShape):
+    """
+    A step load shape
+
+
+    Keyword arguments:
+
+        step_time -- Time between steps
+        step_load -- User increase amount at each step
+        hatch_rate -- Hatch rate to use at every step
+        time_limit -- Time limit in seconds
+
+    """
+
+    step_time = 30
+    step_load = 10
+    hatch_rate = 10
+    time_limit = 600
+
+    def tick(self):
+        run_time = self.get_run_time()
+
+        if run_time > self.time_limit:
+            return None
+
+        current_step = math.floor(run_time / self.step_time) + 1
+        return (current_step * self.step_load, self.hatch_rate)

locust/__init__.py

@@ -8,6 +8,7 @@
 from .user.task import task, tag, TaskSet
 from .user.users import HttpUser, User
 from .user.wait_time import between, constant, constant_pacing
+from .shape import LoadTestShape
 
 from .event import Events
 events = Events()

locust/env.py

@@ -4,6 +4,7 @@
 from .runners import LocalRunner, MasterRunner, WorkerRunner
 from .web import WebUI
 from .user.task import filter_tasks_by_tags
+from .shape import LoadTestShape
 
 
 class Environment:
@@ -15,7 +16,10 @@ class Environment:
 
     user_classes = []
     """User classes that the runner will run"""
-
+
+    shape_class = None
+    """A shape class to control the shape of the load test"""
+
     tags = None
     """If set, only tasks that are tagged by tags in this list will be executed"""
 
@@ -63,12 +67,13 @@ class Environment:
     def  __init__(
         self, *,
         user_classes=[],
+        shape_class=None,
         tags=None,
         exclude_tags=None,
-        events=None, 
-        host=None, 
-        reset_stats=False, 
-        step_load=False, 
+        events=None,
+        host=None,
+        reset_stats=False,
+        step_load=False,
         stop_timeout=None,
         catch_exceptions=True,
         parsed_options=None,
@@ -79,6 +84,7 @@ def  __init__(
             self.events = Events()
 
         self.user_classes = user_classes
+        self.shape_class = shape_class
         self.tags = tags
         self.exclude_tags = exclude_tags
         self.stats = RequestStats()

locust/main.py

@@ -21,6 +21,8 @@
 from .user.inspectuser import get_task_ratio_dict, print_task_ratio
 from .util.timespan import parse_timespan
 from .exception import AuthCredentialsError
+from .shape import LoadTestShape
+
 
 version = locust.__version__
 
@@ -36,6 +38,16 @@ def is_user_class(item):
     )
 
 
+def is_shape_class(item):
+    """
+    Check if a class is a LoadTestShape
+    """
+    return bool(
+        inspect.isclass(item)
+        and issubclass(item, LoadTestShape)
+        and item.__dict__['__module__'] != 'locust.shape'
+    )
+
 def load_locustfile(path):
     """
     Import given locustfile path and return (docstring, callables).
@@ -84,15 +96,24 @@ def __import_locustfile__(filename, path):
         del sys.path[0]
     # Return our two-tuple
     user_classes = {name:value for name, value in vars(imported).items() if is_user_class(value)}
-    return imported.__doc__, user_classes
+
+    # Find shape class, if any, return it
+    shape_classes = [value for name, value in vars(imported).items() if is_shape_class(value)]
+    if shape_classes:
+        shape_class = shape_classes[0]()
+    else:
+        shape_class = None
+
+    return imported.__doc__, user_classes, shape_class
 
 
-def create_environment(user_classes, options, events=None):
+def create_environment(user_classes, options, events=None, shape_class=None):
     """
     Create an Environment instance from options
     """
     return Environment(
         user_classes=user_classes,
+        shape_class=shape_class,
         tags=options.tags,
         exclude_tags=options.exclude_tags,
         events=events,
@@ -110,8 +131,8 @@ def main():
     locustfile = parse_locustfile_option()
 
     # import the locustfile
-    docstring, user_classes = load_locustfile(locustfile)
-    
+    docstring, user_classes, shape_class = load_locustfile(locustfile)
+
     # parse all command line options
     options = parse_options()
 
@@ -163,8 +184,12 @@ def main():
         logger.warning("System open file limit setting is not high enough for load testing, and the OS didn't allow locust to increase it by itself. See https://docs.locust.io/en/stable/installation.html#increasing-maximum-number-of-open-files-limit for more info.")
 
     # create locust Environment
-    environment = create_environment(user_classes, options, events=locust.events)
-
+    environment = create_environment(user_classes, options, events=locust.events, shape_class=shape_class)
+
+    if shape_class and (options.num_users or options.hatch_rate or options.step_load):
+        logger.error("The specified locustfile contains a shape class but a conflicting argument was specified: users, hatch-rate or step-load")
+        sys.exit(1)
+
     if options.show_task_ratio:
         print("\n Task ratio per User class")
         print( "-" * 80)

locust/runners.py

@@ -49,6 +49,8 @@ def __init__(self, environment):
         self.state = STATE_INIT
         self.hatching_greenlet = None
         self.stepload_greenlet = None
+        self.shape_greenlet = None
+        self.shape_last_state = None
         self.current_cpu_usage = 0
         self.cpu_warning_emitted = False
         self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
@@ -288,12 +290,37 @@ def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
         while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
             current_num_users += step_users_growth
             if current_num_users > int(self.total_users):
-                logger.info('Step Load is finished.')
+                logger.info("Step Load is finished")
                 break
             self.start(current_num_users, hatch_rate)
-            logger.info('Step loading: start hatch job of %d user.' % (current_num_users))
+            logger.info("Step loading: start hatch job of %d user" % (current_num_users))
             gevent.sleep(step_duration)
 
+    def start_shape(self):
+        if self.shape_greenlet:
+            logger.info("There is an ongoing shape test running. Editing is disabled")
+            return
+
+        logger.info("Shape test starting. User count and hatch rate are ignored for this type of load test")
+        self.state = STATE_INIT
+        self.shape_greenlet = self.greenlet.spawn(self.shape_worker)
+        self.shape_greenlet.link_exception(greenlet_exception_handler)
+
+    def shape_worker(self):
+        logger.info("Shape worker starting")
+        while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
+            new_state = self.environment.shape_class.tick()
+            if new_state is None:
+                logger.info("Shape test stopping")
+                self.stop()
+            elif self.shape_last_state == new_state:
+                gevent.sleep(1)
+            else:
+                user_count, hatch_rate = new_state
+                logger.info("Shape test updating to %d users at %.2f hatch rate" % (user_count, hatch_rate))
+                self.start(user_count=user_count, hatch_rate=hatch_rate)
+                self.shape_last_state = new_state
+
     def stop(self):
         """
         Stop a running load test by stopping all running users