test: add a load test shape

SYNC-3828

tests/load/Dockerfile

@@ -29,10 +29,10 @@ RUN poetry install --without dev --no-interaction --no-ansi
 RUN useradd --create-home locust
 WORKDIR /home/locust
 
-COPY ./locustfile.py locustfile.py
+COPY ./locustfiles ./locustfiles
 
 # Expose ports for the web UI and the locust master
 EXPOSE 8089 5557
 
 USER locust
-ENTRYPOINT locust -f locustfile.py --websocket_url ${SERVER_URL} --endpoint_url ${ENDPOINT_URL}
+ENTRYPOINT locust -f "locustfiles/locustfile.py,locustfiles/load.py" --websocket_url ${SERVER_URL} --endpoint_url ${ENDPOINT_URL}

tests/load/README.md

@@ -62,6 +62,7 @@ make load
 
 * In a browser navigate to `http://localhost:8089/`
 * Set up the load test parameters:
+    * ShapeClass: Default
     * Number of users: 1
     * Spawn rate: 1
     * Host: 'https://updates-autopush.stage.mozaws.net'
@@ -148,10 +149,8 @@ The load tests can be executed from the [contextual-services-test-eng cloud shel
   ```
 
 * Set up the load test parameters:
-    * Number of users: 83300
-    * Spawn rate: 694
+    * ShapeClass: 'AutopushLoadTestShape'
     * Host: 'https://updates-autopush.stage.mozaws.net'
-    * Duration (Optional): 10m
     * Websocket URL: 'wss://autopush.stage.mozaws.net'
     * Endpoint URL: 'https://updates-autopush.stage.mozaws.net'
 * Select "Start Swarming"

tests/load/docker-compose.yml

@@ -7,10 +7,8 @@ services:
     ports:
       - "8089:8089"
     environment:
+      LOCUST_USERCLASS_PICKER: true
       LOCUST_HOST: ${AUTOPUSH_HOST}
-      LOCUST_USERS: "1"
-      LOCUST_SPAWN_RATE: "1"
-      LOCUST_RUN_TIME: "10m"
       LOCUST_LOGLEVEL: "INFO"
       SERVER_URL: ${SERVER_URL}
       ENDPOINT_URL: ${ENDPOINT_URL}

tests/load/kubernetes-config/locust-master-controller.yml

@@ -24,12 +24,6 @@ spec:
               value:
             - name: LOCUST_CSV
               value:
-            - name: LOCUST_USERS
-              value:
-            - name: LOCUST_SPAWN_RATE
-              value:
-            - name: LOCUST_RUN_TIME
-              value:
             - name: LOCUST_LOGLEVEL
               value:
             - name: LOCUST_LOGFILE
@@ -54,4 +48,4 @@ spec:
               memory: 2Gi
             requests:
               cpu: 1
-              memory: 1Gi
+              memory: 1Gi

tests/load/locustfiles/load.py

@@ -0,0 +1,81 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+"""Load test shape module."""
+import math
+from typing import Callable, Type
+
+import numpy
+from locust import LoadTestShape, User
+from locustfile import AutopushUser
+
+TickTuple = tuple[int, float, list[Type[User]] | None]
+
+
+class QuadraticTrend:
+    """A class that defines a quadratic LoadTestShape trend."""
+
+    a: float
+    b: float
+    c: float
+
+    def __init__(self, max_run_time: int, max_users: int):
+        self.a, self.b, self.c = numpy.polyfit(
+            [0, (max_run_time / 2), max_run_time], [0, max_users, 0], 2
+        )
+
+    def calculate_users(self, run_time: int) -> int:
+        """Determined the number of active users given a run time.
+
+        Returns:
+            int: The number of users
+        """
+        return int(round((self.a * math.pow(run_time, 2)) + (self.b * run_time) + self.c))
+
+
+class AutopushLoadTestShape(LoadTestShape):
+    """A load test shape class for Autopush (Duration: 10 minutes, Users: 55500).
+    Note: The Shape class assumes that the workers can support the generated spawn rates.
+    """
+
+    MAX_RUN_TIME: int = 600  # 10 minutes
+    MAX_USERS: int = 83300
+    trend: QuadraticTrend
+    user_classes: list[Type[User]] = [AutopushUser]
+
+    def __init__(self):
+        super(LoadTestShape, self).__init__()
+
+        self.trend = QuadraticTrend(self.MAX_RUN_TIME, self.MAX_USERS)
+
+    def _get_calculate_users_function(self) -> Callable[[int], int]:
+        a, b, c = numpy.polyfit(
+            [0, (self.MAX_RUN_TIME / 2), self.MAX_RUN_TIME], [0, self.MAX_USERS, 0], 2
+        )
+
+        def calculation(x: int) -> int:
+            return int(round((a * math.pow(x, 2)) + (b * x) + c))
+
+        return calculation
+
+    def tick(self) -> TickTuple | None:
+        """Override defining the desired distribution for Autopush load testing.
+
+        Returns:
+            TickTuple: Distribution parameters
+                user_count: Total user count
+                spawn_rate: Number of users to start/stop per second when changing
+                            number of users
+                user_classes: None or a List of user classes to be spawned
+            None: Instruction to stop the load test
+        """
+        run_time: int = self.get_run_time()
+        if run_time > self.MAX_RUN_TIME:
+            return None
+
+        users: int = self.trend.calculate_users(run_time)
+        # The spawn_rate minimum value is 1
+        spawn_rate: int = max(abs(users - self.get_current_user_count()), 1)
+
+        return users, spawn_rate, self.user_classes

tests/load/pyproject.toml

@@ -30,6 +30,7 @@ license = "Mozilla Public License Version 2.0"
 [tool.poetry.dependencies]
 python = "^3.10"
 locust = "^2.15.1"
+numpy = "^1.25.2"
 websocket-client = "1.5.2"
 
 [tool.poetry.group.dev.dependencies]

tests/load/setup_k8s.sh

@@ -28,9 +28,6 @@ echo "Image tag for locust is set to: ${LOCUST_IMAGE_TAG}"
 ENVIRONMENT_VARIABLES=(
   "TARGET_HOST,$TARGET"
   'LOCUST_CSV,autopush'
-  'LOCUST_USERS,"83300"'
-  'LOCUST_SPAWN_RATE,"694"'
-  'LOCUST_RUN_TIME,"600"' # 10 minutes
   'LOCUST_LOGLEVEL,INFO'
   'LOCUST_LOGFILE,autopush.log'
   'SERVER_URL,wss://autopush.stage.mozaws.net'