Implement sweat.claim load in Locust

This PR introduces sweat.claim contract necessary for #11107 and uses it
to implement sweat::defer_batch and
sweat.claim::get_claimable_balance_for_account calls.

I originally intended to use sweat.claim::claim but it requires the
balance to be locked for some time to be claimed, so to avoid changing
sweat.claim contract, I've used get_claimable_balance_for_account method
instead, which does a similar amount of work and can be called
unconditionally.

pytest/tests/loadtest/locust/README.md

@@ -129,6 +129,7 @@ Currently supported load types:
 | Congestion | congestion.py | (`--congestion-wasm $WASM_PATH`) | Creates a single instance of Congestion contract. Users run large and long transactions. |
 | Sweat (normal load) | sweat.py | (`--sweat-wasm $WASM_PATH`) | Creates a single instance of the SWEAT contract. A mix of FT transfers and batch minting with batch sizes comparable to mainnet observations in summer 2023. |
 | Sweat (storage stress test) | sweat.py | `--tags=storage-stress-test` <br> (`--sweat-wasm $WASM_PATH`) | Creates a single instance of the SWEAT contract. Sends maximally large batches to mint more tokens, thereby touching many storage nodes per receipt. This load will take a while to initialize enough Sweat users on chain. |
+| Sweat (claim) | sweat.py | `--tags=claim-test` <br> (`--sweat-wasm $WASM_PATH`) <br> (`--sweat-claim-wasm $WASM_PATH`) | Creates a single instance of the SWEAT and SWEAT.CLAIM contract. Sends deferred batches to mint more tokens, thereby touching many storage nodes per receipt. Then calls balance checks that iterate through populated state. |
 
 ## Notes on Storage Stress Test
 

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

@@ -1,6 +1,6 @@
 import typing
 from common.ft import FTContract, InitFTAccount
-from common.base import Account, NearNodeProxy, NearUser, FunctionCall, MultiFunctionCall, INIT_DONE
+from common.base import Account, Deploy, NearNodeProxy, NearUser, FunctionCall, MultiFunctionCall, INIT_DONE
 import locust
 import sys
 import pathlib
@@ -16,9 +16,10 @@
 
 class SweatContract(FTContract):
 
-    def __init__(self, main_account: Account, oracle_account: Account,
-                 code: str):
+    def __init__(self, main_account: Account, claim_account: Account,
+                 oracle_account: Account, code: str):
         super().__init__(main_account, oracle_account, code)
+        self.claim = claim_account
         self.oracle = oracle_account
 
     def init_contract(self, node: NearNodeProxy):
@@ -79,10 +80,20 @@ def args(self) -> dict:
         return {"postfix": None}
 
 
+class InitClaim(FunctionCall):
+
+    def __init__(self, claim_account: Account, token_account_id: str):
+        super().__init__(claim_account, claim_account.key.account_id, "init")
+        self.token_account_id = token_account_id
+
+    def args(self) -> dict:
+        return {"token_account_id": self.token_account_id}
+
+
 class SweatAddOracle(FunctionCall):
     """
     Oracle accounts are allowed to mint new tokens and can only be added by the
-    account id  of the contract itself.
+    account id of the contract itself.
     """
 
     def __init__(self, sweat_account: Account, oracle_id: str):
@@ -96,7 +107,7 @@ def args(self) -> dict:
 
 class SweatMint(FunctionCall):
     """
-    A call to `tge_mint`.
+    A call to `sweat::tge_mint`.
     Token Generation Event (TGE) was day 0 when SWEAT launched.
     This is the transaction to get initial balance into accounts.
     """
@@ -115,7 +126,7 @@ def args(self) -> dict:
 
 class SweatMintBatch(MultiFunctionCall):
     """
-    A call to `record_batch`.
+    A call to `sweat::record_batch`.
     Mints new tokens for walked steps for a batch of users.
     Might get split into multiple function calls to avoid log output limits.
     """
@@ -137,6 +148,43 @@ def args(self) -> typing.List[dict]:
         return [{"steps_batch": chunk} for chunk in chunks]
 
 
+class SweatDeferBatch(FunctionCall):
+    """
+    A call to `sweat::defer_batch`.
+    """
+
+    def __init__(self, sweat_id: str, oracle: Account, holding_account_id: str,
+                 steps_batch: typing.List[RecipientSteps]):
+        super().__init__(oracle, sweat_id, "defer_batch")
+        self.holding_account_id = holding_account_id
+        self.steps_batch = steps_batch
+
+    def args(self) -> dict:
+        return {
+            "holding_account_id": self.holding_account_id,
+            "steps_batch": self.steps_batch,
+        }
+
+
+class SweatGetClaimableBalanceForAccount(FunctionCall):
+    """
+    A call to `sweat.claim::get_claimable_balance_for_account`.
+
+    We use it instead of `sweat.claim::claim` as it does not require to wait for funds to become
+    claimable and performs similar amount of computation.
+    """
+
+    def __init__(self, sweat_claim_id: str, user: Account, account_id: str):
+        super().__init__(user, sweat_claim_id,
+                         "get_claimable_balance_for_account")
+        self.account_id = account_id
+
+    def args(self) -> dict:
+        return {
+            "account_id": self.account_id,
+        }
+
+
 @events.init.add_listener
 def on_locust_init(environment, **kwargs):
     INIT_DONE.wait()
@@ -146,23 +194,40 @@ def on_locust_init(environment, **kwargs):
 
     funding_account = NearUser.funding_account
     funding_account.refresh_nonce(node.node)
-    sweat_contract_code = environment.parsed_options.sweat_wasm
     sweat_account_id = f"sweat{run_id}.{environment.master_funding_account.key.account_id}"
+    sweat_claim_account_id = f"sweat-claim{run_id}.{environment.master_funding_account.key.account_id}"
     oracle_account_id = worker_oracle_id(worker_id, run_id,
                                          environment.master_funding_account)
 
     sweat_account = Account(key.Key.from_seed_testonly(sweat_account_id))
+    sweat_claim_account = Account(
+        key.Key.from_seed_testonly(sweat_claim_account_id))
     oracle_account = Account(key.Key.from_seed_testonly(oracle_account_id))
 
-    environment.sweat = SweatContract(sweat_account, oracle_account,
-                                      sweat_contract_code)
+    environment.sweat = SweatContract(sweat_account, sweat_claim_account,
+                                      oracle_account,
+                                      environment.parsed_options.sweat_wasm)
 
     # Create Sweat contract, unless we are a worker, in which case the master already did it
     if not isinstance(environment.runner, locust.runners.WorkerRunner):
         node.prepare_account(oracle_account, environment.master_funding_account,
-                             FTContract.INIT_BALANCE, "create contract account")
+                             FTContract.INIT_BALANCE, "create oracle account")
         environment.sweat.install(node, environment.master_funding_account)
 
+        existed = node.prepare_account(sweat_claim_account, funding_account,
+                                       100000, "create contract account")
+        if not existed:
+            node.send_tx_retry(
+                Deploy(sweat_claim_account,
+                       environment.parsed_options.sweat_claim_wasm,
+                       "Sweat-Claim"), "deploy Sweat-Claim contract")
+            node.send_tx_retry(InitClaim(sweat_claim_account, sweat_account_id),
+                               "init Sweat-Claim contract")
+            # `sweat` account also must be registered as oracle to call methods on `sweat.claim`.
+            node.send_tx_retry(
+                SweatAddOracle(sweat_claim_account, sweat_account_id),
+                "add sweat.claim oracle")
+
     # on master, register oracles for workers
     if isinstance(environment.runner, locust.runners.MasterRunner):
         num_oracles = int(environment.parsed_options.max_workers)
@@ -178,8 +243,10 @@ def on_locust_init(environment, **kwargs):
                               "create contract account")
         for oracle in oracle_accounts:
             id = oracle.key.account_id
-            environment.sweat.register_oracle(node, id)
             environment.sweat.top_up(node, id)
+            environment.sweat.register_oracle(node, id)
+            node.send_tx_retry(SweatAddOracle(sweat_claim_account, id),
+                               "add sweat.claim oracle")
 
 
 def worker_oracle_id(worker_id, run_id, funding_account):
@@ -191,3 +258,6 @@ def _(parser):
     parser.add_argument("--sweat-wasm",
                         default="res/sweat.wasm",
                         help="Path to the compiled Sweat contract")
+    parser.add_argument("--sweat-claim-wasm",
+                        default="res/sweat_claim.wasm",
+                        help="Path to the compiled Sweat-Claim contract")

pytest/tests/loadtest/locust/download_contracts.sh

@@ -14,5 +14,6 @@ SOURCE_CONTRACTS_DIR="${SCRIPT_DIR}/../../../../runtime/near-test-contracts/res"
 
 wget https://raw.githubusercontent.com/NearSocial/social-db/master/res/social_db_release.wasm -O ${TARGET_CONTRACTS_DIR}/social_db.wasm
 wget https://raw.githubusercontent.com/sweatco/sweat-near/main/res/sweat.wasm -O ${TARGET_CONTRACTS_DIR}/sweat.wasm
+wget https://raw.githubusercontent.com/sweatco/sweat-near/main/res/sweat_claim.wasm -O ${TARGET_CONTRACTS_DIR}/sweat_claim.wasm
 ln -s ${SOURCE_CONTRACTS_DIR}/fungible_token.wasm ${TARGET_CONTRACTS_DIR}/fungible_token.wasm
 ln -s ${SOURCE_CONTRACTS_DIR}/backwards_compatible_rs_contract.wasm ${TARGET_CONTRACTS_DIR}/congestion.wasm

pytest/tests/loadtest/locust/locustfiles/sweat.py

@@ -11,11 +11,10 @@
   - Periodic batches that adds steps (mints new tokens)
 """
 
-from common.sweat import RecipientSteps, SweatContract, SweatMintBatch
+from common.sweat import RecipientSteps, SweatGetClaimableBalanceForAccount, SweatMintBatch, SweatDeferBatch
 from common.ft import TransferFT
 from common.base import Account, AddFullAccessKey, NearUser
 from locust import between, tag, task
-import copy
 import logging
 import pathlib
 import random
@@ -29,20 +28,12 @@
 logger = new_logger(level=logging.WARN)
 
 
-class SweatUser(NearUser):
+class SweatOracle(NearUser):
     """
-    Registers itself on an FT contract in the setup phase, then just sends Sweat to
-    random users.
-
-    Also includes a task to mint and distribute tokens in batches.
+    Registers itself as an oracle in the setup phase, then records Sweat minting events.
     """
     wait_time = between(1, 3)  # random pause between transactions
-
-    @task(3)
-    def ft_transfer(self):
-        receiver = self.sweat.random_receiver(self.account_id)
-        tx = TransferFT(self.sweat.account, self.account, receiver)
-        self.send_tx(tx, locust_name="Sweat transfer")
+    fixed_count = 1  # Oracle can do only one operation at a time.
 
     @task(1)
     def record_single_batch(self):
@@ -83,6 +74,22 @@ def record_batch_of_large_batches(self):
             [[account_id, rng.randint(1000, 3000)] for account_id in receivers])
         self.send_tx(tx, locust_name="Sweat record batch (stress test)")
 
+    @tag("claim-test")
+    @task
+    def defer_batch(self):
+        rng = random.Random()
+        # just around the log limit
+        batch_size = min(rng.randint(100, 150),
+                         len(self.sweat.registered_users))
+        receivers = self.sweat.random_receivers(self.account_id, batch_size)
+        tx = SweatDeferBatch(
+            self.sweat.account.key.account_id, self.oracle,
+            self.sweat.claim.key.account_id, [
+                RecipientSteps(account_id, steps=rng.randint(1000, 3000))
+                for account_id in receivers
+            ])
+        self.send_tx(tx, locust_name="Sweat defer batch")
+
     def on_start(self):
         super().on_start()
         # We have one oracle account per worker. Sharing a single access key
@@ -94,10 +101,39 @@ def on_start(self):
         user_oracle_key = key.Key.from_random(oracle.key.account_id)
         self.send_tx_retry(AddFullAccessKey(oracle, user_oracle_key),
                            "add user key to oracle")
-
         self.oracle = Account(user_oracle_key)
         self.oracle.refresh_nonce(self.node.node)
 
+        logger.debug(
+            f"{self.account_id} ready to use Sweat contract {self.sweat.account.key.account_id}"
+        )
+
+
+class SweatUser(NearUser):
+    """
+    Registers itself on an FT contract in the setup phase, then does one of two things:
+      - Sends Sweat to random users.
+      - Checks its' balance in sweat.claim contract.
+    """
+    wait_time = between(1, 3)  # random pause between transactions
+
+    @task(3)
+    def ft_transfer(self):
+        receiver = self.sweat.random_receiver(self.account_id)
+        tx = TransferFT(self.sweat.account, self.account, receiver)
+        self.send_tx(tx, locust_name="Sweat transfer")
+
+    @tag("claim-test")
+    @task
+    def get_claimable_balance_for_account(self):
+        tx = SweatGetClaimableBalanceForAccount(self.sweat.claim.key.account_id,
+                                                self.account,
+                                                self.account.key.account_id)
+        self.send_tx(tx, locust_name="Sweat.claim get balance")
+
+    def on_start(self):
+        super().on_start()
+        self.sweat = self.environment.sweat
         self.sweat.register_user(self)
         logger.debug(
             f"{self.account_id} ready to use Sweat contract {self.sweat.account.key.account_id}"