Str-961: Introduce load testing env in fn-tests. (#639)

* fntests: introduce locust as a load testing tool.

* Drafty structure for the load testing.

* Disable lint for entry.py because of monkey-patching.

* Minor adjustments and comments.

* Refactor, prettify, lints, docstings.

functional-tests/entry.py

@@ -1,4 +1,10 @@
 #!/usr/bin/env python3
+from gevent import monkey
+
+# This is important for locust to work with flexitest.
+# Because of this line, ruff linter is disabled for the whole file :(
+# Currently, it's not possible to disable ruff for the block of code.
+monkey.patch_all()
 
 import argparse
 import os
@@ -10,6 +16,8 @@
 from factory import factory
 from utils import *
 from utils.constants import *
+from load.cfg import RethLoadConfigBuilder
+from load.jobs import EthJob
 
 TEST_DIR: str = "tests"
 
@@ -19,6 +27,13 @@
 parser.add_argument("-t", "--tests", nargs="*", help="Define individual tests to execute")
 
 
+def disabled_tests() -> list[str]:
+    """
+    Helper to disable some tests.
+    Useful during debugging or when the test becomes flaky.
+    """
+    return frozenset(["basic_load"])
+
 def filter_tests(parsed_args, modules):
     """
     Filters test modules against parsed args supplied from the command line.
@@ -31,6 +46,7 @@ def filter_tests(parsed_args, modules):
     )
 
     filtered = dict()
+    disabled = disabled_tests()
     for test, path in modules.items():
         # Drop the prefix of the path before TEST_DIR
         test_path_parts = os.path.normpath(path).split(os.path.sep)
@@ -41,8 +57,9 @@ def filter_tests(parsed_args, modules):
         test_groups = frozenset(test_path_parts[:-1])
 
         # Filtering logic:
-        # if groups or tests were specified (non-empty) as args, then check for exclusion
-        take = True
+        # - check if the test is currently disabled
+        # - if groups or tests were specified (non-empty) as args, then check for exclusion.
+        take = test not in disabled
         if arg_groups and not (arg_groups & test_groups):
             take = False
         if arg_tests and test not in arg_tests:
@@ -53,7 +70,6 @@ def filter_tests(parsed_args, modules):
 
     return filtered
 
-
 def main(argv):
     """
     The main entrypoint for running functional tests.
@@ -72,6 +88,7 @@ def main(argv):
     reth_fac = factory.RethFactory([12600 + i for i in range(100 * 3)])
     prover_client_fac = factory.ProverClientFactory([12900 + i for i in range(100 * 3)])
     bridge_client_fac = factory.BridgeClientFactory([13200 + i for i in range(100)])
+    load_gen_fac = factory.LoadGeneratorFactory([13300 + i for i in range(100)])
     seq_signer_fac = factory.StrataSequencerFactory()
 
     factories = {
@@ -82,8 +99,12 @@ def main(argv):
         "reth": reth_fac,
         "prover_client": prover_client_fac,
         "bridge_client": bridge_client_fac,
+        "load_generator": load_gen_fac,
     }
 
+    reth_load_env = testenv.LoadEnvConfig()
+    reth_load_env.with_load_builder(RethLoadConfigBuilder().with_jobs([EthJob]).with_rate(15))
+
     global_envs = {
         # Basic env is the default env for all tests.
         "basic": testenv.BasicEnvConfig(101),
@@ -99,6 +120,7 @@ def main(argv):
             2
         ),  # TODO: Need to generate at least horizon blocks, based on params
         "prover": testenv.BasicEnvConfig(101),
+        "load_reth": reth_load_env,
     }
 
     setup_root_logger()
@@ -109,6 +131,7 @@ def main(argv):
     results = rt.run_tests(tests)
     rt.save_json_file("results.json", results)
     flexitest.dump_results(results)
+    # TODO(load): dump load test stats into separate file.
 
     flexitest.fail_on_error(results)
 

functional-tests/envs/testenv.py

@@ -9,6 +9,7 @@
 )
 
 from envs.rollup_params_cfg import RollupConfig
+from load.cfg import LoadConfig, LoadConfigBuilder
 from utils import *
 from utils.constants import *
 
@@ -403,3 +404,29 @@ def init(self, ctx: flexitest.EnvContext) -> flexitest.LiveEnv:
             svcs[name] = br
 
         return BasicLiveEnv(svcs, bridge_pk, rollup_cfg)
+
+
+# TODO: Maybe, we need to make it dynamic to enhance any EnvConfig with load testing capabilities.
+class LoadEnvConfig(BasicEnvConfig):
+    _load_cfgs: list[Callable[[dict[str, flexitest.Service]], LoadConfig]] = []
+
+    def with_load_builder(self, builder: LoadConfigBuilder):
+        self._load_cfgs.append(builder)
+        return self
+
+    def init(self, ctx: flexitest.EnvContext) -> flexitest.LiveEnv:
+        basic_live_env = super().init(ctx)
+
+        if not self._load_cfgs:
+            raise Exception(
+                "LoadEnv has no load builders! Specify load builders or just use BasicEnv."
+            )
+
+        # Create load generator services for all the builders.
+        svcs = basic_live_env.svcs
+        load_fac = ctx.get_factory("load_generator")
+        for builder in self._load_cfgs:
+            load_cfg: LoadConfig = builder(svcs)
+            svcs[f"load_generator.{builder.name}"] = load_fac.create_simple_loadgen(load_cfg)
+
+        return BasicLiveEnv(svcs, basic_live_env._bridge_pk, basic_live_env._rollup_cfg)

functional-tests/factory/factory.py

@@ -7,6 +7,8 @@
 from bitcoinlib.services.bitcoind import BitcoindClient
 
 from factory import seqrpc
+from load.cfg import LoadConfig
+from load.service import LoadGeneratorService
 from utils import *
 from utils.constants import *
 
@@ -398,6 +400,30 @@ def create_operator(
         return svc
 
 
+class LoadGeneratorFactory(flexitest.Factory):
+    def __init__(self, port_range: list[int]):
+        super().__init__(port_range)
+
+    @flexitest.with_ectx("ctx")
+    def create_simple_loadgen(
+        self,
+        load_cfg: LoadConfig,
+        ctx: flexitest.EnvContext,
+    ) -> flexitest.Service:
+        name = "load_generator"
+
+        datadir = ctx.make_service_dir(name)
+        rpc_port = self.next_port()
+        logfile = os.path.join(datadir, "service.log")
+
+        rpc_url = f"ws://localhost:{rpc_port}"
+
+        svc = LoadGeneratorService(logfile, load_cfg)
+        svc.start()
+        _inject_service_create_rpc(svc, rpc_url, name)
+        return svc
+
+
 def _inject_service_create_rpc(svc: flexitest.service.ProcService, rpc_url: str, name: str):
     """
     Injects a `create_rpc` method using JSON-RPC onto a `ProcService`, checking

functional-tests/load/cfg.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+import flexitest
+
+from load.job import StrataLoadJob
+
+
+@dataclass(frozen=True)
+class LoadConfig:
+    """
+    Config for the load service.
+    """
+
+    jobs: list[StrataLoadJob]
+    """A set of jobs that emit the load towards the host."""
+
+    host: str
+    """The host that will accept load requests."""
+
+    spawn_rate: int
+    """The rate at which all the jobs will emit the load."""
+
+
+class LoadConfigBuilder:
+    """
+    An abstract builder of the `LoadConfig`.
+    """
+
+    jobs: list[StrataLoadJob] = []
+    """A set of jobs that emit the load towards the host."""
+
+    spawn_rate: int = 10
+    """The rate at which all the jobs will emit the load."""
+
+    service_name: str | None = None
+    """The name of the service to emit the load."""
+
+    def __init__(self):
+        if not self.service_name:
+            raise Exception("LoadConfigBuilder: missing service_name attribute.")
+
+    def with_jobs(self, jobs: list[StrataLoadJob]):
+        self.jobs.extend(jobs)
+        return self
+
+    def with_rate(self, rate: int):
+        self.spawn_rate = rate
+        return self
+
+    def __call__(self, svcs) -> LoadConfig:
+        if not self.jobs:
+            raise Exception("LoadConfigBuilder: load jobs list is empty")
+
+        host = self.host_url(svcs)
+        # Patch jobs by the host.
+        for job in self.jobs:
+            job.host = host
+
+        return LoadConfig(self.jobs, host, self.spawn_rate)
+
+    def host_url(self, _svcs: dict[str, flexitest.Service]) -> str:
+        raise NotImplementedError()
+
+    @property
+    def name(self):
+        return self.service_name
+
+
+class RethLoadConfigBuilder(LoadConfigBuilder):
+    service_name: str = "reth"
+    spawn_rate: int = 20
+
+    def host_url(self, svcs: dict[str, flexitest.Service]) -> str:
+        reth = svcs["reth"]
+        web3_port = reth.get_prop("eth_rpc_http_port")
+        return f"http://localhost:{web3_port}"

functional-tests/load/job.py

@@ -0,0 +1,65 @@
+import web3
+import web3.middleware
+from locust import HttpUser
+
+
+class StrataLoadJob(HttpUser):
+    """
+    A common layer for all the load jobs in the load tests.
+    """
+
+    pass
+
+
+# TODO(load): configure the structured logging as we do in the tests.
+class BaseRethLoadJob(StrataLoadJob):
+    fund_amount: int = 1_000_000_000_000_000_000_000  # 1000 ETH
+
+    def on_start(self):
+        root_w3, genesis_acc = self.w3_with_genesis_acc()
+        self._root_w3 = root_w3
+        self._genesis_acc = genesis_acc
+
+    def w3_with_genesis_acc(self):
+        """
+        Return w3 with prefunded "root" account as specified in the chain config.
+        """
+        return self._init_w3(
+            lambda w3: w3.eth.account.from_key(
+                "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
+            )
+        )
+
+    def w3_with_new_acc(self):
+        """
+        Return w3 with a fresh account.
+        Also, funds this account, so it's able to sign and send some txns.
+        """
+        w3, new_acc = self._init_w3(lambda w3: w3.eth.account.create())
+        self._fund_account(new_acc.address)
+
+        return w3, new_acc
+
+    def _init_w3(self, init):
+        # Reuse the http session by locust internals, so the stats are measured correctly.
+        w3 = web3.Web3(web3.Web3.HTTPProvider(self.host, session=self.client))
+        # Init the account according to lambda
+        account = init(w3)
+        # Set the account onto web3 and init the signing middleware.
+        w3.address = account.address
+        w3.middleware_onion.add(web3.middleware.SignAndSendRawMiddlewareBuilder.build(account))
+
+        return w3, account
+
+    def _fund_account(self, acc):
+        print(f"FUNDING ACCOUNT {acc}")
+        source = self._root_w3.address
+        tx_hash = self._root_w3.eth.send_transaction(
+            {"to": acc, "value": hex(self.fund_amount), "gas": hex(100000), "from": source}
+        )
+
+        tx_receipt = self._root_w3.eth.wait_for_transaction_receipt(tx_hash, timeout=120)
+        print(f"FUNDING SUCCESS: {tx_receipt}")
+
+    def _balance(self, acc):
+        return self._root_w3.eth.get_balance(acc)

functional-tests/load/jobs/__init__.py

@@ -0,0 +1 @@
+from .reth import *

functional-tests/load/jobs/reth.py

@@ -0,0 +1,53 @@
+from locust import task
+
+from load.job import BaseRethLoadJob
+
+
+class EthJob(BaseRethLoadJob):
+    def on_start(self):
+        super().on_start()
+
+        w3, new_acc = self.w3_with_new_acc()
+        self.w3 = w3
+        self.new_acc = new_acc
+
+        b = self._balance(new_acc.address)
+        print(f"BALANCE AFTER START: {b}")
+
+    @task
+    def get_block(self):
+        print("GET_BLOCK REQUEST")
+        for i in range(1, 10):
+            block = self.w3.eth.get_block(hex(i))
+            num = block["number"]
+            txn_cnt = len(block["transactions"])
+            hash = block["hash"]
+            print(f"BLOCK DATA \t\t\t\t\t {hash}, {num}, {txn_cnt}")
+
+    @task
+    def block_num(self):
+        print("BLOCK_NUM REQUEST")
+        # Pure json-rpc without web3 with middleware.
+        method = "eth_blockNumber"
+        params = []
+        payload = {"jsonrpc": "2.0", "method": method, "params": params, "id": 1}
+        headers = {"Content-type": "application/json"}
+        # response = session.post(self.host, json=payload, headers=headers)
+        response = self.client.post("", json=payload, headers=headers)
+        # print(f"raw json response: {response.json()}")
+        print("BLOCK_NUMBER: {}".format(response.json()["result"]))
+
+    @task(5)
+    def send(self):
+        print("TRANSFER TRANSACTION")
+
+        source = self.w3.address
+        dest = self.w3.to_checksum_address("0x0000000000000000000000000000000000000001")
+        to_transfer = 1_000_000_000_000_000_000
+        try:
+            tx_hash = self.w3.eth.send_transaction(
+                {"to": dest, "value": hex(to_transfer), "gas": hex(100000), "from": source}
+            )
+            print(f"transfer transaction hash: {tx_hash}")
+        except Exception as e:
+            print(e)