From f5ec14ab64f0b7ad8164030f017a249a371ab31f Mon Sep 17 00:00:00 2001
From: Bushstar <bushsolo@gmail.com>
Date: Tue, 21 Nov 2023 08:12:46 +0000
Subject: [PATCH] Generate snapshot on ConnectTip

---
 src/dbwrapper.h         |  4 ++--
 src/dfi/masternodes.cpp |  2 --
 src/flushablestorage.h  | 50 ++++++++++++++++++++++++++++++++++++++---
 src/validation.cpp      |  8 +++++++
 4 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/src/dbwrapper.h b/src/dbwrapper.h
index 060e11e671..35d3dc5da8 100644
--- a/src/dbwrapper.h
+++ b/src/dbwrapper.h
@@ -362,8 +362,8 @@ class CDBWrapper
         return new CDBIterator(*this, pdb->NewIterator(readOptions));
     }
 
-    [[nodiscard]] std::unique_ptr<CStorageSnapshot> GetSnapshot() const {
-        return std::make_unique<CStorageSnapshot>(pdb);
+    [[nodiscard]] std::shared_ptr<CStorageSnapshot> GetSnapshot() const {
+        return std::make_shared<CStorageSnapshot>(pdb);
     }
 
     /**
diff --git a/src/dfi/masternodes.cpp b/src/dfi/masternodes.cpp
index c8b8d49b89..35f7885a11 100644
--- a/src/dfi/masternodes.cpp
+++ b/src/dfi/masternodes.cpp
@@ -1401,8 +1401,6 @@ void CalcMissingRewardTempFix(CCustomCSView &mnview, const uint32_t targetHeight
 }
 
 std::unique_ptr<CCustomCSView> GetViewSnapshot() {
-    LOCK(cs_main);
-
     // Get database snapshot
     auto snapshotDB = pcustomcsDB->GetSnapshotDB();
 
diff --git a/src/flushablestorage.h b/src/flushablestorage.h
index b5f89a90c8..9449981465 100644
--- a/src/flushablestorage.h
+++ b/src/flushablestorage.h
@@ -14,6 +14,8 @@
 
 #include <optional>
 
+extern CCriticalSection cs_main;
+
 using TBytes = std::vector<unsigned char>;
 using MapKV = std::map<TBytes, std::optional<TBytes>>;
 
@@ -141,8 +143,8 @@ class CStorageLevelDB : public CStorageKV {
         : db{std::make_shared<CDBWrapper>(dbName, cacheSize, fMemory, fWipe)}, batch(*db) {}
 
     // Snapshot constructor
-    explicit CStorageLevelDB(std::shared_ptr<CDBWrapper> &db, std::unique_ptr<CStorageSnapshot> &otherSnapshot)
-            : db(db), batch(*db), snapshot(std::move(otherSnapshot)) {
+    CStorageLevelDB(std::shared_ptr<CDBWrapper> &db, std::shared_ptr<CStorageSnapshot> &otherSnapshot)
+            : db(db), batch(*db), snapshot(otherSnapshot) {
         options.snapshot = snapshot->GetLevelDBSnapshot();
     }
 
@@ -200,15 +202,57 @@ class CStorageLevelDB : public CStorageKV {
         return db->IsEmpty();
     }
     std::unique_ptr<CStorageLevelDB> GetSnapshotDB() {
+        if (snapshotUpdated.load()) {
+            // Lock cs_main when updating from tipSnapshot to avoid
+            // race with Dis/ConnectTip.
+            LOCK(cs_main);
+            // Double check bool for safety now we are under lock.
+            if (snapshotUpdated.load()) {
+                snapshotUpdated.store(false);
+                if (tipSnapshot) {
+                    tipCopySnapshot = tipSnapshot;
+                } else {
+                    tipCopySnapshot.reset();
+                }
+            }
+        }
+
+        if (tipCopySnapshot) {
+            return std::make_unique<CStorageLevelDB>(db, tipCopySnapshot);
+        }
+
         auto dbSnapshot = db->GetSnapshot();
         return std::make_unique<CStorageLevelDB>(db, dbSnapshot);
     }
+    void GenerateSnapshot() {
+        snapshotUpdated.store(true);
+        tipSnapshot = db->GetSnapshot();
+    }
+    void ReleaseSnapshot() {
+        snapshotUpdated.store(true);
+        tipSnapshot.reset();
+    }
 
 private:
     std::shared_ptr<CDBWrapper> db;
     CDBBatch batch;
     leveldb::ReadOptions options;
-    std::unique_ptr<CStorageSnapshot> snapshot;
+
+    // If this snapshot is set it will be used when
+    // reading from the DB.
+    std::shared_ptr<CStorageSnapshot> snapshot;
+
+    // Snapshot updated at tip. Accessed under cs_main lock.
+    std::shared_ptr<CStorageSnapshot> tipSnapshot;
+
+    // Used to create another CStorageLevelDB object set
+    // as the snapshot member. Lock free.
+    std::shared_ptr<CStorageSnapshot> tipCopySnapshot;
+
+    // Used to determine whether the tipSnapshot has been
+    // updated and should be copied over to the lock free
+    // tipCopySnapshot.
+    std::atomic<bool> snapshotUpdated{};
 };
 
 // Flushable storage
diff --git a/src/validation.cpp b/src/validation.cpp
index 627c05fb68..9ba05388e5 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -3814,6 +3814,9 @@ bool CChainState::DisconnectTip(CValidationState &state,
         assert(flushed);
         mnview.GetHistoryWriters().FlushDB();
 
+        // Release snapshot
+        pcustomcsDB->ReleaseSnapshot();
+
         if (!disconnectedConfirms.empty()) {
             for (const auto &confirm : disconnectedConfirms) {
                 panchorAwaitingConfirms->Add(confirm);
@@ -3991,6 +3994,11 @@ bool CChainState::ConnectTip(CValidationState &state,
         assert(flushed);
         mnview.GetHistoryWriters().FlushDB();
 
+        // Generate snapshot
+        if (!IsInitialBlockDownload()) {
+            pcustomcsDB->GenerateSnapshot();
+        }
+
         // Delete all other confirms from memory
         if (rewardedAnchors) {
             std::vector<uint256> oldConfirms;