From 263942e0291bbb6db244d53122026c3f883a8719 Mon Sep 17 00:00:00 2001
From: Thomas Brady <thomas.brady@stellar.org>
Date: Mon, 25 Nov 2024 21:28:13 -0800
Subject: [PATCH] Fix timing out 'generate soroban load' by special casing
 two-part upgrade sequence retries

---
 src/simulation/LoadGenerator.cpp | 38 ++++++++++++++++++++++++++------
 src/test/TestUtils.cpp           |  5 +++--
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/src/simulation/LoadGenerator.cpp b/src/simulation/LoadGenerator.cpp
index 1e74dc5700..78c5a1f70e 100644
--- a/src/simulation/LoadGenerator.cpp
+++ b/src/simulation/LoadGenerator.cpp
@@ -783,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;
@@ -864,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 ||
@@ -882,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)
         {
@@ -1594,7 +1623,6 @@ GeneratedLoadConfig::createSorobanUpgradeSetupLoad()
     cfg.nAccounts = 1;
     cfg.getMutSorobanConfig().nInstances = 1;
     cfg.txRate = 1;
-    cfg.useRootAccountForSorobanUpgradeFlow = true;
     return cfg;
 }
 
@@ -1604,10 +1632,6 @@ GeneratedLoadConfig::txLoad(LoadGenMode mode, uint32_t nAccounts, uint32_t nTxs,
                             std::optional<uint32_t> maxFee)
 {
     GeneratedLoadConfig cfg;
-    if (mode == LoadGenMode::SOROBAN_CREATE_UPGRADE)
-    {
-        cfg.useRootAccountForSorobanUpgradeFlow = true;
-    }
     cfg.mode = mode;
     cfg.nAccounts = nAccounts;
     cfg.nTxs = nTxs;
diff --git a/src/test/TestUtils.cpp b/src/test/TestUtils.cpp
index 985164f8e2..61f1452a99 100644
--- a/src/test/TestUtils.cpp
+++ b/src/test/TestUtils.cpp
@@ -185,21 +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");
 
     // Create upload wasm transaction.
     auto createUploadCfg = GeneratedLoadConfig::createSorobanUpgradeSetupLoad();
+    createUploadCfg.useRootAccountForSorobanUpgradeFlow = true;
     lg.generateLoad(createUploadCfg);
     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.useRootAccountForSorobanUpgradeFlow = true;
     // Get current network config.
     auto cfg = nodes[0]->getLedgerManager().getSorobanNetworkConfig();
     modifyFn(cfg);