add support for load shapes (#105)

- add support for load shapes
- add load shapes for step load and spike patterns
- update dependencies
- update README.md
- fix bug in beacon user

Author: erwin-wee

README.md

@@ -112,7 +112,9 @@ This will run a load test with a general BSC profile.
 
 ### Parameters and Flags
 - `-p, --profile`: Specifies the profile to use for the benchmark. Available profiles can be found in the profile directory. Sample usage `-p bsc.general`
+- `-s, --shape`: Specifies the shape of the load pattern. List available shapes with `chainbench list shapes`.
 - `-u, --users`: Sets the number of simulated users to use for the benchmark.
+- `-r, --spawn-rate`: Sets the spawn rate of users per second.
 - `-w, --workers`: Sets the number of worker threads to use for the benchmark.
 - `-t, --test-time`: Sets the duration of the test to run.
 - `--target`: Specifies the target blockchain node URL that the benchmark will connect to.
@@ -155,6 +157,15 @@ Here's an example of how to run a load test for Ethereum using the `evm.light` p
 chainbench start --profile evm.light --users 50 --workers 2 --test-time 12h --target https://node-url --headless --autoquit
 ```
 
+## Load Pattern Shapes
+Load pattern shapes are used to define how the load will be distributed over time. You may specify the shape of the load pattern using the `-s` or `--shape` flag.
+This is an optional flag and if not specified, the default shape will be used. The default shape is `ramp-up` which means the load will increase linearly over time at
+the spawn-rate until the specified number of users is reached, after that it will maintain the number of users until test duration is over.
+
+Other available shapes are:
+- `step` - The load will increase in steps. `--spawn-rate` flag is required to specify the step size. The number of steps will be calculated based on `--users` divided by `--spawn-rate`. The duration of each step will be calculated based on `--test-time` divided by the number of steps.
+- `spike` - The load will run in a spike pattern. The load will ramp up to 10% of the total users for 40% of the test duration and then spike to 100% of the total users as specified by `--users` for 20% of test duration and then reduce back to 10% of total users until the test duration is over.
+
 ### Test Data Size
 You may specify the test data size using the `--size` flag. This will determine how much data is used in the test.
 Take note that larger data size will result in longer test data generation time before the test starts.

chainbench/main.py

@@ -1,4 +1,3 @@
-import logging
 import os
 import shlex
 import subprocess
@@ -34,8 +33,6 @@
 DEFAULT_PROFILE = "ethereum.general"
 NOTIFY_URL_TEMPLATE = "https://ntfy.sh/{topic}"
 
-logger = logging.getLogger(__name__)
-
 
 @click.group(
     help="Tool for flexible blockchain infrastructure benchmarking.",
@@ -108,6 +105,12 @@ def validate_profile(ctx: Context, param: Parameter, value: str) -> str:
     help="Profile to run",
     show_default=True,
 )
+@click.option(
+    "-s",
+    "--shape",
+    default=None,
+    help="Shape of load pattern",
+)
 @click.option("-H", "--host", default=MASTER_HOST, help="Host to run on", show_default=True)
 @click.option("-P", "--port", default=MASTER_PORT, help="Port to run on", show_default=True)
 @click.option(
@@ -188,6 +191,7 @@ def start(
     ctx: Context,
     profile: str,
     profile_dir: Path | None,
+    shape: str | None,
     host: str,
     port: int,
     workers: int,
@@ -248,12 +252,12 @@ def start(
         sys.exit(1)
 
     if test_by_directory:
-        from locust.argument_parser import find_locustfiles
+        from locust.argument_parser import parse_locustfile_paths
         from locust.util.load_locustfile import load_locustfile
 
         user_classes = {}
         test_data_types = set()
-        for locustfile in find_locustfiles([profile_path.__str__()], True):
+        for locustfile in parse_locustfile_paths([profile_path.__str__()]):
             _, _user_classes, _ = load_locustfile(locustfile)
             for key, value in _user_classes.items():
                 user_classes[key] = value
@@ -286,6 +290,13 @@ def start(
         click.echo(f"Testing profile: {profile}")
         test_plan = profile
 
+    if shape is not None:
+        shapes_dir = get_base_path(__file__) / "shapes"
+        shape_path = get_profile_path(shapes_dir, shape)
+        click.echo(f"Using load shape: {shape}")
+    else:
+        shape_path = None
+
     results_dir = Path(results_dir).resolve()
     results_path = ensure_results_dir(profile=profile, parent_dir=results_dir, run_id=run_id)
 
@@ -331,6 +342,7 @@ def start(
         target=target,
         custom_tags=custom_tags,
         exclude_tags=custom_exclude_tags,
+        shape_path=shape_path,
         timescale=timescale,
         pg_host=pg_host,
         pg_port=pg_port,
@@ -480,6 +492,14 @@ def profiles(profile_dir: Path) -> None:
         click.echo(profile)
 
 
+@_list.command(
+    help="Lists all available load shapes.",
+)
+def shapes() -> None:
+    for shape in get_profiles(get_base_path(__file__) / "shapes"):
+        click.echo(shape)
+
+
 @_list.command(
     help="Lists all available methods.",
 )

chainbench/shapes/spike.py

@@ -0,0 +1,28 @@
+from locust import LoadTestShape
+
+
+class SpikeLoadShape(LoadTestShape):
+    """
+    A step load shape class that has the following shape:
+    10% of users start at the beginning for 40% of the test duration, then 100% of users for 20% of the test duration,
+    then 10% of users until the end of the test duration.
+    """
+
+    use_common_options = True
+
+    def tick(self):
+        run_time = self.get_run_time()
+        total_run_time = self.runner.environment.parsed_options.run_time
+        period_duration = round(total_run_time / 10)
+        spike_run_time_start = period_duration * 4
+        spike_run_time_end = period_duration * 6
+
+        if run_time < spike_run_time_start:
+            user_count = round(self.runner.environment.parsed_options.num_users / 10)
+            return user_count, self.runner.environment.parsed_options.spawn_rate
+        elif run_time < spike_run_time_end:
+            return self.runner.environment.parsed_options.num_users, self.runner.environment.parsed_options.spawn_rate
+        elif run_time < total_run_time:
+            user_count = round(self.runner.environment.parsed_options.num_users / 10)
+            return user_count, self.runner.environment.parsed_options.spawn_rate
+        return None

chainbench/shapes/step.py

@@ -0,0 +1,25 @@
+import math
+
+from locust import LoadTestShape
+
+
+class StepLoadShape(LoadTestShape):
+    """
+    This load shape determines the number of steps by using the total number of users divided by the spawn rate.
+    Duration of each step is calculated by dividing the total run time by the number of steps equally.
+    """
+
+    use_common_options = True
+
+    def tick(self):
+        run_time = self.get_run_time()
+        total_run_time = self.runner.environment.parsed_options.run_time
+
+        if run_time < total_run_time:
+            step = self.runner.environment.parsed_options.spawn_rate
+            users = self.runner.environment.parsed_options.num_users
+            no_of_steps = round(users / step)
+            step_time = total_run_time / no_of_steps
+            user_count = min(step * math.ceil(run_time / step_time), users)
+            return user_count, step
+        return None

chainbench/user/http.py

@@ -58,11 +58,6 @@ def check_fatal(self, response: ResponseContextManager) -> None:
             self.logger.critical(f"Redirect error: {response.url}")
 
     def check_http_error(self, response: ResponseContextManager) -> None:
-        if response.request is not None:
-            self.logger.debug(f"Request: {response.request.method} {response.request.url_split}")
-            if response.request.body is not None:
-                self.logger.debug(f"{response.request.body}")
-
         """Check the response for errors."""
         if response.status_code != 200:
             self.logger.error(f"Request failed with {response.status_code} code")

chainbench/user/jsonrpc.py

@@ -1,3 +1,4 @@
+import logging
 import random
 import typing as t
 
@@ -94,7 +95,7 @@ def make_rpc_call(
         rpc_call: RpcCall | None = None,
         method: str | None = None,
         params: list[t.Any] | dict | None = None,
-        name: str = "",
+        name: str | None = None,
         path: str = "",
     ) -> None:
         """Make a JSON-RPC call."""
@@ -103,15 +104,18 @@ def make_rpc_call(
                 raise ValueError("Either rpc_call or method must be provided")
             else:
                 rpc_call = RpcCall(method, params)
-                name = method
-        else:
+        if name is None:
             name = rpc_call.method
 
         with self.client.request(
             "POST", self.rpc_path + path, json=rpc_call.request_body(), name=name, catch_response=True
         ) as response:
             self.check_http_error(response)
             self.check_json_rpc_response(response, name=name)
+            if logging.getLogger("locust").level == logging.DEBUG:
+                self.logger.debug(f"jsonrpc: {rpc_call.method} - params: {rpc_call.params}, response: {response.text}")
+            else:
+                self.logger.info(f"jsonrpc: {rpc_call.method} - params: {rpc_call.params}")
 
     def make_batch_rpc_call(self, rpc_calls: list[RpcCall], name: str = "", path: str = "") -> None:
         """Make a Batch JSON-RPC call."""

chainbench/user/protocol/ethereum.py

@@ -110,7 +110,10 @@ def task_to_method(task_name: str) -> str:
 
     @classmethod
     def method_to_task_function(cls, method: str) -> t.Callable:
-        return getattr(cls, f"{method}_task")
+        if not method.endswith("task"):
+            return getattr(cls, f"{method}_task")
+        else:
+            return getattr(cls, method)
 
     def eth_v1_beacon_states_head_fork_task(self):
         self.get(

chainbench/util/cli.py

@@ -30,10 +30,10 @@ def get_profile_path(base_path: Path, profile: str) -> Path:
 
 def get_profiles(profile_dir: Path) -> list[str]:
     """Get list of profiles in given directory."""
-    from locust.argument_parser import find_locustfiles
+    from locust.argument_parser import parse_locustfile_paths
 
     result = []
-    for locustfile in find_locustfiles([profile_dir.__str__()], True):
+    for locustfile in parse_locustfile_paths([profile_dir.__str__()]):
         locustfile_path = Path(locustfile).relative_to(profile_dir)
         if locustfile_path.parent.__str__() != ".":
             result.append(".".join(locustfile_path.parts[:-1]) + "." + locustfile_path.parts[-1][:-3])
@@ -121,6 +121,7 @@ class LocustOptions:
     custom_tags: list[str]
     exclude_tags: list[str]
     target: str
+    shape_path: Path | None = None
     headless: bool = False
     timescale: bool = False
     pg_host: str | None = None
@@ -136,8 +137,9 @@ class LocustOptions:
 
     def get_master_command(self) -> str:
         """Generate master command."""
+        profile_args = f"{self.profile_path},{self.shape_path}" if self.shape_path else self.profile_path
         command = (
-            f"locust -f {self.profile_path} --master "
+            f"locust -f {profile_args} --master "
             f"--master-bind-host {self.host} --master-bind-port {self.port} "
             f"--web-host {self.host} "
             f"-u {self.users} -r {self.spawn_rate} --run-time {self.test_time} "
@@ -154,8 +156,9 @@ def get_master_command(self) -> str:
 
     def get_worker_command(self, worker_id: int = 0) -> str:
         """Generate worker command."""
+        profile_args = f"{self.profile_path},{self.shape_path}" if self.shape_path else self.profile_path
         command = (
-            f"locust -f {self.profile_path} --worker --master-host {self.host} --master-port {self.port} "
+            f"locust -f {profile_args} --worker --master-host {self.host} --master-port {self.port} "
             f"--logfile {self.results_path}/worker_{worker_id}.log --loglevel {self.log_level} --stop-timeout 30"
         )
         return self.get_extra_options(command)

chainbench/util/http.py

@@ -49,7 +49,6 @@ def json(self) -> dict[str, t.Any]:
             logger.error("Response is not json: %s", self.content)
             raise
         else:
-            logger.debug(f"Response: {self.content}")
             return data
 
     def check_http_error(self, request_uri: str = "", error_level: HttpErrorLevel = HttpErrorLevel.ClientError) -> None: