feat: swarmable FT loadtest (near#9111)

To make the locust-based loadtest swarmable:
- Create N accounts with deterministic name and keys
- N workers pick one each to fund their users

Note that the naive approach to share the same funding account across workers doesn't work. They will constantly invalidate each others transactions due to nonces.

This commit also moves some code out of `locustfile.py` to `base.py` and `ft.py` for a better separation of concerns.

And it includes small fixes/improvements:
- allow multiple FT contracts per worker
- refresh_nonce sometimes raises a `KeyError` when the account was just created, catch it crudely in a try-catch retry loop
- Provide `locust_name` arguments to show better statistics on requests

pytest/tests/loadtest/locust/README.md

@@ -17,7 +17,7 @@ For a local test setup, this works just fine.
 # This assumes your shell is are in nearcore directory
 CONTRACT="${PWD}/runtime/near-test-contracts/res/fungible_token.wasm"
 # This assumes you are running against localnet
-KEY=KEY=~/.near/test0/validator_key.json
+KEY=~/.near/localnet/node0/validator_key.json
 ```
 
 For a quick demo, you can also run a localnet using [nearup](https://github.com/near/nearup).
@@ -30,9 +30,54 @@ Then to actually run it, this is the command. (Update ports and IP according to
 cd pytest/tests/loadtest/locust/
 locust -H 127.0.0.1:3030 \
   --fungible-token-wasm=$CONTRACT \
-  --contract-key=$KEY
+  --funding-key=$KEY
 ```
 
 This will print a link to a web UI where the loadtest can be started and statistics & graphs can be observed.
 But if you only want the pure numbers, it can also be run headless.
 (add `--headless` and something like `-t 120s -u 1000` to the command to specify total time and max concurrent users)
+
+## Running with full load (multi-threaded or even multi-machine)
+
+Each locust process will only use a single core. This leads to unwanted
+throttling on the load generator side if you try to use more than a couple
+hundred of users.
+
+Luckily, locust has the ability to swarm the load generation across many processes.
+To use it, start one process with the `--master` argument and as many as you
+like with `--worker`. (If they run on different machines, you also need to
+provide `--master-host` and `--master-port`, if running on the same machine it
+will work automagically.)
+
+Start the master
+```sh
+locust -H 127.0.0.1:3030 \
+  --fungible-token-wasm=$CONTRACT \
+  --funding-key=$KEY \
+  --master
+```
+
+On the worker
+```sh
+# Increase soft limit of open files for the current OS user
+# (each locust user opens a separate socket = file)
+ulimit -S -n 100000
+# Run the worker, key must include the same account id as used on master
+locust -H 127.0.0.1:3030 \
+  --fungible-token-wasm=$CONTRACT \
+  --funding-key=$KEY \
+  --worker
+```
+
+Spawning N workers in a single shell:
+```sh
+for i in {1..16}
+do
+   locust -H 127.0.0.1:3030 \
+    --fungible-token-wasm=$CONTRACT \
+    --funding-key=$KEY \
+    --worker &
+done
+```
+
+Use `pkill -P $$` to stop all workers.

pytest/tests/loadtest/locust/common/base.py

@@ -1,3 +1,10 @@
+from configured_logger import new_logger
+from locust import HttpUser, events, runners
+import utils
+import mocknet_helpers
+import key
+import transaction
+import cluster
 import base64
 import json
 import base58
@@ -10,28 +17,32 @@
 import requests
 import sys
 import time
+import retry
 
 sys.path.append(str(pathlib.Path(__file__).resolve().parents[4] / 'lib'))
 
-import cluster
-import transaction
-import key
-import mocknet_helpers
-import utils
-
-from locust import HttpUser, events
-from configured_logger import new_logger
-
 DEFAULT_TRANSACTION_TTL_SECONDS = 20
 logger = new_logger(level=logging.WARN)
 
 
+def is_key_error(exception):
+    return isinstance(exception, KeyError)
+
+
 class Account:
 
     def __init__(self, key):
         self.key = key
         self.current_nonce = multiprocessing.Value(ctypes.c_ulong, 0)
 
+    # Race condition: maybe the account was created but the RPC interface
+    # doesn't display it yet, which returns an empty result and raises a
+    # `KeyError`.
+    # (not quite sure how this happens but it did happen to me on localnet)
+    @retry(wait_exponential_multiplier=500,
+           wait_exponential_max=10000,
+           stop_max_attempt_number=5,
+           retry_on_exception=is_key_error)
     def refresh_nonce(self, node):
         with self.current_nonce.get_lock():
             self.current_nonce.value = mocknet_helpers.get_nonce_for_key(
@@ -110,8 +121,7 @@ def sign_and_serialize(self, block_hash):
 class NearUser(HttpUser):
     abstract = True
     id_counter = 0
-    # initialized in `on_locust_init`
-    funding_account = None
+    INIT_BALANCE = 100.0
 
     @classmethod
     def get_next_id(cls):
@@ -121,6 +131,7 @@ def get_next_id(cls):
     @classmethod
     def account_id(cls, id):
         # Pseudo-random 6-digit prefix to spread the users in the state tree
+        # TODO: Also make sure these are spread evenly across shards
         prefix = str(hash(str(id)))[-6:]
         return f"{prefix}_user{id}.{cls.funding_account.key.account_id}"
 
@@ -139,10 +150,11 @@ def on_start(self):
         self.send_tx_retry(
             CreateSubAccount(NearUser.funding_account,
                              self.account.key,
-                             balance=5000.0))
+                             balance=NearUser.INIT_BALANCE))
+
         self.account.refresh_nonce(self.node)
 
-    def send_tx(self, tx: Transaction):
+    def send_tx(self, tx: Transaction, locust_name="generic send_tx"):
         """
         Send a transaction and return the result, no retry attempted.
         """
@@ -160,14 +172,12 @@ def send_tx(self, tx: Transaction):
             "jsonrpc": "2.0"
         }
 
-        # with self.node.session.post(
         # This is tracked by locust
-        with self.client.post(
-                url="http://%s:%s" % self.node.rpc_addr(),
-                json=j,
-                timeout=DEFAULT_TRANSACTION_TTL_SECONDS,
-                catch_response=True,
-        ) as response:
+        with self.client.post(url="http://%s:%s" % self.node.rpc_addr(),
+                              json=j,
+                              timeout=DEFAULT_TRANSACTION_TTL_SECONDS,
+                              catch_response=True,
+                              name=locust_name) as response:
             try:
                 rpc_result = json.loads(response.content)
                 tx_result = evaluate_rpc_result(rpc_result)
@@ -180,7 +190,9 @@ def send_tx(self, tx: Transaction):
                 response.failure(err.message)
         return response
 
-    def send_tx_retry(self, tx: Transaction):
+    def send_tx_retry(self,
+                      tx: Transaction,
+                      locust_name="generic send_tx_retry"):
         """
         Send a transaction and retry until it succeeds
         """
@@ -190,7 +202,7 @@ def send_tx_retry(self, tx: Transaction):
         # error, as long as it is one defined by us (inherits from NearError)
         while True:
             try:
-                result = self.send_tx(tx)
+                result = self.send_tx(tx, locust_name=locust_name)
                 return result
             except NearError as error:
                 logger.warn(
@@ -293,14 +305,55 @@ def evaluate_rpc_result(rpc_result):
 # called once per process before user initialization
 @events.init.add_listener
 def on_locust_init(environment, **kwargs):
-    funding_key = key.Key.from_json_file(environment.parsed_options.funding_key)
-    NearUser.funding_account = Account(funding_key)
+    # Note: These setup requests are not tracked by locust because we use our own http session
+    host, port = environment.host.split(":")
+    node = cluster.RpcNode(host, port)
 
+    master_funding_key = key.Key.from_json_file(
+        environment.parsed_options.funding_key)
+    master_funding_account = Account(master_funding_key)
 
-# CLI args
+    funding_account = None
+    # every worker needs a funding account to create its users, eagerly create them in the master
+    if isinstance(environment.runner, runners.MasterRunner):
+        num_funding_accounts = environment.parsed_options.max_workers
+        funding_balance = 10000 * NearUser.INIT_BALANCE
+        # TODO: Create accounts in parallel
+        for id in range(num_funding_accounts):
+            account_id = f"funds_worker_{id}.{master_funding_account.key.account_id}"
+            worker_key = key.Key.from_seed_testonly(account_id, account_id)
+            logger.info(f"Creating {account_id}")
+            send_transaction(
+                node,
+                CreateSubAccount(master_funding_account,
+                                 worker_key,
+                                 balance=funding_balance))
+        funding_account = master_funding_account
+    elif isinstance(environment.runner, runners.WorkerRunner):
+        worker_id = environment.runner.worker_index
+        worker_account_id = f"funds_worker_{worker_id}.{master_funding_account.key.account_id}"
+        worker_key = key.Key.from_seed_testonly(worker_account_id,
+                                                worker_account_id)
+        funding_account = Account(worker_key)
+    elif isinstance(environment.runner, runners.LocalRunner):
+        funding_account = master_funding_account
+    else:
+        raise SystemExit(
+            f"unexpected runner class {environment.runner.__class__.__name__}")
+
+    NearUser.funding_account = funding_account
+
+
+# Add custom CLI args here, will be available in `environment.parsed_options`
 @events.init_command_line_parser.add_listener
 def _(parser):
     parser.add_argument(
         "--funding-key",
         required=True,
         help="account to use as source of NEAR for account creation")
+    parser.add_argument(
+        "--max-workers",
+        type=int,
+        required=False,
+        default=16,
+        help="How many funding accounts to generate for workers")

pytest/tests/loadtest/locust/common/ft.py

@@ -1,18 +1,66 @@
 import json
+import random
 import sys
 import pathlib
+from locust import events
 
 sys.path.append(str(pathlib.Path(__file__).resolve().parents[4] / 'lib'))
 
+import cluster
+import key
 import transaction
 
 from account import TGAS
-from common.base import Transaction
+from common.base import Account, CreateSubAccount, Deploy, NearUser, send_transaction, Transaction
+
+
+class FTContract:
+    # NEAR balance given to contracts, doesn't have to be much since users are
+    # going to pay for storage
+    INIT_BALANCE = NearUser.INIT_BALANCE
+
+    def __init__(self, account: Account, code):
+        self.account = account
+        self.registered_users = []
+        self.code = code
+
+    def install(self, node):
+        """
+        Deploy and initialize the contract on chain.
+        The account should already exist at this point.
+        """
+        self.account.refresh_nonce(node)
+        send_transaction(node, Deploy(self.account, self.code, "FT"))
+        send_transaction(node, InitFT(self.account))
+
+    def register_user(self, user: NearUser):
+        user.send_tx(InitFTAccount(self.account, user.account),
+                     locust_name="Init FT Account")
+        user.send_tx(TransferFT(self.account,
+                                self.account,
+                                user.account_id,
+                                how_much=1E8),
+                     locust_name="FT Funding")
+        self.registered_users.append(user.account_id)
+
+    def random_receiver(self, sender: str) -> str:
+        rng = random.Random()
+        receiver = rng.choice(self.registered_users)
+        # Sender must be != receiver but maybe there is no other registered user
+        # yet, so we just send to the contract account which is registered
+        # implicitly from the start
+        if receiver == sender:
+            receiver = self.account.key.account_id
+        return receiver
 
 
 class TransferFT(Transaction):
 
-    def __init__(self, ft, sender, recipient_id, how_much=1):
+    def __init__(self,
+                 ft: Account,
+                 sender: Account,
+                 recipient_id: str,
+                 how_much=1):
         super().__init__()
         self.ft = ft
         self.sender = sender
@@ -39,7 +87,7 @@ def sign_and_serialize(self, block_hash):
 
 class InitFT(Transaction):
 
-    def __init__(self, contract):
+    def __init__(self, contract: Account):
         super().__init__()
         self.contract = contract
 
@@ -60,7 +108,7 @@ def sign_and_serialize(self, block_hash):
 
 class InitFTAccount(Transaction):
 
-    def __init__(self, contract, account):
+    def __init__(self, contract: Account, account: Account):
         super().__init__()
         self.contract = contract
         self.account = account
@@ -75,3 +123,54 @@ def sign_and_serialize(self, block_hash):
                                                  int(3E14), int(1E23),
                                                  account.use_nonce(),
                                                  block_hash)
+
+
+@events.init.add_listener
+def on_locust_init(environment, **kwargs):
+    # Note: These setup requests are not tracked by locust because we use our own http session
+    host, port = environment.host.split(":")
+    node = cluster.RpcNode(host, port)
+
+    ft_contract_code = environment.parsed_options.fungible_token_wasm
+    num_ft_contracts = environment.parsed_options.num_ft_contracts
+    funding_account = NearUser.funding_account
+    parent_id = funding_account.key.account_id
+    worker_id = getattr(environment.runner, "worker_id", "local_id")
+
+    funding_account.refresh_nonce(node)
+
+    environment.ft_contracts = []
+    # TODO: Create accounts in parallel
+    for i in range(num_ft_contracts):
+        # Prefix that makes accounts unique across workers
+        # Shuffling with a hash avoids locality in the state trie.
+        # TODO: Also make sure these are spread evenly across shards
+        prefix = str(hash(str(worker_id) + str(i)))[-6:]
+        contract_key = key.Key.from_random(f"{prefix}_ft.{parent_id}")
+        ft_account = Account(contract_key)
+        send_transaction(
+            node,
+            CreateSubAccount(funding_account,
+                             ft_account.key,
+                             balance=FTContract.INIT_BALANCE))
+
+        ft_contract = FTContract(ft_account, ft_contract_code)
+        ft_contract.install(node)
+        environment.ft_contracts.append(ft_contract)
+
+
+# FT specific CLI args
+@events.init_command_line_parser.add_listener
+def _(parser):
+    parser.add_argument("--fungible-token-wasm",
+                        type=str,
+                        required=True,
+                        help="Path to the compiled Fungible Token contract")
+    parser.add_argument(
+        "--num-ft-contracts",
+        type=int,
+        required=False,
+        default=4,
+        help=
+        "How many different FT contracts to spawn from this worker (FT contracts are never shared between workers)"
+    )