Use root account for soroban upgrades in loadgen

src/simulation/LoadGenerator.cpp

@@ -657,7 +657,11 @@ LoadGenerator::generateLoad(GeneratedLoadConfig cfg)
                 return;
             }
 
-            uint64_t sourceAccountId = getNextAvailableAccount(ledgerNum);
+            uint64_t sourceAccountId = 0;
+            if (!cfg.useRootAccountForSorobanUpgradeFlow)
+            {
+                sourceAccountId = getNextAvailableAccount(ledgerNum);
+            }
 
             std::function<std::pair<TxGenerator::TestAccountPtr,
                                     TransactionFrameBaseConstPtr>()>
@@ -753,6 +757,14 @@ LoadGenerator::generateLoad(GeneratedLoadConfig cfg)
                     auto upgradeBytes =
                         mTxGenerator.getConfigUpgradeSetFromLoadConfig(
                             cfg.getSorobanUpgradeConfig());
+                    if (cfg.useRootAccountForSorobanUpgradeFlow)
+                    {
+                        return mTxGenerator
+                            .invokeSorobanCreateUpgradeTransaction(
+                                mRoot, upgradeBytes, *mCodeKey,
+                                *mContractInstanceKeys.begin(),
+                                cfg.maxGeneratedFeeRate);
+                    }
                     return mTxGenerator.invokeSorobanCreateUpgradeTransaction(
                         ledgerNum, sourceAccountId, upgradeBytes, *mCodeKey,
                         *mContractInstanceKeys.begin(),
@@ -771,6 +783,19 @@ LoadGenerator::generateLoad(GeneratedLoadConfig cfg)
             {
                 --cfg.nTxs;
             }
+            else if (cfg.mode == LoadGenMode::SOROBAN_UPGRADE_SETUP &&
+                     cfg.useRootAccountForSorobanUpgradeFlow)
+            {
+                // If submission failed during SOROBAN_UPGRADE_SETUP, the
+                // contract instance key must be regenerated
+                // so we reset the nInstances to 1 to prevent
+                // subtracting past zero.
+                auto& sorobanCfg = cfg.getMutSorobanConfig();
+                if (sorobanCfg.nInstances == 0)
+                {
+                    sorobanCfg.nInstances = 1;
+                }
+            }
             else if (mFailed)
             {
                 break;
@@ -852,11 +877,9 @@ LoadGenerator::submitTx(GeneratedLoadConfig const& cfg,
     TransactionResultCode code;
     TransactionQueue::AddResultCode status;
     uint32_t numTries = 0;
-
     while ((status = execute(tx, cfg.mode, code)) !=
            TransactionQueue::AddResultCode::ADD_STATUS_PENDING)
     {
-
         if (cfg.skipLowFeeTxs &&
             (status ==
                  TransactionQueue::AddResultCode::ADD_STATUS_TRY_AGAIN_LATER ||
@@ -870,6 +893,24 @@ LoadGenerator::submitTx(GeneratedLoadConfig const& cfg,
                       tx->getInclusionFee());
             return false;
         }
+        // If we are using the root account to perform a soroban upgrade flow,
+        // retry the transaction.
+        if (cfg.mode == LoadGenMode::SOROBAN_UPGRADE_SETUP &&
+            cfg.useRootAccountForSorobanUpgradeFlow)
+        {
+            if (status == TransactionQueue::AddResultCode::
+                              ADD_STATUS_TRY_AGAIN_LATER ||
+                (status == TransactionQueue::AddResultCode::ADD_STATUS_ERROR &&
+                 code == txBAD_SEQ))
+            {
+                // The next attempt will regenerate a contract instance key.
+                mContractInstanceKeys.clear();
+                maybeHandleFailedTx(tx, from, status, code); // Update seq num
+                return false;
+            }
+            mFailed = true;
+            return false;
+        }
         if (++numTries >= TX_SUBMIT_MAX_TRIES ||
             status != TransactionQueue::AddResultCode::ADD_STATUS_ERROR)
         {
@@ -1083,9 +1124,17 @@ LoadGenerator::createUploadWasmTransaction(GeneratedLoadConfig const& cfg,
     // instance and ContractCode LE overhead
     mContactOverheadBytes = wasmBytes.size() + 160;
 
-    return mTxGenerator.createUploadWasmTransaction(ledgerNum, sourceAccountId,
-                                                    wasmBytes, *mCodeKey,
-                                                    cfg.maxGeneratedFeeRate);
+    if (cfg.useRootAccountForSorobanUpgradeFlow)
+    {
+        return mTxGenerator.createUploadWasmTransaction(
+            mRoot, wasmBytes, *mCodeKey, cfg.maxGeneratedFeeRate);
+    }
+    else
+    {
+        return mTxGenerator.createUploadWasmTransaction(
+            ledgerNum, sourceAccountId, wasmBytes, *mCodeKey,
+            cfg.maxGeneratedFeeRate);
+    }
 }
 
 std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
@@ -1096,9 +1145,19 @@ LoadGenerator::createInstanceTransaction(GeneratedLoadConfig const& cfg,
     auto salt = sha256("upgrade" +
                        std::to_string(++mNumCreateContractTransactionCalls));
 
-    auto txPair = mTxGenerator.createContractTransaction(
-        ledgerNum, sourceAccountId, *mCodeKey, mContactOverheadBytes, salt,
-        cfg.maxGeneratedFeeRate);
+    std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr> txPair;
+    if (cfg.useRootAccountForSorobanUpgradeFlow)
+    {
+        txPair = mTxGenerator.createContractTransaction(
+            mRoot, *mCodeKey, mContactOverheadBytes, salt,
+            cfg.maxGeneratedFeeRate);
+    }
+    else
+    {
+        txPair = mTxGenerator.createContractTransaction(
+            ledgerNum, sourceAccountId, *mCodeKey, mContactOverheadBytes, salt,
+            cfg.maxGeneratedFeeRate);
+    }
 
     auto const& instanceLk =
         txPair.second->sorobanResources().footprint.readWrite.back();

src/simulation/LoadGenerator.h

@@ -117,6 +117,7 @@ struct GeneratedLoadConfig
     Json::Value getStatus() const;
 
     LoadGenMode mode = LoadGenMode::CREATE;
+    bool useRootAccountForSorobanUpgradeFlow = false;
     uint32_t nAccounts = 0;
     uint32_t offset = 0;
     uint32_t nTxs = 0;
@@ -329,6 +330,10 @@ class LoadGenerator
     std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
     createUploadWasmTransaction(GeneratedLoadConfig const& cfg,
                                 uint32_t ledgerNum, uint64_t sourceAccountId);
+    std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
+    createUploadWasmTransaction(GeneratedLoadConfig const& cfg,
+                                uint32_t ledgerNum,
+                                TxGenerator::TestAccountPtr sourceAccount);
 
     std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
     createInstanceTransaction(GeneratedLoadConfig const& cfg,

src/simulation/TxGenerator.cpp

@@ -271,8 +271,18 @@ TxGenerator::createUploadWasmTransaction(
     std::optional<uint32_t> maxGeneratedFeeRate,
     std::optional<SorobanResources> uploadResources)
 {
-    auto account = findAccount(accountId, ledgerNum);
+    return createUploadWasmTransaction(findAccount(accountId, ledgerNum), wasm,
+                                       contractCodeLedgerKey,
+                                       maxGeneratedFeeRate, uploadResources);
+}
 
+std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
+TxGenerator::createUploadWasmTransaction(
+    TxGenerator::TestAccountPtr account, xdr::opaque_vec<> const& wasm,
+    LedgerKey const& contractCodeLedgerKey,
+    std::optional<uint32_t> maxGeneratedFeeRate,
+    std::optional<SorobanResources> uploadResources)
+{
     if (!uploadResources)
     {
         uploadResources = SorobanResources{};
@@ -307,7 +317,16 @@ TxGenerator::createContractTransaction(
     uint64_t contractOverheadBytes, uint256 const& salt,
     std::optional<uint32_t> maxGeneratedFeeRate)
 {
-    auto account = findAccount(accountId, ledgerNum);
+    return createContractTransaction(findAccount(accountId, ledgerNum), codeKey,
+                                     contractOverheadBytes, salt,
+                                     maxGeneratedFeeRate);
+}
+std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
+TxGenerator::createContractTransaction(
+    TxGenerator::TestAccountPtr account, LedgerKey const& codeKey,
+    uint64_t contractOverheadBytes, uint256 const& salt,
+    std::optional<uint32_t> maxGeneratedFeeRate)
+{
     SorobanResources createResources{};
     createResources.instructions = 1'000'000;
     createResources.readBytes = contractOverheadBytes;
@@ -774,7 +793,18 @@ TxGenerator::invokeSorobanCreateUpgradeTransaction(
     std::optional<uint32_t> maxGeneratedFeeRate,
     std::optional<SorobanResources> resources)
 {
-    auto account = findAccount(accountId, ledgerNum);
+    return invokeSorobanCreateUpgradeTransaction(
+        findAccount(accountId, ledgerNum), upgradeBytes, codeKey, instanceKey,
+        maxGeneratedFeeRate, resources);
+}
+
+std::pair<TxGenerator::TestAccountPtr, TransactionFrameBaseConstPtr>
+TxGenerator::invokeSorobanCreateUpgradeTransaction(
+    TxGenerator::TestAccountPtr account, SCBytes const& upgradeBytes,
+    LedgerKey const& codeKey, LedgerKey const& instanceKey,
+    std::optional<uint32_t> maxGeneratedFeeRate,
+    std::optional<SorobanResources> resources)
+{
     auto const& contractID = instanceKey.contractData().contract;
 
     LedgerKey upgradeLK(CONTRACT_DATA);

src/simulation/TxGenerator.h

@@ -99,13 +99,25 @@ class TxGenerator
         std::optional<uint32_t> maxGeneratedFeeRate,
         std::optional<SorobanResources> resources = std::nullopt);
     std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
+    createUploadWasmTransaction(
+        TestAccountPtr account, xdr::opaque_vec<> const& wasm,
+        LedgerKey const& contractCodeLedgerKey,
+        std::optional<uint32_t> maxGeneratedFeeRate,
+        std::optional<SorobanResources> resources = std::nullopt);
+
+    std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
     createContractTransaction(uint32_t ledgerNum, uint64_t accountId,
                               LedgerKey const& codeKey,
                               uint64_t contractOverheadBytes,
                               uint256 const& salt,
                               std::optional<uint32_t> maxGeneratedFeeRate);
 
     std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
+    createContractTransaction(TestAccountPtr account, LedgerKey const& codeKey,
+                              uint64_t contractOverheadBytes,
+                              uint256 const& salt,
+                              std::optional<uint32_t> maxGeneratedFeeRate);
+    std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
     invokeSorobanLoadTransaction(uint32_t ledgerNum, uint64_t accountId,
                                  TxGenerator::ContractInstance const& instance,
                                  uint64_t contractOverheadBytes,
@@ -117,6 +129,12 @@ class TxGenerator
         std::optional<uint32_t> maxGeneratedFeeRate,
         std::optional<SorobanResources> resources = std::nullopt);
     std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
+    invokeSorobanCreateUpgradeTransaction(
+        TestAccountPtr account, SCBytes const& upgradeBytes,
+        LedgerKey const& codeKey, LedgerKey const& instanceKey,
+        std::optional<uint32_t> maxGeneratedFeeRate,
+        std::optional<SorobanResources> resources = std::nullopt);
+    std::pair<TestAccountPtr, TransactionFrameBaseConstPtr>
     sorobanRandomWasmTransaction(uint32_t ledgerNum, uint64_t accountId,
                                  uint32_t inclusionFee);
 

src/test/TestUtils.cpp

@@ -185,40 +185,22 @@ upgradeSorobanNetworkConfig(std::function<void(SorobanNetworkConfig&)> modifyFn,
     auto nodes = simulation->getNodes();
     auto& lg = nodes[0]->getLoadGenerator();
     auto& app = *nodes[0];
-
     auto& complete =
         app.getMetrics().NewMeter({"loadgen", "run", "complete"}, "run");
-    auto completeCount = complete.count();
-    // Only create an account if there are none aleady created.
-    uint32_t offset = 0;
-    if (app.getMetrics()
-            .NewMeter({"loadgen", "account", "created"}, "account")
-            .count() == 0)
-    {
-        auto createAccountsLoadConfig =
-            GeneratedLoadConfig::createAccountsLoad(1, 1);
-        offset = std::numeric_limits<uint32_t>::max() - 1;
-        createAccountsLoadConfig.offset = offset;
-
-        lg.generateLoad(createAccountsLoadConfig);
-        simulation->crankUntil(
-            [&]() { return complete.count() == completeCount + 1; },
-            300 * Herder::EXP_LEDGER_TIMESPAN_SECONDS, false);
-    }
 
     // Create upload wasm transaction.
     auto createUploadCfg = GeneratedLoadConfig::createSorobanUpgradeSetupLoad();
-    createUploadCfg.offset = offset;
+    createUploadCfg.useRootAccountForSorobanUpgradeFlow = true;
     lg.generateLoad(createUploadCfg);
-    completeCount = complete.count();
+    auto completeCount = complete.count();
     simulation->crankUntil(
-        [&]() { return complete.count() == completeCount + 1; },
+        [&]() { return complete.count() >= completeCount + 1; },
         300 * Herder::EXP_LEDGER_TIMESPAN_SECONDS, false);
 
     // Create upgrade transaction.
     auto createUpgradeLoadGenConfig = GeneratedLoadConfig::txLoad(
         LoadGenMode::SOROBAN_CREATE_UPGRADE, 1, 1, 1);
-    createUpgradeLoadGenConfig.offset = offset;
+    createUpgradeLoadGenConfig.useRootAccountForSorobanUpgradeFlow = true;
     // Get current network config.
     auto cfg = nodes[0]->getLedgerManager().getSorobanNetworkConfig();
     modifyFn(cfg);