'prune-discarded' mode (#1702)

* Disable pruner for archive mode, add 'prune-discarded' mode

core/application/app_configuration.hpp

@@ -254,6 +254,8 @@ namespace kagome::application {
 
     virtual std::optional<size_t> statePruningDepth() const = 0;
 
+    virtual bool shouldPruneDiscardedStates() const = 0;
+
     /**
      * @return database state cache size in MiB
      */

core/application/impl/app_configuration_impl.cpp

@@ -790,7 +790,7 @@ namespace kagome::application {
         ("db-cache", po::value<uint32_t>()->default_value(def_db_cache_size), "Limit the memory the database cache can use <MiB>")
         ("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' or the number of finalized blocks to keep.")
+        ("state-pruning", po::value<std::string>()->default_value("archive"), "state pruning policy. 'archive', 'prune-discarded', or the number of finalized blocks to keep.")
         ;
 
     po::options_description network_desc("Network options");
@@ -1449,9 +1449,12 @@ namespace kagome::application {
     if (auto state_pruning_opt =
             find_argument<std::string>(vm, "state-pruning");
         state_pruning_opt.has_value()) {
-      const auto& val = state_pruning_opt.value();
+      const auto &val = state_pruning_opt.value();
       if (val == "archive") {
         state_pruning_depth_ = std::nullopt;
+      } else if (val == "prune-discarded") {
+        state_pruning_depth_ = std::nullopt;
+        prune_discarded_states_ = true;
       } else {
         uint32_t depth{};
         auto [_, err] = std::from_chars(&*val.begin(), &*val.end(), depth);

core/application/impl/app_configuration_impl.hpp

@@ -187,6 +187,9 @@ namespace kagome::application {
     std::optional<size_t> statePruningDepth() const override {
       return state_pruning_depth_;
     }
+    bool shouldPruneDiscardedStates() const override {
+      return state_pruning_depth_.has_value() || prune_discarded_states_;
+    }
     std::optional<std::string_view> devMnemonicPhrase() const override {
       if (dev_mnemonic_phrase_) {
         return *dev_mnemonic_phrase_;
@@ -341,6 +344,7 @@ namespace kagome::application {
     StorageBackend storage_backend_ = StorageBackend::RocksDB;
     uint32_t db_cache_size_;
     std::optional<size_t> state_pruning_depth_;
+    bool prune_discarded_states_ = false;
     std::optional<std::string> dev_mnemonic_phrase_;
     std::string node_wss_pem_;
     std::optional<BenchmarkConfigSection> benchmark_config_;

core/blockchain/impl/block_tree_impl.cpp

@@ -17,7 +17,6 @@
 #include "crypto/blake2/blake2b.h"
 #include "log/profiling_logger.hpp"
 #include "storage/database_error.hpp"
-#include "storage/trie_pruner/recover_pruner_state.hpp"
 #include "storage/trie_pruner/trie_pruner.hpp"
 
 namespace {
@@ -306,8 +305,7 @@ namespace kagome::blockchain {
           log, "Existing non-finalized block {} is added to block tree", block);
     }
 
-    OUTCOME_TRY(
-        storage::trie_pruner::recoverPrunerState(*state_pruner, *block_tree));
+    OUTCOME_TRY(state_pruner->recoverState(*block_tree));
 
     return block_tree;
   }

core/consensus/babe/impl/babe_config_repository_impl.cpp

@@ -96,10 +96,11 @@ namespace kagome::consensus::babe {
     std::unique_lock lock{indexer_mutex_};
     auto finalized = block_tree_->getLastFinalized();
     auto finalized_header = block_tree_->getBlockHeader(finalized.hash).value();
+
     if (finalized.number - indexer_.last_finalized_indexed_.number
             > kMaxUnindexedBlocksNum
         and trie_storage_->getEphemeralBatchAt(finalized_header.state_root)) {
-      warp(finalized);
+      warp(lock, finalized);
     }
 
     auto genesis_res = config({block_tree_->getGenesisBlockHash(), 0}, false);
@@ -237,9 +238,14 @@ namespace kagome::consensus::babe {
     return 0;
   }
 
+  void BabeConfigRepositoryImpl::warp(std::unique_lock<std::mutex> &lock,
+                                      const primitives::BlockInfo &block) {
+    indexer_.put(block, {}, true);
+  }
+
   void BabeConfigRepositoryImpl::warp(const primitives::BlockInfo &block) {
     std::unique_lock lock{indexer_mutex_};
-    indexer_.put(block, {}, true);
+    warp(lock, block);
   }
 
   outcome::result<std::shared_ptr<const primitives::BabeConfiguration>>

core/consensus/babe/impl/babe_config_repository_impl.hpp

@@ -52,7 +52,7 @@ namespace kagome::consensus::babe {
     std::optional<std::shared_ptr<const primitives::BabeConfiguration>> state;
     /**
      * Next epoch lazily computed from `config` and digests.
-    */
+     */
     std::optional<std::shared_ptr<const primitives::BabeConfiguration>>
         next_state;
   };
@@ -125,6 +125,9 @@ namespace kagome::consensus::babe {
     outcome::result<std::shared_ptr<const primitives::BabeConfiguration>>
     loadPrev(const std::optional<primitives::BlockInfo> &prev) const;
 
+    void warp(std::unique_lock<std::mutex> &lock,
+              const primitives::BlockInfo &block);
+
     std::shared_ptr<storage::BufferStorage> persistent_storage_;
     bool config_warp_sync_;
     std::shared_ptr<blockchain::BlockTree> block_tree_;

core/injector/application_injector.cpp

@@ -85,6 +85,7 @@
 #include "injector/bind_by_lambda.hpp"
 #include "injector/calculate_genesis_state.hpp"
 #include "injector/get_peer_keypair.hpp"
+#include "injector/idle_trie_pruner.hpp"
 #include "log/configurator.hpp"
 #include "log/logger.hpp"
 #include "metrics/impl/exposer_impl.hpp"
@@ -729,7 +730,20 @@ namespace {
             di::bind<storage::trie::PolkadotTrieFactory>.template to<storage::trie::PolkadotTrieFactoryImpl>(),
             di::bind<storage::trie::Codec>.template to<storage::trie::PolkadotCodec>(),
             di::bind<storage::trie::TrieSerializer>.template to<storage::trie::TrieSerializerImpl>(),
-            di::bind<storage::trie_pruner::TriePruner>.template to<storage::trie_pruner::TriePrunerImpl>(),
+            bind_by_lambda<storage::trie_pruner::TriePruner>(
+                [](const auto &injector)
+                    -> sptr<storage::trie_pruner::TriePruner> {
+                  const application::AppConfiguration &config =
+                      injector.template create<
+                          application::AppConfiguration const &>();
+                  if (config.statePruningDepth() == std::nullopt
+                      && !config.shouldPruneDiscardedStates()) {
+                    return std::make_shared<
+                        storage::trie_pruner::IdleTriePruner>();
+                  }
+                  return injector.template create<
+                      sptr<storage::trie_pruner::TriePrunerImpl>>();
+                }),
             di::bind<runtime::RuntimeCodeProvider>.template to<runtime::StorageCodeProvider>(),
             bind_by_lambda<application::ChainSpec>([](const auto &injector) {
               const application::AppConfiguration &config =

core/injector/idle_trie_pruner.hpp

@@ -0,0 +1,52 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "storage/trie_pruner/trie_pruner.hpp"
+
+namespace kagome::storage::trie_pruner {
+
+  /**
+   * IdleTriePruner does nothing, so it can be used to disabled pruning
+   */
+  class IdleTriePruner : public storage::trie_pruner::TriePruner {
+   public:
+    outcome::result<void> addNewState(
+        const storage::trie::RootHash &state_root,
+        storage::trie::StateVersion version) override {
+      return outcome::success();
+    }
+
+    outcome::result<void> addNewState(
+        const storage::trie::PolkadotTrie &new_trie,
+        storage::trie::StateVersion version) override {
+      return outcome::success();
+    }
+
+    outcome::result<void> pruneFinalized(
+        const primitives::BlockHeader &state) override {
+      return outcome::success();
+    }
+
+    outcome::result<void> pruneDiscarded(
+        const primitives::BlockHeader &state) override {
+      return outcome::success();
+    }
+
+    outcome::result<void> recoverState(
+        const blockchain::BlockTree &block_tree) override {
+      return outcome::success();
+    }
+
+    std::optional<primitives::BlockInfo> getLastPrunedBlock() const override {
+      return std::nullopt;
+    }
+
+    std::optional<uint32_t> getPruningDepth() const override {
+      return std::nullopt;
+    }
+  };
+}  // namespace kagome::storage::trie_pruner

core/runtime/wavm/module.cpp

@@ -115,7 +115,7 @@ namespace kagome::runtime::wavm {
 
   WAVM::Runtime::ImportBindings ModuleImpl::link(
       IntrinsicResolver &resolver) const {
-    auto &ir_module = WAVM::Runtime::getModuleIR(module_);
+    const WAVM::IR::Module &ir_module = WAVM::Runtime::getModuleIR(module_);
 
     auto link_result = WAVM::Runtime::linkModule(ir_module, resolver);
     if (!link_result.success) {

core/storage/CMakeLists.txt

@@ -25,7 +25,6 @@ add_library(storage
     trie/serialization/trie_serializer_impl.cpp
     trie/serialization/polkadot_codec.cpp
     trie_pruner/impl/trie_pruner_impl.cpp
-    trie_pruner/impl/recover_pruner_state.cpp
     )
 target_link_libraries(storage
     blob

core/storage/trie_pruner/impl/trie_pruner_impl.cpp

@@ -267,7 +267,7 @@ namespace kagome::storage::trie_pruner {
           if (value_ref_it == value_ref_count_.end()) {
             values_unknown++;
           } else {
-            auto& value_ref_count = value_ref_it->second;
+            auto &value_ref_count = value_ref_it->second;
             value_ref_count--;
             if (value_ref_count == 0) {
               OUTCOME_TRY(batch.remove(hash));
@@ -421,7 +421,55 @@ namespace kagome::storage::trie_pruner {
     return *root_hash.asHash();
   }
 
-  outcome::result<void> TriePrunerImpl::restoreState(
+  outcome::result<void> TriePrunerImpl::recoverState(
+      const blockchain::BlockTree &block_tree) {
+    static log::Logger logger =
+        log::createLogger("PrunerStateRecovery", "storage");
+    auto last_pruned_block = getLastPrunedBlock();
+    if (!last_pruned_block.has_value()) {
+      OUTCOME_TRY(first_hash_opt, block_tree.getBlockHash(1));
+      if (first_hash_opt.has_value()) {
+        SL_WARN(logger,
+                "Running pruner on a non-empty non-pruned storage may lead to "
+                "skipping some stored states.");
+        OUTCOME_TRY(
+            last_finalized,
+            block_tree.getBlockHeader(block_tree.getLastFinalized().hash));
+
+        if (auto res = restoreStateAt(last_finalized, block_tree);
+            res.has_error()) {
+          SL_ERROR(logger,
+                   "Failed to restore trie pruner state starting from last "
+                   "finalized "
+                   "block: {}",
+                   res.error().message());
+          return res.as_failure();
+        }
+      } else {
+        OUTCOME_TRY(
+            genesis_header,
+            block_tree.getBlockHeader(block_tree.getGenesisBlockHash()));
+        OUTCOME_TRY(
+            addNewState(genesis_header.state_root, trie::StateVersion::V0));
+      }
+    } else {
+      OUTCOME_TRY(base_block_header,
+                  block_tree.getBlockHeader(last_pruned_block.value().hash));
+      BOOST_ASSERT(block_tree.getLastFinalized().number
+                   >= last_pruned_block.value().number);
+      if (auto res = restoreStateAt(base_block_header, block_tree);
+          res.has_error()) {
+        SL_WARN(logger,
+                "Failed to restore trie pruner state starting from base "
+                "block {}: {}",
+                last_pruned_block.value(),
+                res.error().message());
+      }
+    }
+    return outcome::success();
+  }
+
+  outcome::result<void> TriePrunerImpl::restoreStateAt(
       const primitives::BlockHeader &last_pruned_block,
       const blockchain::BlockTree &block_tree) {
     KAGOME_PROFILE_START_L(logger_, restore_state);

core/storage/trie_pruner/impl/trie_pruner_impl.hpp

@@ -90,10 +90,6 @@ namespace kagome::storage::trie_pruner {
     virtual outcome::result<void> pruneDiscarded(
         const primitives::BlockHeader &state) override;
 
-    virtual outcome::result<void> restoreState(
-        const primitives::BlockHeader &last_pruned_block,
-        const blockchain::BlockTree &block_tree) override;
-
     std::optional<primitives::BlockInfo> getLastPrunedBlock() const override {
       return last_pruned_block_;
     }
@@ -118,7 +114,14 @@ namespace kagome::storage::trie_pruner {
       return pruning_depth_;
     }
 
+    outcome::result<void> recoverState(
+        const blockchain::BlockTree &block_tree) override;
+
    private:
+    outcome::result<void> restoreStateAt(
+        const primitives::BlockHeader &last_pruned_block,
+        const blockchain::BlockTree &block_tree);
+
     outcome::result<void> prune(BufferBatch &batch,
                                 const storage::trie::RootHash &state);