feat(loadtest): allow rerunning on the same state (near#9231)

Previously the second run on the same state would fail due to the
funding accounts already existing.

Now one can choose between reusing old state (default) or starting
a new run with a new state (set a `--social-db-index` to a unique value)

Already existing accounts are simply not created, so things should work
fine as long as the initial balance per account doesn't run out.

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

@@ -242,6 +242,10 @@ def poll_tx_result(self, meta, params):
         # poll for tx result, using "EXPERIMENTAL_tx_status" which waits for
         # all receipts to finish rather than just the first one, as "tx" would do
         result_response = self.post_json("EXPERIMENTAL_tx_status", params)
+        # very verbose, but very useful to see what's happening when things are stuck
+        logger.debug(
+            f"polling, got: {result_response.status_code} {result_response.json()}"
+        )
 
         try:
             meta["response"] = evaluate_rpc_result(result_response.json())
@@ -250,6 +254,9 @@ def poll_tx_result(self, meta, params):
             meta["response"] = result_response.content
             raise
 
+    def account_exists(self, account_id: str) -> bool:
+        return "error" not in self.node.get_account(account_id, do_assert=False)
+
 
 class NearUser(User):
     abstract = True
@@ -281,10 +288,11 @@ def on_start(self):
         Called once per user, creating the account on chain
         """
         self.account = Account(key.Key.from_random(self.account_id))
-        self.send_tx_retry(
-            CreateSubAccount(NearUser.funding_account,
-                             self.account.key,
-                             balance=NearUser.INIT_BALANCE))
+        if not self.node.account_exists(self.account_id):
+            self.send_tx_retry(
+                CreateSubAccount(NearUser.funding_account,
+                                 self.account.key,
+                                 balance=NearUser.INIT_BALANCE))
 
         self.account.refresh_nonce(self.node.node)
 
@@ -419,11 +427,12 @@ def on_locust_init(environment, **kwargs):
         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)
-            node.send_tx_retry(
-                CreateSubAccount(master_funding_account,
-                                 worker_key,
-                                 balance=funding_balance),
-                "create funding account")
+            if not node.account_exists(account_id):
+                node.send_tx_retry(
+                    CreateSubAccount(master_funding_account,
+                                     worker_key,
+                                     balance=funding_balance),
+                    "create funding account")
         funding_account = master_funding_account
     elif isinstance(environment.runner, runners.WorkerRunner):
         worker_id = environment.runner.worker_index
@@ -454,3 +463,9 @@ def _(parser):
         required=False,
         default=16,
         help="How many funding accounts to generate for workers")
+    parser.add_argument(
+        "--run-id",
+        default="",
+        help="Unique index to append to static account ids. "
+        "Change between runs if you need a new state. Keep at default if you want to reuse the old state"
+    )

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

@@ -92,9 +92,11 @@ def on_locust_init(environment, **kwargs):
     account = base.Account(
         key.Key.from_seed_testonly(environment.congestion_account_id,
                                    environment.congestion_account_id))
-    node.send_tx_retry(
-        base.CreateSubAccount(funding_account, account.key, balance=50000.0),
-        "create congestion funding account")
+    if not node.account_exists(account.key.account_id):
+        node.send_tx_retry(
+            base.CreateSubAccount(funding_account, account.key,
+                                  balance=50000.0),
+            "create congestion funding account")
     account.refresh_nonce(node.node)
     node.send_tx_retry(
         base.Deploy(

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

@@ -160,13 +160,14 @@ def on_locust_init(environment, **kwargs):
         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)
-        node.send_tx_retry(
-            CreateSubAccount(funding_account,
-                             ft_account.key,
-                             balance=FTContract.INIT_BALANCE),
-            "create ft funding account")
-        ft_contract = FTContract(ft_account, ft_contract_code)
-        ft_contract.install(node)
+        if not node.account_exists(ft_account.key.account_id):
+            node.send_tx_retry(
+                CreateSubAccount(funding_account,
+                                 ft_account.key,
+                                 balance=FTContract.INIT_BALANCE),
+                "create ft funding account")
+            ft_contract = FTContract(ft_account, ft_contract_code)
+            ft_contract.install(node)
         environment.ft_contracts.append(ft_contract)
 
 

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

@@ -280,7 +280,7 @@ def on_locust_init(environment, **kwargs):
     # `master_funding_account` is the same on all runners, allowing to share a
     # single instance of SocialDB in its `social` sub account
     funding_account = environment.master_funding_account
-    environment.social_account_id = f"social.{funding_account.key.account_id}"
+    environment.social_account_id = f"social{environment.parsed_options.run_id}.{funding_account.key.account_id}"
 
     # Create SocialDB account, unless we are a worker, in which case the master already did it
     if not isinstance(environment.runner, runners.WorkerRunner):
@@ -291,21 +291,23 @@ def on_locust_init(environment, **kwargs):
             )
 
         social_contract_code = environment.parsed_options.social_db_wasm
-        contract_key = key.Key.from_random(environment.social_account_id)
+        contract_key = key.Key.from_seed_testonly(environment.social_account_id,
+                                                  environment.social_account_id)
         social_account = Account(contract_key)
 
         node = NearNodeProxy(environment)
-        node.send_tx_retry(
-            CreateSubAccount(funding_account,
-                             social_account.key,
-                             balance=50000.0),
-            "create socialDB funding account")
-        social_account.refresh_nonce(node.node)
-        node.send_tx_retry(
-            Deploy(social_account, social_contract_code, "Social DB"),
-            "deploy socialDB contract")
-        node.send_tx_retry(InitSocialDB(social_account),
-                           "init socialDB contract")
+        if not node.account_exists(social_account.key.account_id):
+            node.send_tx_retry(
+                CreateSubAccount(funding_account,
+                                 social_account.key,
+                                 balance=50000.0),
+                "create socialDB funding account")
+            social_account.refresh_nonce(node.node)
+            node.send_tx_retry(
+                Deploy(social_account, social_contract_code, "Social DB"),
+                "deploy socialDB contract")
+            node.send_tx_retry(InitSocialDB(social_account),
+                               "init socialDB contract")
 
 
 # Social specific CLI args