Fast pruner start (#1756)

* Fast pruner startup

* Add option for thorough pruning

---------

Co-authored-by: kamilsa <kamilsa16@gmail.com>

core/application/app_configuration.hpp

@@ -258,6 +258,8 @@ namespace kagome::application {
 
     virtual bool shouldPruneDiscardedStates() const = 0;
 
+    virtual bool enableThoroughPruning() const = 0;
+
     /**
      * @return database state cache size in MiB
      */

core/application/impl/app_configuration_impl.cpp

@@ -250,7 +250,7 @@ namespace kagome::application {
         recovery_state_{def_block_to_recover},
         db_cache_size_{def_db_cache_size},
         state_pruning_depth_{} {
-    SL_INFO(logger_, "Soramitsu Kagome started. Version: {} ", buildVersion());
+    SL_INFO(logger_, "Kagome started. Version: {} ", buildVersion());
   }
 
   fs::path AppConfigurationImpl::chainSpecPath() const {
@@ -792,6 +792,7 @@ namespace kagome::application {
         ("enable-offchain-indexing", po::value<bool>(), "enable Offchain Indexing API, which allow block import to write to offchain DB)")
         ("recovery", po::value<std::string>(), "recovers block storage to state after provided block presented by number or hash, and stop after that")
         ("state-pruning", po::value<std::string>()->default_value("archive"), "state pruning policy. 'archive', 'prune-discarded', or the number of finalized blocks to keep.")
+        ("enable-thorough-pruning", po::bool_switch(), "Makes trie node pruner more efficient, but the node starts slowly")
         ;
 
     po::options_description network_desc("Network options");
@@ -884,6 +885,8 @@ namespace kagome::application {
     }
 
     if (vm.count("help") > 0) {
+      std::cout
+          << "Available subcommands: storage-explorer db-editor benchmark\n";
       std::cout << desc << std::endl;
       return false;
     }
@@ -1478,6 +1481,10 @@ namespace kagome::application {
           return false;
         }
       }
+
+      if (find_argument(vm, "enable-thorough-pruning")) {
+        enable_thorough_pruning_ = true;
+      }
     }
 
     // if something wrong with config print help message

core/application/impl/app_configuration_impl.hpp

@@ -194,6 +194,9 @@ namespace kagome::application {
     bool shouldPruneDiscardedStates() const override {
       return state_pruning_depth_.has_value() || prune_discarded_states_;
     }
+    bool enableThoroughPruning() const override {
+      return enable_thorough_pruning_;
+    }
     std::optional<std::string_view> devMnemonicPhrase() const override {
       if (dev_mnemonic_phrase_) {
         return *dev_mnemonic_phrase_;
@@ -349,6 +352,7 @@ namespace kagome::application {
     uint32_t db_cache_size_;
     std::optional<size_t> state_pruning_depth_;
     bool prune_discarded_states_ = false;
+    bool enable_thorough_pruning_ = false;
     std::optional<std::string> dev_mnemonic_phrase_;
     std::string node_wss_pem_;
     std::optional<BenchmarkConfigSection> benchmark_config_;

core/runtime/common/memory_allocator.cpp

@@ -101,8 +101,7 @@ namespace kagome::runtime {
     }
 
     resize(new_size);
-    auto current_size = memory_.getSize();
-    BOOST_ASSERT(current_size >= new_size);
+    BOOST_ASSERT(memory_.getSize() >= new_size);
     return allocate(allocation_sz);
   }
 

core/storage/trie/impl/persistent_trie_batch_impl.cpp

@@ -43,8 +43,8 @@ namespace kagome::storage::trie {
   outcome::result<RootHash> PersistentTrieBatchImpl::commit(
       StateVersion version) {
     OUTCOME_TRY(commitChildren(version));
-    OUTCOME_TRY(root, serializer_->storeTrie(*trie_, version));
     OUTCOME_TRY(state_pruner_->addNewState(*trie_, version));
+    OUTCOME_TRY(root, serializer_->storeTrie(*trie_, version));
     SL_TRACE_FUNC_CALL(logger_, root);
     return std::move(root);
   }

core/storage/trie_pruner/impl/trie_pruner_impl.cpp

@@ -70,7 +70,8 @@ namespace kagome::storage::trie_pruner {
         codec_{codec},
         storage_{storage},
         hasher_{hasher},
-        pruning_depth_{config->statePruningDepth()} {
+        pruning_depth_{config->statePruningDepth()},
+        thorough_pruning_{config->enableThoroughPruning()} {
     BOOST_ASSERT(trie_storage_ != nullptr);
     BOOST_ASSERT(serializer_ != nullptr);
     BOOST_ASSERT(codec_ != nullptr);
@@ -249,20 +250,28 @@ namespace kagome::storage::trie_pruner {
       }
 
       auto &ref_count = ref_count_it->second;
-      BOOST_ASSERT(ref_count != 0);
+      if (ref_count == 0) {
+        SL_WARN(logger_,
+                "Pruner encountered an unindexed node {} while pruning, this "
+                "indicates a bug",
+                hash);
+        continue;
+      }
       ref_count--;
       SL_TRACE(logger_,
                "Prune - {} - Node {}, ref count {}",
                depth,
                hash,
                ref_count);
 
-      if (ref_count == 0) {
+      if (immortal_nodes_.find(hash) == immortal_nodes_.end()
+          && ref_count == 0) {
         nodes_removed++;
         ref_count_.erase(ref_count_it);
         OUTCOME_TRY(batch.remove(hash));
         auto hash_opt = node->getValue().hash;
         if (hash_opt.has_value()) {
+          auto &hash = *hash_opt;
           auto value_ref_it = value_ref_count_.find(hash);
           if (value_ref_it == value_ref_count_.end()) {
             values_unknown++;
@@ -330,15 +339,13 @@ namespace kagome::storage::trie_pruner {
 
   outcome::result<void> TriePrunerImpl::addNewState(
       const storage::trie::RootHash &state_root, trie::StateVersion version) {
-    std::optional<std::scoped_lock<std::mutex>> lock;
     OUTCOME_TRY(trie, serializer_->retrieveTrie(state_root));
     OUTCOME_TRY(addNewStateWith(*trie, version));
     return outcome::success();
   }
 
   outcome::result<void> TriePrunerImpl::addNewState(
       const trie::PolkadotTrie &new_trie, trie::StateVersion version) {
-    std::optional<std::scoped_lock<std::mutex>> lock;
     OUTCOME_TRY(addNewStateWith(new_trie, version));
     return outcome::success();
   }
@@ -373,7 +380,21 @@ namespace kagome::storage::trie_pruner {
     while (!queued_nodes.empty()) {
       auto [node, hash] = queued_nodes.back();
       queued_nodes.pop_back();
-      const size_t ref_count = ++ref_count_[hash];
+      auto &ref_count = ref_count_[hash];
+      if (ref_count == 0 && !thorough_pruning_) {
+        OUTCOME_TRY(hash_is_in_storage, trie_storage_->contains(hash));
+        if (hash_is_in_storage) {
+          // the node is present in storage but pruner has not indexed it
+          // because pruner has been initialized on a newer state
+          SL_TRACE(
+              logger_,
+              "Node {} is unindexed, but already in storage, make it immortal",
+              hash.toHex());
+          ref_count++;
+          immortal_nodes_.emplace(hash);
+        }
+      }
+      ref_count++;
       SL_TRACE(logger_, "Add node {}, ref count {}", hash.toHex(), ref_count);
 
       referenced_nodes_num++;
@@ -396,6 +417,8 @@ namespace kagome::storage::trie_pruner {
             OUTCOME_TRY(child, serializer_->retrieveNode(opaque_child));
             OUTCOME_TRY(child_merkle_val,
                         encoder.getMerkleValue(*child, version));
+            // otherwise it is not stored as a separated node, but as a part of
+            // the branch
             if (child_merkle_val.isHash()) {
               SL_TRACE(logger_, "Queue child {}", child_merkle_val.asBuffer());
               queued_nodes.push_back({child, *child_merkle_val.asHash()});

core/storage/trie_pruner/impl/trie_pruner_impl.hpp

@@ -134,6 +134,7 @@ namespace kagome::storage::trie_pruner {
     std::mutex ref_count_mutex_;
     std::unordered_map<common::Hash256, size_t> ref_count_;
     std::unordered_map<common::Hash256, size_t> value_ref_count_;
+    std::unordered_set<common::Hash256> immortal_nodes_;
 
     std::optional<primitives::BlockInfo> last_pruned_block_;
     std::shared_ptr<storage::trie::TrieStorageBackend> trie_storage_;
@@ -143,6 +144,7 @@ namespace kagome::storage::trie_pruner {
     std::shared_ptr<const crypto::Hasher> hasher_;
 
     const std::optional<uint32_t> pruning_depth_{};
+    const bool thorough_pruning_{false};
     log::Logger logger_ = log::createLogger("TriePruner", "trie_pruner");
   };
 

test/core/storage/trie_pruner/trie_pruner_test.cpp

@@ -165,6 +165,7 @@ class TriePrunerTest : public testing::Test {
     auto config_mock =
         std::make_shared<kagome::application::AppConfigurationMock>();
     ON_CALL(*config_mock, statePruningDepth()).WillByDefault(Return(16));
+    ON_CALL(*config_mock, enableThoroughPruning()).WillByDefault(Return(true));
 
     trie_storage_mock.reset(
         new testing::NiceMock<trie::TrieStorageBackendMock>());
@@ -203,6 +204,7 @@ class TriePrunerTest : public testing::Test {
     auto config_mock =
         std::make_shared<kagome::application::AppConfigurationMock>();
     ON_CALL(*config_mock, statePruningDepth()).WillByDefault(Return(16));
+    ON_CALL(*config_mock, enableThoroughPruning()).WillByDefault(Return(true));
     trie_pruner::TriePrunerImpl::TriePrunerInfo info{.last_pruned_block =
                                                          last_pruned};
 
@@ -628,6 +630,9 @@ TEST_F(TriePrunerTest, RestoreStateFromGenesis) {
   ON_CALL(*block_tree, bestLeaf())
       .WillByDefault(Return(BlockInfo{6, hash_from_header(headers.at(6))}));
 
+  ON_CALL(*block_tree, bestLeaf())
+      .WillByDefault(Return(BlockInfo{6, hash_from_header(headers.at(6))}));
+
   auto mock_block = [&](unsigned int number) {
     auto str_number = std::to_string(number);
 

test/mock/core/application/app_configuration_mock.hpp

@@ -156,6 +156,8 @@ namespace kagome::application {
 
     MOCK_METHOD(bool, shouldPruneDiscardedStates, (), (const, override));
 
+    MOCK_METHOD(bool, enableThoroughPruning, (), (const, override));
+
     MOCK_METHOD(StorageBackend, storageBackend, (), (const, override));
 
     MOCK_METHOD(uint32_t, dbCacheSize, (), (const, override));