Lumen snapshot invariant (#5059)

# Description

Resolves #5028

<!---

Describe what this pull request does, which issue it's resolving
(usually applicable for code changes).

--->

# Checklist
- [ ] Reviewed the
[contributing](https://github.com/stellar/stellar-core/blob/master/CONTRIBUTING.md#submitting-changes)
document
- [ ] Rebased on top of master (no merge commits)
- [ ] Ran `clang-format` v8.0.0 (via `make format` or the Visual Studio
extension)
- [ ] Compiles
- [ ] Ran all tests
- [ ] If change impacts performance, include supporting evidence per the
[performance
document](https://github.com/stellar/stellar-core/blob/master/performance-eval/performance-eval.md)

Author: sisuresh

src/invariant/ArchivedStateConsistency.cpp

@@ -30,7 +30,8 @@ ArchivedStateConsistency::ArchivedStateConsistency() : Invariant(true)
 std::string
 ArchivedStateConsistency::checkSnapshot(
     CompleteConstLedgerStatePtr ledgerState,
-    InMemorySorobanState const& inMemorySnapshot)
+    InMemorySorobanState const& inMemorySnapshot,
+    std::function<bool()> isStopping)
 {
     LogSlowExecution logSlow("ArchivedStateConsistency::stateSnapshotInvariant",
                              LogSlowExecution::Mode::AUTOMATIC_RAII, "took",
@@ -54,33 +55,39 @@ ArchivedStateConsistency::checkSnapshot(
     // For each entry in the Live BucketList, check if we have seen it in a
     // previous level. If not, this entry is the newest version, so check if it
     // exists in the Hot Archive.
-    auto checkIfLiveEntryInArchive =
-        [&seenKeys, &errorMsg, &hotArchiveSnapshot](BucketEntry const& be) {
-            if (be.type() == LIVEENTRY || be.type() == INITENTRY)
-            {
-                auto lk = LedgerEntryKey(be.liveEntry());
-                auto [_, wasInserted] = seenKeys.emplace(lk);
-
-                // If this BucketEntry is not shadowed, and the key exists in
-                // the Hot Archive, we have an error.
-                if (wasInserted && hotArchiveSnapshot->load(lk))
-                {
-                    errorMsg = fmt::format(
-                        FMT_STRING("ArchivedStateConsistency invariant failed: "
-                                   "Live entry is present in both live and "
-                                   "archived state: {}"),
-                        xdrToCerealString(lk, "entry_key"));
-                    return Loop::COMPLETE;
-                }
-            }
-            // Mark DEADENTRY as seen, but the key does not exist wrt ledger
-            // state, so we don't need to check the Hot Archive.
-            else if (be.type() == DEADENTRY)
+    auto checkIfLiveEntryInArchive = [&seenKeys, &errorMsg, &hotArchiveSnapshot,
+                                      &isStopping](BucketEntry const& be) {
+        // Allow early termination if application is stopping
+        if (isStopping())
+        {
+            return Loop::COMPLETE;
+        }
+
+        if (be.type() == LIVEENTRY || be.type() == INITENTRY)
+        {
+            auto lk = LedgerEntryKey(be.liveEntry());
+            auto [_, wasInserted] = seenKeys.emplace(lk);
+
+            // If this BucketEntry is not shadowed, and the key exists in
+            // the Hot Archive, we have an error.
+            if (wasInserted && hotArchiveSnapshot->load(lk))
             {
-                seenKeys.emplace(be.deadEntry());
+                errorMsg = fmt::format(
+                    FMT_STRING("ArchivedStateConsistency invariant failed: "
+                               "Live entry is present in both live and "
+                               "archived state: {}"),
+                    xdrToCerealString(lk, "entry_key"));
+                return Loop::COMPLETE;
             }
-            return Loop::INCOMPLETE;
-        };
+        }
+        // Mark DEADENTRY as seen, but the key does not exist wrt ledger
+        // state, so we don't need to check the Hot Archive.
+        else if (be.type() == DEADENTRY)
+        {
+            seenKeys.emplace(be.deadEntry());
+        }
+        return Loop::INCOMPLETE;
+    };
 
     // We just need to check for Soroban types that are stored in the Hot
     // Archive

src/invariant/ArchivedStateConsistency.h

@@ -44,6 +44,7 @@ class ArchivedStateConsistency : public Invariant
 
     virtual std::string
     checkSnapshot(CompleteConstLedgerStatePtr ledgerState,
-                  InMemorySorobanState const& inMemorySnapshot) override;
+                  InMemorySorobanState const& inMemorySnapshot,
+                  std::function<bool()> isStopping) override;
 };
 }

src/invariant/ConservationOfLumens.cpp

@@ -3,13 +3,12 @@
 // of this distribution or at http://www.apache.org/licenses/LICENSE-2.0
 
 #include "invariant/ConservationOfLumens.h"
-#include "crypto/SHA.h"
 #include "invariant/InvariantManager.h"
+#include "ledger/LedgerManager.h"
 #include "ledger/LedgerTxn.h"
 #include "main/Application.h"
 #include "transactions/TransactionUtils.h"
-#include "util/GlobalChecks.h"
-#include <fmt/format.h>
+#include "util/LogSlowExecution.h"
 #include <numeric>
 
 namespace stellar
@@ -175,4 +174,222 @@ ConservationOfLumens::checkOnOperationApply(
     }
     return {};
 }
+
+// Helper function that processes an entry if it hasn't been seen before.
+// Returns true on success, false on error (with error set in errorMsg).
+static bool
+processEntryIfNew(LedgerEntry const& entry, LedgerKey const& key,
+                  std::unordered_set<LedgerKey>& countedKeys,
+                  Asset const& asset,
+                  AssetContractInfo const& assetContractInfo,
+                  int64_t& sumBalance, std::string& errorMsg)
+{
+    if (countedKeys.count(key) != 0)
+    {
+        return true;
+    }
+
+    auto result = getAssetBalance(entry, asset, assetContractInfo);
+
+    if (result.overflowed)
+    {
+        errorMsg = fmt::format(
+            FMT_STRING(
+                "ConservationOfLumens: getAssetBalance overflow for key: {}"),
+            xdrToCerealString(key, "ledger_key"));
+        return false;
+    }
+
+    if (!result.assetMatched)
+    {
+        return true; // Asset doesn't match, skip
+    }
+
+    if (!result.balance || !addBalance(sumBalance, *result.balance))
+    {
+        errorMsg = fmt::format(
+            FMT_STRING(
+                "ConservationOfLumens: Overflow adding balance for key: {}"),
+            xdrToCerealString(key, "ledger_key"));
+        return false;
+    }
+
+    countedKeys.emplace(key);
+
+    return true;
+}
+
+// Scan live bucket list for entries that can hold the native asset
+static void
+scanLiveBuckets(
+    std::shared_ptr<SearchableLiveBucketListSnapshot const> const& liveSnapshot,
+    Asset const& asset, AssetContractInfo const& assetContractInfo,
+    int64_t& sumBalance, std::string& errorMsg,
+    std::function<bool()> const& isStopping)
+{
+    // Scan all entry types that can hold the native asset
+    for (auto let : xdr::xdr_traits<LedgerEntryType>::enum_values())
+    {
+        LedgerEntryType type = static_cast<LedgerEntryType>(let);
+        if (!canHoldAsset(type, asset))
+        {
+            continue;
+        }
+
+        std::unordered_set<LedgerKey> countedKeys;
+
+        liveSnapshot->scanForEntriesOfType(
+            type, [&](BucketEntry const& be) -> Loop {
+                if (isStopping())
+                {
+                    return Loop::COMPLETE;
+                }
+
+                if (be.type() == LIVEENTRY || be.type() == INITENTRY)
+                {
+                    if (!processEntryIfNew(
+                            be.liveEntry(), LedgerEntryKey(be.liveEntry()),
+                            countedKeys, asset, assetContractInfo, sumBalance,
+                            errorMsg))
+                    {
+                        return Loop::COMPLETE;
+                    }
+                }
+                else if (be.type() == DEADENTRY)
+                {
+                    countedKeys.emplace(be.deadEntry());
+                }
+                return Loop::INCOMPLETE;
+            });
+
+        if (!errorMsg.empty())
+        {
+            return;
+        }
+    }
+}
+
+static Loop
+scanHotArchiveBucket(HotArchiveBucketSnapshot const& bucket,
+                     std::unordered_set<LedgerKey>& countedKeys,
+                     Asset const& asset,
+                     AssetContractInfo const& assetContractInfo,
+                     int64_t& sumBalance, std::string& errorMsg,
+                     std::function<bool()> const& isStopping)
+{
+    for (HotArchiveBucketInputIterator iter(bucket.getRawBucket()); iter;
+         ++iter)
+    {
+        // Allow early termination if application is stopping
+        if (isStopping())
+        {
+            return Loop::COMPLETE;
+        }
+
+        auto const& be = *iter;
+        if (be.type() == HOT_ARCHIVE_ARCHIVED)
+        {
+            if (!canHoldAsset(be.archivedEntry().data.type(), asset))
+            {
+                continue;
+            }
+            if (!processEntryIfNew(be.archivedEntry(),
+                                   LedgerEntryKey(be.archivedEntry()),
+                                   countedKeys, asset, assetContractInfo,
+                                   sumBalance, errorMsg))
+            {
+                return Loop::COMPLETE;
+            }
+        }
+        else if (be.type() == HOT_ARCHIVE_LIVE &&
+                 canHoldAsset(be.key().type(), asset))
+        {
+            // HOT_ARCHIVE_LIVE means entry was restored from archive,
+            // so mark it as seen (shadowing any archived versions)
+            countedKeys.emplace(be.key());
+        }
+    }
+    return Loop::INCOMPLETE;
+}
+
+std::string
+ConservationOfLumens::checkSnapshot(
+    CompleteConstLedgerStatePtr ledgerState,
+    InMemorySorobanState const& inMemorySnapshot,
+    std::function<bool()> isStopping)
+{
+    LogSlowExecution logSlow("ConservationOfLumens::checkSnapshot",
+                             LogSlowExecution::Mode::AUTOMATIC_RAII, "took",
+                             std::chrono::seconds(90));
+
+    auto liveSnapshot = ledgerState->getBucketSnapshot();
+    auto hotArchiveSnapshot = ledgerState->getHotArchiveSnapshot();
+    auto const& header = liveSnapshot->getLedgerHeader();
+
+    // This invariant can fail prior to v24 due to bugs
+    if (protocolVersionIsBefore(header.ledgerVersion, ProtocolVersion::V_24))
+    {
+        return std::string{};
+    }
+
+    Asset nativeAsset(ASSET_TYPE_NATIVE);
+
+    int64_t sumBalance = 0;
+    std::string errorMsg;
+
+    // Start with the fee pool from the ledger header
+    if (!addBalance(sumBalance, header.feePool))
+    {
+        return fmt::format(
+            FMT_STRING("ConservationOfLumens invariant failed: "
+                       "Fee pool balance overflowed when added to total. "
+                       "Current sum: {}, Fee pool: {}"),
+            sumBalance, header.feePool);
+    }
+
+    // Scan the Live BucketList for native balances using loopAllBuckets
+
+    scanLiveBuckets(liveSnapshot, nativeAsset, mLumenContractInfo, sumBalance,
+                    errorMsg, isStopping);
+
+    if (!errorMsg.empty())
+    {
+        return errorMsg;
+    }
+
+    // Scan the Hot Archive for native balances using loopAllBuckets
+    {
+        std::unordered_set<LedgerKey> countedKeys;
+        hotArchiveSnapshot->loopAllBuckets(
+            [&countedKeys, &nativeAsset, &sumBalance, &errorMsg, &isStopping,
+             this](HotArchiveBucketSnapshot const& bucket) {
+                return scanHotArchiveBucket(bucket, countedKeys, nativeAsset,
+                                            mLumenContractInfo, sumBalance,
+                                            errorMsg, isStopping);
+            });
+
+        if (!errorMsg.empty())
+        {
+            return errorMsg;
+        }
+    }
+
+    // We stopped early, so it's likely we didn't finish scanning everything
+    if (isStopping())
+    {
+        return std::string{};
+    }
+
+    // Compare the calculated total with totalCoins from the ledger header
+    if (sumBalance != header.totalCoins)
+    {
+        return fmt::format(
+            FMT_STRING("ConservationOfLumens invariant failed: "
+                       "Total native asset supply mismatch. "
+                       "Calculated from buckets: {}, Expected (totalCoins): "
+                       "{}, Difference: {}"),
+            sumBalance, header.totalCoins, header.totalCoins - sumBalance);
+    }
+    return std::string{};
+}
 }

src/invariant/ConservationOfLumens.h

@@ -33,6 +33,11 @@ class ConservationOfLumens : public Invariant
         LedgerTxnDelta const& ltxDelta,
         std::vector<ContractEvent> const& events, AppConnector& app) override;
 
+    virtual std::string
+    checkSnapshot(CompleteConstLedgerStatePtr ledgerState,
+                  InMemorySorobanState const& inMemorySnapshot,
+                  std::function<bool()> isStopping) override;
+
   private:
     AssetContractInfo const mLumenContractInfo;
 };

src/invariant/Invariant.h

@@ -86,7 +86,8 @@ class Invariant
 
     virtual std::string
     checkSnapshot(CompleteConstLedgerStatePtr ledgerState,
-                  InMemorySorobanState const& inMemorySnapshot)
+                  InMemorySorobanState const& inMemorySnapshot,
+                  std::function<bool()> isStopping)
     {
         return std::string{};
     }

src/invariant/InvariantManager.h

@@ -68,7 +68,8 @@ class InvariantManager
     // INVARIANT_EXTRA_CHECKS is enabled.
     virtual void
     runStateSnapshotInvariant(CompleteConstLedgerStatePtr ledgerState,
-                              InMemorySorobanState const& inMemorySnapshot) = 0;
+                              InMemorySorobanState const& inMemorySnapshot,
+                              std::function<bool()> isStopping) = 0;
 
     virtual void registerInvariant(std::shared_ptr<Invariant> invariant) = 0;
 

src/invariant/InvariantManagerImpl.cpp

@@ -331,7 +331,8 @@ InvariantManagerImpl::handleInvariantFailure(bool isStrict,
 void
 InvariantManagerImpl::runStateSnapshotInvariant(
     CompleteConstLedgerStatePtr ledgerState,
-    InMemorySorobanState const& inMemorySnapshot)
+    InMemorySorobanState const& inMemorySnapshot,
+    std::function<bool()> isStopping)
 {
     // Reset our trigger flag and mark the invariant as running.
     mStateSnapshotInvariantRunning = true;
@@ -342,7 +343,8 @@ InvariantManagerImpl::runStateSnapshotInvariant(
 
     for (auto const& invariant : mEnabled)
     {
-        auto result = invariant->checkSnapshot(ledgerState, inMemorySnapshot);
+        auto result =
+            invariant->checkSnapshot(ledgerState, inMemorySnapshot, isStopping);
         if (!result.empty())
         {
             auto ledgerSeq =

src/invariant/InvariantManagerImpl.h

@@ -80,9 +80,9 @@ class InvariantManagerImpl : public InvariantManager
 
     bool shouldRunInvariantSnapshot() const override;
 
-    void runStateSnapshotInvariant(
-        CompleteConstLedgerStatePtr ledgerState,
-        InMemorySorobanState const& inMemorySnapshot) override;
+    void runStateSnapshotInvariant(CompleteConstLedgerStatePtr ledgerState,
+                                   InMemorySorobanState const& inMemorySnapshot,
+                                   std::function<bool()> isStopping) override;
 
 #ifdef BUILD_TESTS
     void snapshotForFuzzer() override;