Use snapshottable map for auditing in 2PC shard

Signed-off-by: James Lovejoy <[email protected]>

Author: metalicjames

src/uhs/twophase/locking_shard/controller.cpp

@@ -169,6 +169,7 @@ namespace cbdc::locking_shard {
                    || audit_epoch <= m_last_audit_epoch) {
                     continue;
                 }
+                m_logger->info("Starting audit for", audit_epoch);
 
                 auto maybe_total = m_shard->audit(audit_epoch);
                 if(!maybe_total.has_value()) {
@@ -186,6 +187,8 @@ namespace cbdc::locking_shard {
                 m_last_audit_epoch = audit_epoch;
 
                 m_shard->prune(audit_epoch);
+
+                m_logger->info("Prune completed for", audit_epoch);
             }
         }
     }

src/uhs/twophase/locking_shard/locking_shard.cpp

@@ -32,18 +32,13 @@ namespace cbdc::locking_shard {
           m_logger(std::move(logger)),
           m_completed_txs(completed_txs_cache_size),
           m_opts(std::move(opts)) {
-        m_uhs.max_load_factor(std::numeric_limits<float>::max());
         m_applied_dtxs.max_load_factor(std::numeric_limits<float>::max());
         m_prepared_dtxs.max_load_factor(std::numeric_limits<float>::max());
-        m_locked.max_load_factor(std::numeric_limits<float>::max());
 
         static constexpr auto dtx_buckets = 100000;
         m_applied_dtxs.rehash(dtx_buckets);
         m_prepared_dtxs.rehash(dtx_buckets);
 
-        static constexpr auto locked_buckets = 10000000;
-        m_locked.rehash(locked_buckets);
-
         if(!preseed_file.empty()) {
             m_logger->info("Reading preseed file into memory");
             if(!read_preseed_file(preseed_file)) {
@@ -65,12 +60,9 @@ namespace cbdc::locking_shard {
             }
             in.seekg(0, std::ios::beg);
             auto deser = istream_serializer(in);
-            m_uhs.clear();
-            static constexpr auto uhs_size_factor = 2;
-            auto bucket_count = static_cast<unsigned long>(sz / cbdc::hash_size
-                                                           * uhs_size_factor);
-            m_uhs.rehash(bucket_count);
-            deser >> m_uhs;
+            auto uhs = std::map<hash_t, uhs_element>();
+            deser >> uhs;
+            m_uhs = std::move(uhs);
             return true;
         }
         return false;
@@ -215,22 +207,17 @@ namespace cbdc::locking_shard {
         }
     }
 
-    auto locking_shard::audit(uint64_t epoch) const
-        -> std::optional<uint64_t> {
-        auto uhs = decltype(m_uhs)();
-        auto locked = decltype(m_locked)();
-        auto spent = decltype(m_spent)();
+    auto locking_shard::audit(uint64_t epoch) -> std::optional<uint64_t> {
         {
-            std::shared_lock l(m_mut);
-            uhs = m_uhs;
-            locked = m_locked;
-            spent = m_spent;
+            std::unique_lock l(m_mut);
+            m_uhs.snapshot();
+            m_locked.snapshot();
+            m_spent.snapshot();
         }
 
         auto check_uhs
-            = [epoch](
-                  const std::unordered_map<hash_t, uhs_element, hashing::null>&
-                      uhs_map) -> std::optional<uint64_t> {
+            = [epoch](const snapshot_map<hash_t, uhs_element>& uhs_map)
+            -> std::optional<uint64_t> {
             uint64_t tot{};
             for(const auto& [id, elem] : uhs_map) {
                 if(elem.m_creation_epoch <= epoch
@@ -247,24 +234,34 @@ namespace cbdc::locking_shard {
             return tot;
         };
 
+        // TODO: use bignum because tot will probably overflow
         uint64_t tot{};
-        auto uhs_tot = check_uhs(uhs);
-        if(!uhs_tot.has_value()) {
-            return std::nullopt;
+        auto uhs_tot = check_uhs(m_uhs);
+        if(uhs_tot.has_value()) {
+            tot += uhs_tot.value();
         }
-        tot += uhs_tot.value();
 
-        auto locked_tot = check_uhs(locked);
-        if(!locked_tot.has_value()) {
-            return std::nullopt;
+        auto locked_tot = check_uhs(m_locked);
+        if(locked_tot.has_value()) {
+            tot += locked_tot.value();
+        }
+
+        auto spent_tot = check_uhs(m_spent);
+        if(spent_tot.has_value()) {
+            tot += spent_tot.value();
+        }
+
+        {
+            std::unique_lock l(m_mut);
+            m_uhs.release_snapshot();
+            m_locked.release_snapshot();
+            m_spent.release_snapshot();
         }
-        tot += locked_tot.value();
 
-        auto spent_tot = check_uhs(spent);
-        if(!spent_tot.has_value()) {
+        if(!spent_tot.has_value() || !locked_tot.has_value()
+           || !uhs_tot.has_value()) {
             return std::nullopt;
         }
-        tot += spent_tot.value();
 
         return tot;
     }
@@ -277,7 +274,7 @@ namespace cbdc::locking_shard {
     void locking_shard::prune(uint64_t epoch) {
         std::unique_lock l(m_mut);
         for(auto it = m_spent.begin(); it != m_spent.end();) {
-            auto& elem = it->second;
+            const auto& elem = it->second;
             if(elem.m_deletion_epoch.value() < epoch) {
                 it = m_spent.erase(it);
             } else {

src/uhs/twophase/locking_shard/locking_shard.hpp

@@ -14,6 +14,7 @@
 #include "util/common/hash.hpp"
 #include "util/common/hashmap.hpp"
 #include "util/common/logging.hpp"
+#include "util/common/snapshot_map.hpp"
 
 #include <filesystem>
 #include <future>
@@ -143,7 +144,7 @@ namespace cbdc::locking_shard {
         /// \param epoch the epoch to audit the supply at.
         /// \return total value of coins in this shard's UHS, or std::nullopt
         ///         if any of the UHS elements do not match their UHS ID.
-        auto audit(uint64_t epoch) const -> std::optional<uint64_t>;
+        auto audit(uint64_t epoch) -> std::optional<uint64_t>;
 
         /// Prunes any spent UHS elements spent prior to the given epoch.
         /// \param epoch epoch to prune prior to.
@@ -167,9 +168,9 @@ namespace cbdc::locking_shard {
 
         std::shared_ptr<logging::log> m_logger;
         mutable std::shared_mutex m_mut;
-        std::unordered_map<hash_t, uhs_element, hashing::null> m_uhs;
-        std::unordered_map<hash_t, uhs_element, hashing::null> m_locked;
-        std::unordered_map<hash_t, uhs_element, hashing::null> m_spent;
+        snapshot_map<hash_t, uhs_element> m_uhs;
+        snapshot_map<hash_t, uhs_element> m_locked;
+        snapshot_map<hash_t, uhs_element> m_spent;
         std::unordered_map<hash_t, prepared_dtx, hashing::null>
             m_prepared_dtxs;
         std::unordered_set<hash_t, hashing::null> m_applied_dtxs;

src/util/common/snapshot_map.hpp

@@ -0,0 +1,130 @@
+// Copyright (c) 2022 MIT Digital Currency Initiative,
+//                    Federal Reserve Bank of Boston
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or http://www.opensource.org/licenses/mit-license.php.
+
+#ifndef CBDC_UNIVERSE0_SRC_COMMON_SNAPSHOT_MAP_H_
+#define CBDC_UNIVERSE0_SRC_COMMON_SNAPSHOT_MAP_H_
+
+#include <cassert>
+#include <map>
+#include <set>
+
+namespace cbdc {
+    template<typename K, typename V, typename... Ts>
+    class snapshot_map {
+      public:
+        using map_type = std::map<K, V, Ts...>;
+        using set_type = std::set<K, Ts...>;
+
+        auto operator=(map_type&& m) -> snapshot_map<K, V, Ts...>& {
+            m_map = std::move(m);
+            return *this;
+        }
+
+        void release_snapshot() {
+            m_snapshot = false;
+        }
+
+        void snapshot() {
+            m_snapshot = true;
+        }
+
+        [[nodiscard]] auto find(const K& key) const ->
+            typename map_type::const_iterator {
+            auto added_it = m_added.find(key);
+            if(added_it != m_added.end()) {
+                return added_it;
+            }
+
+            auto removed_it = m_removed.find(key);
+            if(removed_it != m_removed.end()) {
+                return m_map.end();
+            }
+
+            return m_map.find(key);
+        }
+
+        [[nodiscard]] auto size() const -> size_t {
+            return m_added.size() + m_removed.size() + m_map.size();
+        }
+
+        [[nodiscard]] auto end() const -> typename map_type::const_iterator {
+            return m_map.end();
+        }
+
+        [[nodiscard]] auto begin() const -> typename map_type::const_iterator {
+            return m_map.begin();
+        }
+
+        template<class... Args>
+        auto emplace(Args&&... args)
+            -> std::pair<typename map_type::iterator, bool> {
+            gc();
+            auto ret = [&]() {
+                if(m_snapshot) {
+                    return m_added.emplace(std::forward<Args>(args)...);
+                }
+                auto it = m_map.emplace(std::forward<Args>(args)...);
+                m_added.erase(it.first->first);
+                return it;
+            }();
+            m_removed.erase(ret.first->first);
+            return ret;
+        }
+
+        auto erase(typename map_type::const_iterator it) ->
+            typename map_type::iterator {
+            assert(!m_snapshot);
+            m_added.erase(it->first);
+            m_removed.erase(it->first);
+            return m_map.erase(it);
+        }
+
+        void erase(const K& key) {
+            gc();
+            m_added.erase(key);
+            if(m_snapshot) {
+                m_removed.emplace(key);
+            } else {
+                m_removed.erase(key);
+                m_map.erase(key);
+            }
+        }
+
+      private:
+        map_type m_map;
+        map_type m_added;
+        set_type m_removed;
+        bool m_snapshot{false};
+
+        void gc() {
+            if(m_snapshot) {
+                return;
+            }
+            constexpr size_t factor = 1000000;
+            constexpr size_t one = 1;
+            auto added_elems = std::min(std::max(m_added.size() / factor, one),
+                                        m_added.size());
+            size_t count = 0;
+            for(auto it = m_added.begin();
+                it != m_added.end() && count < added_elems;
+                count++) {
+                auto n = m_added.extract(it++);
+                m_map.insert(std::move(n));
+            }
+            count = 0;
+            auto removed_elems
+                = std::min(std::max(m_removed.size() / factor, one),
+                           m_removed.size());
+            for(auto it = m_removed.begin();
+                it != m_removed.end() && count < removed_elems;
+                count++) {
+                auto n = m_removed.extract(it++);
+                m_map.erase(n.value());
+            }
+        }
+    };
+}
+
+#endif

src/util/serialization/format.hpp

@@ -284,6 +284,32 @@ namespace cbdc {
         return deser;
     }
 
+    /// Deserializes a map of key-value pairs.
+    template<typename K, typename V, typename... Ts>
+    auto operator>>(serializer& deser, std::map<K, V, Ts...>& map)
+        -> serializer& {
+        auto len = uint64_t();
+        if(!(deser >> len)) {
+            return deser;
+        }
+
+        for(uint64_t i = 0; i < len; i++) {
+            auto key = K();
+            if(!(deser >> key)) {
+                return deser;
+            }
+
+            auto val = V();
+            if(!(deser >> val)) {
+                return deser;
+            }
+
+            map.emplace(std::move(key), std::move(val));
+        }
+
+        return deser;
+    }
+
     /// Serializes the count of items, and then each item statically-casted.
     /// \see \ref cbdc::operator<<(serializer&, T)
     template<typename K, typename... Ts>