refactor: use CountedBucketMap for counting total amount of pending sigs

knst · thepastaclaw · commit bdca48cde091 · 2026-07-10T04:14:47.000-05:00
diff --git a/src/llmq/signing_shares.cpp b/src/llmq/signing_shares.cpp
@@ -206,14 +206,14 @@ bool CSigSharesNodeState::GetSessionInfoByRecvId(uint32_t sessionId, SessionInfo
     return true;
 }
 
-size_t CSigSharesNodeState::RemoveSession(const uint256& signHash)
+void CSigSharesNodeState::RemoveSession(const uint256& signHash)
 {
     if (const auto it = sessions.find(signHash); it != sessions.end()) {
         sessionByRecvId.erase(it->second.recvSessionId);
         sessions.erase(it);
     }
     requestedSigShares.EraseAllForSignHash(signHash);
-    return pendingIncomingSigShares.EraseAllForSignHash(signHash);
+    pendingIncomingSigShares.EraseAllForSignHash(signHash);
 }
 
 //////////////////////
@@ -493,16 +493,18 @@ bool CSigSharesManager::TryAddPendingIncomingSigShare(NodeId nodeId, CSigSharesN
                  __func__, MAX_PENDING_SIG_SHARES_PER_NODE, nodeId);
         return false;
     }
-    if (m_pending_sig_shares_total >= MAX_PENDING_SIG_SHARES_TOTAL) {
+    size_t total{0};
+    for (const auto& [_, ns] : nodeStates) {
+        // the size of nodeStates is limited by DEFAULT_MAX_PEER_CONNECTIONS(125) so it should not be performance issue
+        // The name of variable is intentionally mentioned in comment to make this code snippet relevant for possible changes in future
+        total += ns.pendingIncomingSigShares.Size();
+    }
+    if (total >= MAX_PENDING_SIG_SHARES_TOTAL) {
         LogPrint(BCLog::LLMQ_SIGS, "CSigSharesManager::%s -- global pending sig shares cap reached (%d), dropping sigShare. node=%d\n",
                  __func__, MAX_PENDING_SIG_SHARES_TOTAL, nodeId);
         return false;
     }
-    if (!nodeState.pendingIncomingSigShares.Add(sigShare.GetKey(), sigShare)) {
-        return false;
-    }
-    ++m_pending_sig_shares_total;
-    return true;
+    return nodeState.pendingIncomingSigShares.Add(sigShare.GetKey(), sigShare);
 }
 
 bool CSigSharesManager::CollectPendingSigSharesToVerify(
@@ -544,7 +546,7 @@ bool CSigSharesManager::CollectPendingSigSharesToVerify(
                     retSigShares[nodeId].emplace_back(sigShare);
                     ++sharesAdded;
                 }
-                m_pending_sig_shares_total -= ns.pendingIncomingSigShares.Erase(sigShare.GetKey());
+                ns.pendingIncomingSigShares.Erase(sigShare.GetKey());
                 return !ns.pendingIncomingSigShares.Empty();
             },
             rnd);
@@ -1399,7 +1401,6 @@ void CSigSharesManager::Cleanup()
             AssertLockHeld(cs);
             sigSharesRequested.Erase(k);
         });
-        m_pending_sig_shares_total -= it->second.pendingIncomingSigShares.Size();
         nodeStates.erase(nodeId);
     }
 }
@@ -1409,7 +1410,7 @@ void CSigSharesManager::RemoveSigSharesForSession(const uint256& signHash)
     AssertLockHeld(cs);
 
     for (auto& [_, nodeState] : nodeStates) {
-        m_pending_sig_shares_total -= nodeState.RemoveSession(signHash);
+        nodeState.RemoveSession(signHash);
     }
 
     sigSharesRequested.EraseAllForSignHash(signHash);
@@ -1431,7 +1432,6 @@ void CSigSharesManager::RemoveNodesIf(std::function<bool(NodeId)> predicate)
                 AssertLockHeld(cs);
                 sigSharesRequested.Erase(k);
             });
-            m_pending_sig_shares_total -= it->second.pendingIncomingSigShares.Size();
             it = nodeStates.erase(it);
         } else {
             ++it;
@@ -1460,7 +1460,7 @@ void CSigSharesManager::MarkAsBanned(NodeId nodeId)
         sigSharesRequested.Erase(k);
     });
     nodeState.requestedSigShares.Clear();
-    m_pending_sig_shares_total -= nodeState.pendingIncomingSigShares.Clear();
+    nodeState.pendingIncomingSigShares.Clear();
     nodeState.banned = true;
 }
 
diff --git a/src/llmq/signing_shares.h b/src/llmq/signing_shares.h
@@ -188,36 +188,31 @@ class CountedBucketMap
         return true;
     }
 
-    size_t Erase(const SigShareKey& k)
+    void Erase(const SigShareKey& k)
     {
         auto it = m_data.find(k.first);
         if (it == m_data.end()) {
-            return 0;
+            return;
         }
-        const size_t n = it->second.erase(k.second);
-        m_num_entries -= n;
+        m_num_entries -= it->second.erase(k.second);
         if (it->second.empty()) {
             m_data.erase(it);
         }
-        return n;
     }
 
-    size_t EraseBucket(const uint256& signHash)
+    void EraseBucket(const uint256& signHash)
     {
         auto it = m_data.find(signHash);
         if (it == m_data.end()) {
-            return 0;
+            return;
         }
-        const size_t n = it->second.size();
-        m_num_entries -= n;
+        m_num_entries -= it->second.size();
         m_data.erase(it);
-        return n;
     }
 
     template<typename F>
-    size_t EraseIf(F&& f)
+    void EraseIf(F&& f)
     {
-        size_t erased{0};
         for (auto it = m_data.begin(); it != m_data.end(); ) {
             SigShareKey k;
             k.first = it->first;
@@ -226,7 +221,6 @@ class CountedBucketMap
                 if (f(k, jt->second)) {
                     jt = it->second.erase(jt);
                     --m_num_entries;
-                    ++erased;
                 } else {
                     ++jt;
                 }
@@ -237,15 +231,12 @@ class CountedBucketMap
                 ++it;
             }
         }
-        return erased;
     }
 
-    size_t Clear()
+    void Clear()
     {
-        const size_t prev = m_num_entries;
         m_data.clear();
         m_num_entries = 0;
-        return prev;
     }
 };
 
@@ -261,14 +252,14 @@ class SigShareMap
         return internalMap.Emplace(k, v);
     }
 
-    size_t Erase(const SigShareKey& k)
+    void Erase(const SigShareKey& k)
     {
-        return internalMap.Erase(k);
+        internalMap.Erase(k);
     }
 
-    size_t Clear()
+    void Clear()
     {
-        return internalMap.Clear();
+        internalMap.Clear();
     }
 
     [[nodiscard]] bool Has(const SigShareKey& k) const
@@ -346,15 +337,15 @@ class SigShareMap
         return &it->second;
     }
 
-    size_t EraseAllForSignHash(const uint256& signHash)
+    void EraseAllForSignHash(const uint256& signHash)
     {
-        return internalMap.EraseBucket(signHash);
+        internalMap.EraseBucket(signHash);
     }
 
     template<typename F>
-    size_t EraseIf(F&& f)
+    void EraseIf(F&& f)
     {
-        return internalMap.EraseIf(f);
+        internalMap.EraseIf(f);
     }
 
     template<typename F>
@@ -423,10 +414,7 @@ class CSigSharesNodeState
     Session* GetSessionByRecvId(uint32_t sessionId);
     bool GetSessionInfoByRecvId(uint32_t sessionId, SessionInfo& retInfo);
 
-    // Returns the number of pending incoming sig shares that were dropped
-    // as part of tearing this session down; other per-session structures
-    // are erased unconditionally.
-    size_t RemoveSession(const uint256& signHash);
+    void RemoveSession(const uint256& signHash);
 };
 
 class CSignedSession
@@ -471,10 +459,6 @@ class CSigSharesManager : public llmq::CRecoveredSigsListener
     mutable Mutex cs;
 
     SigShareMap<CSigShare> sigShares GUARDED_BY(cs);
-    // Running total of shares held in nodeStates[*].pendingIncomingSigShares.
-    // Maintained explicitly so the global cap check in TryAddPendingIncomingSigShare
-    // doesn't need to walk every peer on every admission.
-    size_t m_pending_sig_shares_total GUARDED_BY(cs){0};
     Uint256HashMap<CSignedSession> signedSessions GUARDED_BY(cs);
 
     // stores time of last receivedSigShare. Used to detect timeouts
diff --git a/src/test/llmq_utils_tests.cpp b/src/test/llmq_utils_tests.cpp
@@ -114,29 +114,26 @@ BOOST_AUTO_TEST_CASE(sig_share_map_size_tracks_mutations)
     BOOST_CHECK(sig_share_map.Add(sig_share2.GetKey(), sig_share2));
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 2U);
 
-    BOOST_CHECK_EQUAL(sig_share_map.Erase(sig_share1.GetKey()), 1U);
-    BOOST_CHECK_EQUAL(sig_share_map.Erase(sig_share1.GetKey()), 0U);
+    sig_share_map.Erase(sig_share1.GetKey());
+    sig_share_map.Erase(sig_share1.GetKey());
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 1U);
 
-    BOOST_CHECK_EQUAL(sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash()), 1U);
-    BOOST_CHECK_EQUAL(sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash()), 0U);
+    sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash());
+    sig_share_map.EraseAllForSignHash(sig_share2.GetSignHash());
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 0U);
     BOOST_CHECK(sig_share_map.Empty());
 
     BOOST_CHECK(sig_share_map.Add(sig_share1.GetKey(), sig_share1));
     BOOST_CHECK(sig_share_map.Add(sig_share2.GetKey(), sig_share2));
-    BOOST_CHECK_EQUAL(sig_share_map.EraseIf([&](const SigShareKey& k, const CSigShare&) { return k == sig_share1.GetKey(); }), 1U);
+    sig_share_map.EraseIf([&](const SigShareKey& k, const CSigShare&) { return k == sig_share1.GetKey(); });
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 1U);
 
-    BOOST_CHECK_EQUAL(sig_share_map.Clear(), 1U);
-    BOOST_CHECK_EQUAL(sig_share_map.Clear(), 0U);
+    sig_share_map.Clear();
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 0U);
 }
 
-BOOST_AUTO_TEST_CASE(sig_share_map_bucket_erase_returns_bucket_size)
+BOOST_AUTO_TEST_CASE(sig_share_map_bucket_erase_updates_size)
 {
-    // EraseAllForSignHash must report the total number of entries evicted for the bucket,
-    // so callers can maintain accurate running totals (e.g. the global pending sig-share cap).
     SigShareMap<CSigShare> sig_share_map;
     const auto sign_hash = MakeSigShare(1).GetSignHash();
 
@@ -149,7 +146,7 @@ BOOST_AUTO_TEST_CASE(sig_share_map_bucket_erase_returns_bucket_size)
     }
     BOOST_CHECK_EQUAL(sig_share_map.Size(), 5U);
 
-    BOOST_CHECK_EQUAL(sig_share_map.EraseAllForSignHash(sign_hash), 5U);
+    sig_share_map.EraseAllForSignHash(sign_hash);
     BOOST_CHECK(sig_share_map.Empty());
 }
 
@@ -163,18 +160,15 @@ BOOST_AUTO_TEST_CASE(pending_sig_shares_session_removal_updates_count)
     BOOST_CHECK(node_state.pendingIncomingSigShares.Add(sig_share2.GetKey(), sig_share2));
     BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 2U);
 
-    // RemoveSession must report the number of pending shares actually dropped so the manager's
-    // global pending-share counter can be decremented in lockstep.
-    BOOST_CHECK_EQUAL(node_state.RemoveSession(sig_share1.GetSignHash()), 1U);
+    node_state.RemoveSession(sig_share1.GetSignHash());
     BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 1U);
     BOOST_CHECK(!node_state.pendingIncomingSigShares.Has(sig_share1.GetKey()));
     BOOST_CHECK(node_state.pendingIncomingSigShares.Has(sig_share2.GetKey()));
 
-    // Removing the same session twice is idempotent and reports zero on the second call.
-    BOOST_CHECK_EQUAL(node_state.RemoveSession(sig_share1.GetSignHash()), 0U);
-
-    // Removing a session that never had pending shares must also report zero.
-    BOOST_CHECK_EQUAL(node_state.RemoveSession(MakeSigShare(3).GetSignHash()), 0U);
+    // Removing the same session twice, or a session with no pending shares, is a no-op.
+    node_state.RemoveSession(sig_share1.GetSignHash());
+    node_state.RemoveSession(MakeSigShare(3).GetSignHash());
+    BOOST_CHECK_EQUAL(node_state.pendingIncomingSigShares.Size(), 1U);
 }
 
 BOOST_AUTO_TEST_CASE(deterministic_outbound_connection_test)

Original file line number	Diff line number	Diff line change
`@@ -188,36 +188,31 @@ class CountedBucketMap`
`188`	`188`	`return true;`
`189`	`189`	`}`
`190`	`190`
`191`		`- size_t Erase(const SigShareKey& k)`
	`191`	`+ void Erase(const SigShareKey& k)`
`192`	`192`	`{`
`193`	`193`	`auto it = m_data.find(k.first);`
`194`	`194`	`if (it == m_data.end()) {`
`195`		`- return 0;`
	`195`	`+ return;`
`196`	`196`	`}`
`197`		`- const size_t n = it->second.erase(k.second);`
`198`		`- m_num_entries -= n;`
	`197`	`+ m_num_entries -= it->second.erase(k.second);`
`199`	`198`	`if (it->second.empty()) {`
`200`	`199`	`m_data.erase(it);`
`201`	`200`	`}`
`202`		`- return n;`
`203`	`201`	`}`
`204`	`202`
`205`		`- size_t EraseBucket(const uint256& signHash)`
	`203`	`+ void EraseBucket(const uint256& signHash)`
`206`	`204`	`{`
`207`	`205`	`auto it = m_data.find(signHash);`
`208`	`206`	`if (it == m_data.end()) {`
`209`		`- return 0;`
	`207`	`+ return;`
`210`	`208`	`}`
`211`		`- const size_t n = it->second.size();`
`212`		`- m_num_entries -= n;`
	`209`	`+ m_num_entries -= it->second.size();`
`213`	`210`	`m_data.erase(it);`
`214`		`- return n;`
`215`	`211`	`}`
`216`	`212`
`217`	`213`	`template<typename F>`
`218`		`- size_t EraseIf(F&& f)`
	`214`	`+ void EraseIf(F&& f)`
`219`	`215`	`{`
`220`		`- size_t erased{0};`
`221`	`216`	`for (auto it = m_data.begin(); it != m_data.end(); ) {`
`222`	`217`	`SigShareKey k;`
`223`	`218`	`k.first = it->first;`
`@@ -226,7 +221,6 @@ class CountedBucketMap`
`226`	`221`	`if (f(k, jt->second)) {`
`227`	`222`	`jt = it->second.erase(jt);`
`228`	`223`	`--m_num_entries;`
`229`		`- ++erased;`
`230`	`224`	`} else {`
`231`	`225`	`++jt;`
`232`	`226`	`}`
`@@ -237,15 +231,12 @@ class CountedBucketMap`
`237`	`231`	`++it;`
`238`	`232`	`}`
`239`	`233`	`}`
`240`		`- return erased;`
`241`	`234`	`}`
`242`	`235`
`243`		`- size_t Clear()`
	`236`	`+ void Clear()`
`244`	`237`	`{`
`245`		`- const size_t prev = m_num_entries;`
`246`	`238`	`m_data.clear();`
`247`	`239`	`m_num_entries = 0;`
`248`		`- return prev;`
`249`	`240`	`}`
`250`	`241`	`};`
`251`	`242`
`@@ -261,14 +252,14 @@ class SigShareMap`
`261`	`252`	`return internalMap.Emplace(k, v);`
`262`	`253`	`}`
`263`	`254`
`264`		`- size_t Erase(const SigShareKey& k)`
	`255`	`+ void Erase(const SigShareKey& k)`
`265`	`256`	`{`
`266`		`- return internalMap.Erase(k);`
	`257`	`+ internalMap.Erase(k);`
`267`	`258`	`}`
`268`	`259`
`269`		`- size_t Clear()`
	`260`	`+ void Clear()`
`270`	`261`	`{`
`271`		`- return internalMap.Clear();`
	`262`	`+ internalMap.Clear();`
`272`	`263`	`}`
`273`	`264`
`274`	`265`	`[[nodiscard]] bool Has(const SigShareKey& k) const`
`@@ -346,15 +337,15 @@ class SigShareMap`
`346`	`337`	`return &it->second;`
`347`	`338`	`}`
`348`	`339`
`349`		`- size_t EraseAllForSignHash(const uint256& signHash)`
	`340`	`+ void EraseAllForSignHash(const uint256& signHash)`
`350`	`341`	`{`
`351`		`- return internalMap.EraseBucket(signHash);`
	`342`	`+ internalMap.EraseBucket(signHash);`
`352`	`343`	`}`
`353`	`344`
`354`	`345`	`template<typename F>`
`355`		`- size_t EraseIf(F&& f)`
	`346`	`+ void EraseIf(F&& f)`
`356`	`347`	`{`
`357`		`- return internalMap.EraseIf(f);`
	`348`	`+ internalMap.EraseIf(f);`
`358`	`349`	`}`
`359`	`350`
`360`	`351`	`template<typename F>`
`@@ -423,10 +414,7 @@ class CSigSharesNodeState`
`423`	`414`	`Session* GetSessionByRecvId(uint32_t sessionId);`
`424`	`415`	`bool GetSessionInfoByRecvId(uint32_t sessionId, SessionInfo& retInfo);`
`425`	`416`
`426`		`- // Returns the number of pending incoming sig shares that were dropped`
`427`		`- // as part of tearing this session down; other per-session structures`
`428`		`- // are erased unconditionally.`
`429`		`- size_t RemoveSession(const uint256& signHash);`
	`417`	`+ void RemoveSession(const uint256& signHash);`
`430`	`418`	`};`
`431`	`419`
`432`	`420`	`class CSignedSession`
`@@ -471,10 +459,6 @@ class CSigSharesManager : public llmq::CRecoveredSigsListener`
`471`	`459`	`mutable Mutex cs;`
`472`	`460`
`473`	`461`	`SigShareMap<CSigShare> sigShares GUARDED_BY(cs);`
`474`		`- // Running total of shares held in nodeStates[*].pendingIncomingSigShares.`
`475`		`- // Maintained explicitly so the global cap check in TryAddPendingIncomingSigShare`
`476`		`- // doesn't need to walk every peer on every admission.`
`477`		`- size_t m_pending_sig_shares_total GUARDED_BY(cs){0};`
`478`	`462`	`Uint256HashMap<CSignedSession> signedSessions GUARDED_BY(cs);`
`479`	`463`
`480`	`464`	`// stores time of last receivedSigShare. Used to detect timeouts`