Keep root children in cache

Signed-off-by: Anthony Fieroni <bvbfan@abv.bg>

src/claimtrie.cpp

@@ -266,10 +266,9 @@ expirationQueueRowType* CClaimTrieCacheBase::getExpirationQueueCacheRow(int, boo
 
 bool CClaimTrieCacheBase::haveClaim(const std::string& name, const COutPoint& outPoint) const
 {
-    auto it = nodesToAddOrUpdate.find(name);
+    if (auto it = nodesToAddOrUpdate.find(name))
-    if (it && it->haveClaim(outPoint))
+        return it->haveClaim(outPoint);
-        return true;
+    if (nodesToDelete.count(name))
-    if (it || nodesToDelete.count(name))
         return false;
     CClaimTrieData data;
     return base->find(name, data) && data.haveClaim(outPoint);
@@ -325,12 +324,14 @@ bool CClaimTrieCacheBase::haveSupportInQueue(const std::string& name, const COut
     return haveInQueue<CSupportValue>(name, outPoint, nValidAtHeight);
 }
 
-void CClaimTrie::recurseNodes(const std::string& name, const CClaimTrieDataNode& current, std::function<recurseNodesCB> function) const {
+void CClaimTrie::recurseNodes(const std::string& name, const CClaimTrieDataNode& current, std::function<recurseNodesCB> function) const
+{
     CClaimTrieData data;
     find(name, data);
 
-    data.hash = current.hash;
+    if (!(data.flags & CAME_FROM_NODE_CACHE))
-    data.flags |= current.children.empty() ? 0 : CClaimTrieDataFlags::POTENTIAL_CHILDREN;
+        data.hash = current.hash;
+
     function(name, data, current.children);
 
     for (auto& child: current.children) {
@@ -382,10 +383,9 @@ CAmount CClaimTrie::getTotalValueOfClaimsInTrie(bool fControllingOnly) const
 
 bool CClaimTrieCacheBase::getInfoForName(const std::string& name, CClaimValue& claim) const
 {
-    auto it = nodesToAddOrUpdate.find(name);
+    if (auto it = nodesToAddOrUpdate.find(name))
-    if (it && it->getBestClaim(claim))
+        return it->getBestClaim(claim);
-        return true;
+    if (nodesToDelete.count(name))
-    if (it || nodesToDelete.count(name))
         return false;
     CClaimTrieData claims;
     return base->find(name, claims) && claims.getBestClaim(claim);
@@ -492,7 +492,8 @@ bool CClaimTrie::checkConsistency(const uint256& rootHash) const
     return success;
 }
 
-std::vector<std::pair<std::string, CClaimTrieDataNode>> CClaimTrie::nodes(const std::string &key) const {
+std::vector<std::pair<std::string, CClaimTrieDataNode>> CClaimTrie::nodes(const std::string &key) const
+{
     std::vector<std::pair<std::string, CClaimTrieDataNode>> ret;
     CClaimTrieDataNode node;
 
@@ -527,30 +528,54 @@ std::vector<std::pair<std::string, CClaimTrieDataNode>> CClaimTrie::nodes(const
     return ret;
 }
 
-bool CClaimTrie::contains(const std::string &key) const {
+bool CClaimTrie::contains(const std::string& key) const
-    return db->Exists(std::make_pair(TRIE_NODE_CHILDREN, key));
+{
+    return cacheNodes.find(key) || db->Exists(std::make_pair(TRIE_NODE_CHILDREN, key));
 }
 
-bool CClaimTrie::empty() const {
+bool CClaimTrie::empty() const
-    return !contains({});
+{
+    return cacheNodes.empty() && !contains({});
 }
 
-bool CClaimTrie::find(const std::string& key, CClaimTrieDataNode &node) const {
+std::vector<std::string> extractChildren(CClaimPrefixTrie::const_iterator it)
+{
+    std::vector<std::string> children;
+    auto size = it.key().size();
+    for (auto& child : it.children()) // ordering here is important
+        children.emplace_back(child.key().substr(size));
+    return children;
+}
+
+bool CClaimTrie::find(const std::string& key, CClaimTrieDataNode& node) const
+{
+    if (auto it = cacheNodes.find(key)) {
+        if (it.hasChildren()) {
+            node.hash = it->hash;
+            node.children = extractChildren(it);
+            return true;
+        }
+    }
     return db->Read(std::make_pair(TRIE_NODE_CHILDREN, key), node);
 }
 
-bool CClaimTrie::find(const std::string& key, CClaimTrieData &data) const {
+bool CClaimTrie::find(const std::string& key, CClaimTrieData& data) const
+{
+    if (auto it = cacheNodes.find(key)) {
+        data = it.data();
+        data.flags |= CAME_FROM_NODE_CACHE;
+        return true;
+    }
     return db->Read(std::make_pair(TRIE_NODE_CLAIMS, key), data);
 }
 
 template <typename K, typename T>
-void BatchWrite(CDBBatch& batch, uint8_t dbkey, const K& key, const std::vector<T>& value)
+void BatchWrite(CDBBatch& batch, uint8_t dbkey, const K& key, const T& value)
 {
-    if (value.empty()) {
+    if (value.empty())
         batch.Erase(std::make_pair(dbkey, key));
-    } else {
+    else
         batch.Write(std::make_pair(dbkey, key), value);
-    }
 }
 
 template <typename Container>
@@ -579,22 +604,34 @@ bool CClaimTrieCacheBase::flush()
 
     getMerkleHash();
 
-    for (auto it = nodesToAddOrUpdate.begin(); it != nodesToAddOrUpdate.end(); ++it) {
+    if (!nodesToDelete.empty() || !nodesToAddOrUpdate.empty()) {
-        bool removed = forDeleteFromBase.erase(it.key());
+        base->cacheNodes.clear();
-        if (it->flags & CClaimTrieDataFlags::HASH_DIRTY) {
+        if (auto it = nodesToAddOrUpdate.begin()) {
-            CClaimTrieDataNode node;
+            if (it->flags & HASH_DIRTY) {
-            node.hash = it->hash;
+                for (auto& child : it.children())
-            for (auto &child: it.children()) // ordering here is important
+                    base->cacheNodes.copy(child);
-                node.children.push_back(child.key().substr(it.key().size()));
+                base->cacheNodes.copy(it);
-
+            }
-            batch.Write(std::make_pair(TRIE_NODE_CHILDREN, it.key()), node);
-
-            if (removed || (it->flags & CClaimTrieDataFlags::CLAIMS_DIRTY))
-                batch.Write(std::make_pair(TRIE_NODE_CLAIMS, it.key()), it.data());
         }
     }
 
+    for (auto it = nodesToAddOrUpdate.begin(); it != nodesToAddOrUpdate.end(); ++it) {
+        bool removed = forDeleteFromBase.erase(it.key());
+        if (it->flags & HASH_DIRTY) {
+            CClaimTrieDataNode node;
+            node.hash = it->hash;
+            node.children = extractChildren(it);
+
+            batch.Write(std::make_pair(TRIE_NODE_CHILDREN, it.key()), node);
+
+            if (removed || (it->flags & CLAIMS_DIRTY))
+                BatchWrite(batch, TRIE_NODE_CLAIMS, it.key(), it.data());
+        }
+        it->flags = 0;
+    }
+
     for (auto& name: forDeleteFromBase) {
+        base->cacheNodes.erase(name);
         batch.Erase(std::make_pair(TRIE_NODE_CHILDREN, name));
         batch.Erase(std::make_pair(TRIE_NODE_CLAIMS, name));
     }
@@ -614,10 +651,9 @@ bool CClaimTrieCacheBase::flush()
         LogPrintf("TrieCache size: %zu nodes on block %d, batch writes %zu bytes.\n",
                 nodesToAddOrUpdate.height(), nNextHeight, batch.SizeEstimate());
     }
-    auto ret = base->db->WriteBatch(batch);
 
     clear();
-    return ret;
+    return base->db->WriteBatch(batch);
 }
 
 bool CClaimTrieCacheBase::validateTrieConsistency(const CBlockIndex* tip)
@@ -718,22 +754,21 @@ CClaimPrefixTrie::iterator CClaimTrieCacheBase::cacheData(const std::string& nam
             // do not insert nodes that are already present
             CClaimTrieData data;
             base->find(node.first, data);
-            data.hash = node.second.hash;
+            if (!(data.flags & CAME_FROM_NODE_CACHE))
-            data.flags = node.second.children.empty() ? 0 : CClaimTrieDataFlags::POTENTIAL_CHILDREN;
+                data.hash = node.second.hash;
-            nodesToAddOrUpdate.insert(node.first, data);
+            nodesToAddOrUpdate.insert(node.first, std::move(data));
         }
         for (auto& child : node.second.children) {
             auto childKey = node.first + child;
             if (nodesAlreadyCached.insert(childKey).second) {
                 CClaimTrieData childData;
-                if (!base->find(childKey, childData))
+                base->find(childKey, childData);
-                    childData = {};
+                if (!(childData.flags & CAME_FROM_NODE_CACHE)) {
-                CClaimTrieDataNode childNode;
+                    CClaimTrieDataNode childNode;
-                if (base->find(childKey, childNode)) {
+                    if (base->find(childKey, childNode))
-                    childData.hash = childNode.hash;
+                        childData.hash = childNode.hash;
-                    childData.flags = childNode.children.empty() ? 0 : CClaimTrieDataFlags::POTENTIAL_CHILDREN;
                 }
-                nodesToAddOrUpdate.insert(childKey, childData);
+                nodesToAddOrUpdate.insert(childKey, std::move(childData));
             }
         }
     }
@@ -741,8 +776,6 @@ CClaimPrefixTrie::iterator CClaimTrieCacheBase::cacheData(const std::string& nam
     auto it = nodesToAddOrUpdate.find(name);
     if (!it && create) {
         it = nodesToAddOrUpdate.insert(name, CClaimTrieData{});
-        // if (it.hasChildren()) any children should be in the trie (not base alone)
-        //    it->flags |= CClaimTrieDataFlags::POTENTIAL_CHILDREN;
         confirmTakeoverWorkaroundNeeded(name);
     }
 
@@ -778,10 +811,11 @@ bool CClaimTrieCacheBase::getLastTakeoverForName(const std::string& name, uint16
 void CClaimTrieCacheBase::markAsDirty(const std::string& name, bool fCheckTakeover)
 {
     for (auto& node : nodesToAddOrUpdate.nodes(name)) {
-        node->flags |= CClaimTrieDataFlags::HASH_DIRTY;
+        node->flags |= HASH_DIRTY;
+
         node->hash.SetNull();
         if (node.key() == name)
-            node->flags |= CClaimTrieDataFlags::CLAIMS_DIRTY;
+            node->flags |= CLAIMS_DIRTY;
     }
 
     if (fCheckTakeover)
@@ -1490,7 +1524,7 @@ void CClaimTrieCacheBase::recurseNodes(const std::string &name,
     else {
         for (auto it = begin(); it != end(); ++it) {
             function(it.key(), it.data());
-            if ((it->flags & CClaimTrieDataFlags::POTENTIAL_CHILDREN) && !it.hasChildren()) {
+            if (!it.hasChildren()) {
                 CClaimTrieDataNode node;
                 if (base->find(it.key(), node))
                     for (auto& partialKey: node.children) {

src/claimtrie.h

@@ -139,7 +139,7 @@ typedef std::vector<CSupportValue> supportEntryType;
 enum CClaimTrieDataFlags: uint32_t {
     HASH_DIRTY = 1U,
     CLAIMS_DIRTY = 2U,
-    POTENTIAL_CHILDREN = 4U, // existing on disk
+    CAME_FROM_NODE_CACHE = 4U,
 };
 
 struct CClaimTrieData
@@ -188,7 +188,8 @@ struct CClaimTrieData
     }
 };
 
-struct CClaimTrieDataNode {
+struct CClaimTrieDataNode
+{
     uint256 hash;
     // we're using a vector to avoid RAM thrashing and for faster serialization ops.
     // We're assuming its data is inserted in order and never modified.
@@ -364,6 +365,8 @@ struct CClaimSupportToName
     const std::vector<CSupportValue> unmatchedSupports;
 };
 
+typedef CPrefixTrie<std::string, CClaimTrieData> CClaimPrefixTrie;
+
 class CClaimTrie
 {
 public:
@@ -402,6 +405,7 @@ protected:
     int nProportionalDelayFactor = 0;
     std::unique_ptr<CDBWrapper> db;
 
+    CClaimPrefixTrie cacheNodes;
     using recurseNodesCB = void(const std::string&, const CClaimTrieData&, const std::vector<std::string>&);
     void recurseNodes(const std::string& name, const CClaimTrieDataNode& current, std::function<recurseNodesCB> function) const;
 };
@@ -508,8 +512,6 @@ typedef std::map<int, expirationQueueRowType> expirationQueueType;
 typedef std::set<CClaimValue> claimIndexClaimListType;
 typedef std::vector<CClaimIndexElement> claimIndexElementListType;
 
-typedef CPrefixTrie<std::string, CClaimTrieData> CClaimPrefixTrie;
-
 class CClaimTrieCacheBase
 {
 public:

src/prefixtrie.cpp

@@ -108,16 +108,18 @@ template <typename TKey, typename TData>
 template <bool IsConst>
 typename CPrefixTrie<TKey, TData>::template Iterator<IsConst>::reference CPrefixTrie<TKey, TData>::Iterator<IsConst>::operator*() const
 {
-    assert(!node.expired());
+    auto shared = node.lock();
-    return TPair(name, node.lock()->data);
+    assert(shared);
+    return TPair(name, *(shared->data));
 }
 
 template <typename TKey, typename TData>
 template <bool IsConst>
 typename CPrefixTrie<TKey, TData>::template Iterator<IsConst>::pointer CPrefixTrie<TKey, TData>::Iterator<IsConst>::operator->() const
 {
-    assert(!node.expired());
+    auto shared = node.lock();
-    return &(node.lock()->data);
+    assert(shared);
+    return shared->data.get();
 }
 
 template <typename TKey, typename TData>
@@ -129,18 +131,20 @@ const TKey& CPrefixTrie<TKey, TData>::Iterator<IsConst>::key() const
 
 template <typename TKey, typename TData>
 template <bool IsConst>
-TData& CPrefixTrie<TKey, TData>::Iterator<IsConst>::data()
+typename CPrefixTrie<TKey, TData>::template Iterator<IsConst>::TDataRef CPrefixTrie<TKey, TData>::Iterator<IsConst>::data()
 {
-    assert(!node.expired());
+    auto shared = node.lock();
-    return node.lock()->data;
+    assert(shared);
+    return *(shared->data);
 }
 
 template <typename TKey, typename TData>
 template <bool IsConst>
 const TData& CPrefixTrie<TKey, TData>::Iterator<IsConst>::data() const
 {
-    assert(!node.expired());
+    auto shared = node.lock();
-    return node.lock()->data;
+    assert(shared);
+    return *(shared->data);
 }
 
 template <typename TKey, typename TData>
@@ -244,15 +248,16 @@ std::shared_ptr<typename CPrefixTrie<TKey, TData>::Node>& CPrefixTrie<TKey, TDat
         return it->second;
     }
     if (count < it->first.size()) {
-        const TKey prefix(key.begin(), key.begin() + count);
+        TKey prefix(key.begin(), key.begin() + count);
-        const TKey postfix(it->first.begin() + count, it->first.end());
+        TKey postfix(it->first.begin() + count, it->first.end());
         auto nodes = std::move(it->second);
         children.erase(it);
         ++size;
-        it = children.emplace(prefix, std::make_shared<Node>()).first;
+        it = children.emplace(std::move(prefix), std::make_shared<Node>()).first;
-        it->second->children.emplace(postfix, std::move(nodes));
+        it->second->children.emplace(std::move(postfix), std::move(nodes));
         if (key.size() == count)
             return it->second;
+        it->second->data = std::make_shared<TData>();
     }
     return insert(TKey(key.begin() + count, key.end()), it->second);
 }
@@ -269,12 +274,12 @@ void CPrefixTrie<TKey, TData>::erase(const TKey& key, std::shared_ptr<Node>& nod
     if (!find(key, node, cb))
         return;
 
-    nodes.back().second->data = {};
+    nodes.back().second->data = std::make_shared<TData>();
     for (; nodes.size() > 1; nodes.pop_back()) {
         // if we have only one child and no data ourselves, bring them up to our level
         auto& cNode = nodes.back().second;
         auto onlyOneChild = cNode->children.size() == 1;
-        auto noData = cNode->data.empty();
+        auto noData = cNode->data->empty();
         if (onlyOneChild && noData) {
             auto child = cNode->children.begin();
             auto& prefix = nodes.back().first;
@@ -282,7 +287,7 @@ void CPrefixTrie<TKey, TData>::erase(const TKey& key, std::shared_ptr<Node>& nod
             auto& postfix = child->first;
             newKey.insert(newKey.end(), postfix.begin(), postfix.end());
             auto& pNode = nodes[nodes.size() - 2].second;
-            pNode->children.emplace(newKey, std::move(child->second));
+            pNode->children.emplace(std::move(newKey), std::move(child->second));
             pNode->children.erase(prefix);
             --size;
             continue;
@@ -302,6 +307,7 @@ void CPrefixTrie<TKey, TData>::erase(const TKey& key, std::shared_ptr<Node>& nod
 template <typename TKey, typename TData>
 CPrefixTrie<TKey, TData>::CPrefixTrie() : size(0), root(std::make_shared<Node>())
 {
+    root->data = std::make_shared<TData>();
 }
 
 template <typename TKey, typename TData>
@@ -309,7 +315,7 @@ template <typename TDataUni>
 typename CPrefixTrie<TKey, TData>::iterator CPrefixTrie<TKey, TData>::insert(const TKey& key, TDataUni&& data)
 {
     auto& node = key.empty() ? root : insert(key, root);
-    node->data = std::forward<TDataUni>(data);
+    node->data = std::make_shared<TData>(std::forward<TDataUni>(data));
     return key.empty() ? begin() : iterator{key, node};
 }
 
@@ -335,7 +341,7 @@ typename CPrefixTrie<TKey, TData>::iterator CPrefixTrie<TKey, TData>::insert(CPr
         auto& node = insert(key, shared);
         copy = iterator{name, node};
     }
-    copy.node.lock()->data = std::forward<TDataUni>(data);
+    copy.node.lock()->data = std::make_shared<TData>(std::forward<TDataUni>(data));
     return copy;
 }
 
@@ -404,7 +410,7 @@ bool CPrefixTrie<TKey, TData>::erase(const TKey& key)
 {
     auto size_was = height();
     if (key.empty()) {
-        root->data = {};
+        root->data = std::make_shared<TData>();
     } else {
         erase(key, root);
     }
@@ -415,7 +421,7 @@ template <typename TKey, typename TData>
 void CPrefixTrie<TKey, TData>::clear()
 {
     size = 0;
-    root->data = {};
+    root->data = std::make_shared<TData>();
     root->children.clear();
 }
 
@@ -428,7 +434,7 @@ bool CPrefixTrie<TKey, TData>::empty() const
 template <typename TKey, typename TData>
 std::size_t CPrefixTrie<TKey, TData>::height() const
 {
-    return size + (root->data.empty() ? 0 : 1);
+    return size + (root->data->empty() ? 0 : 1);
 }
 
 template <typename TKey, typename TData>

src/prefixtrie.h

@@ -29,7 +29,7 @@ class CPrefixTrie
         Node(Node&& o) noexcept = default;
         Node& operator=(Node&& o) noexcept = default;
         Node& operator=(const Node&) = delete;
-        TData data;
+        std::shared_ptr<TData> data;
     };
 
     using TChildren = decltype(Node::children);
@@ -41,9 +41,10 @@ class CPrefixTrie
         friend class Iterator;
         friend class CPrefixTrie<TKey, TData>;
 
-        using TKeyRef = std::reference_wrapper<const TKey>;
+        using TDataRef = typename std::conditional<IsConst, const TData&, TData&>::type;
-        using TDataRef = std::reference_wrapper<TData>;
+        using TKeyWrap = std::reference_wrapper<const TKey>;
-        using TPair = std::pair<TKeyRef, TDataRef>;
+        using TDataWrap = std::reference_wrapper<TData>;
+        using TPair = std::pair<TKeyWrap, TDataWrap>;
 
         TKey name;
         std::weak_ptr<Node> node;
@@ -94,7 +95,7 @@ class CPrefixTrie
 
         const TKey& key() const;
 
-        TData& data();
+        TDataRef data();
         const TData& data() const;
 
         std::size_t depth() const;