Faster pruner (#2376)

* Several improvements to trie pruner

* Move merkle value caching to trie nodes

* Fix pruner caching

* Add pruner benchmark 

* Fix db editor

* Ignore empty child trees on polkadot

* Fix valueandhash comparator

CMakeLists.txt

@@ -69,7 +69,7 @@ endif ()
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # cmake-format: off
-option(TESTING "Build and run test suite" ON)
+option(TESTING "Build test suite" ON)
 option(CLANG_FORMAT "Enable clang-format target" ON)
 option(CLANG_TIDY "Enable clang-tidy checks during compilation" OFF)
 option(COVERAGE "Enable generation of coverage info" OFF)
@@ -79,6 +79,8 @@ option(BACKWARD "Enable stacktrace logging instruments" ON)
 option(CLEAR_OBJS "Clear object files" OFF)
 option(WERROR "Turn all warnings into errors" OFF)
 
+option(KAGOME_BENCHMARK "Build benchmark suite" ON)
+
 if (NOT ($ENV{CI}) OR NOT ($ENV{GITHUB_ACTIONS}))
   set(_EXTERNAL_PROJECT_DEFAULT ON)
 endif ()
@@ -141,6 +143,10 @@ if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "^(AppleClang|Clang|GNU)$")
     add_flag(-Werror)                # turn all warnings into errors
   endif()
 
+  if (KAGOME_BENCHMARK)
+    add_flag(-fno-omit-frame-pointer)
+  endif()
+
   # cmake-format: on
   if ((("${CMAKE_CXX_COMPILER_ID}" MATCHES "^(AppleClang|Clang)$")
       AND (${CMAKE_CXX_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.0"))
@@ -256,4 +262,11 @@ if(TESTING)
   add_subdirectory(test)
 endif ()
 
+if(KAGOME_BENCHMARK)
+  if (NOT TESTING)
+    add_subdirectory(test/testutil)
+  endif()
+  add_subdirectory(benchmark)
+endif()
+
 add_subdirectory(node)

benchmark/CMakeLists.txt

@@ -0,0 +1,15 @@
+#
+# Copyright Quadrivium LLC
+# All Rights Reserved
+# SPDX-License-Identifier: Apache-2.0
+#
+
+add_executable(trie_pruner_benchmark storage/trie_pruner_benchmark.cpp)
+target_link_libraries(trie_pruner_benchmark 
+    storage 
+    benchmark::benchmark 
+    GTest::gmock_main 
+    hasher 
+    log_configurator
+)
+target_include_directories(trie_pruner_benchmark PRIVATE "${CMAKE_SOURCE_DIR}/test")

benchmark/storage/trie_pruner_benchmark.cpp

@@ -0,0 +1,163 @@
+/**
+ * Copyright Quadrivium LLC
+ * All Rights Reserved
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <benchmark/benchmark.h>
+#include <rocksdb/options.h>
+
+#include <boost/filesystem/operations.hpp>
+#include <cstdlib>
+#include <memory>
+#include <random>
+#include <soralog/level.hpp>
+#include <soralog/macro.hpp>
+
+#include "crypto/hasher/hasher_impl.hpp"
+#include "gmock/gmock.h"
+#include "log/logger.hpp"
+#include "primitives/common.hpp"
+#include "storage/in_memory/in_memory_spaced_storage.hpp"
+#include "storage/rocksdb/rocksdb.hpp"
+#include "storage/trie/impl/trie_storage_backend_impl.hpp"
+#include "storage/trie/polkadot_trie/polkadot_trie_factory.hpp"
+#include "storage/trie/polkadot_trie/polkadot_trie_factory_impl.hpp"
+#include "storage/trie/serialization/polkadot_codec.hpp"
+#include "storage/trie/serialization/trie_serializer_impl.hpp"
+#include "storage/trie/types.hpp"
+#include "storage/trie_pruner/impl/trie_pruner_impl.hpp"
+#include "testutil/prepare_loggers.hpp"
+
+#include "mock/core/application/app_configuration_mock.hpp"
+#include "mock/core/application/app_state_manager_mock.hpp"
+#include "utils/thread_pool.hpp"
+
+namespace storage = kagome::storage;
+namespace trie = storage::trie;
+
+struct TriePrunerBenchmark {
+  TriePrunerBenchmark() {
+    testutil::prepareLoggers(soralog::Level::DEBUG);
+    app_state_manager =
+        std::make_shared<kagome::application::AppStateManagerMock>();
+    EXPECT_CALL(*app_state_manager, atPrepare(testing::_))
+        .WillRepeatedly(testing::Return());
+    EXPECT_CALL(*app_state_manager, atLaunch(testing::_))
+        .WillRepeatedly(testing::Return());
+    EXPECT_CALL(*app_state_manager, atShutdown(testing::_))
+        .WillRepeatedly(testing::Return());
+
+    app_config = std::make_shared<kagome::application::AppConfigurationMock>();
+    EXPECT_CALL(*app_config, statePruningDepth())
+        .WillRepeatedly(testing::Return(100));
+    EXPECT_CALL(*app_config, enableThoroughPruning())
+        .WillRepeatedly(testing::Return(true));
+
+    hasher = std::make_shared<kagome::crypto::HasherImpl>();
+    codec = std::make_shared<trie::PolkadotCodec>();
+    rocksdb::Options options{};
+    options.create_if_missing = true;
+    storage =
+        storage::RocksDb::create(
+            std::filesystem::path((boost::filesystem::temp_directory_path()
+                                   / "kagome_pruner_benchmark"
+                                   / boost::filesystem::unique_path())
+                                      .string()),
+            options)
+            .value();
+    storage_backend = std::make_shared<trie::TrieStorageBackendImpl>(storage);
+    trie_factory = std::make_shared<trie::PolkadotTrieFactoryImpl>();
+    serializer = std::make_shared<trie::TrieSerializerImpl>(
+        trie_factory, codec, storage_backend);
+    thread_pool = std::make_shared<kagome::common::WorkerThreadPool>(
+        kagome::TestThreadPool{});
+  }
+
+  auto createPruner() {
+    return std::make_unique<kagome::storage::trie_pruner::TriePrunerImpl>(
+        app_state_manager,
+        storage_backend,
+        serializer,
+        codec,
+        storage,
+        hasher,
+        app_config,
+        thread_pool);
+  }
+
+  std::shared_ptr<kagome::application::AppStateManagerMock> app_state_manager;
+  std::shared_ptr<kagome::application::AppConfigurationMock> app_config;
+  std::shared_ptr<kagome::crypto::HasherImpl> hasher;
+  std::shared_ptr<trie::PolkadotCodec> codec;
+  std::shared_ptr<storage::SpacedStorage> storage;
+  std::shared_ptr<trie::TrieStorageBackendImpl> storage_backend;
+  std::shared_ptr<trie::PolkadotTrieFactoryImpl> trie_factory;
+  std::shared_ptr<trie::TrieSerializerImpl> serializer;
+  std::shared_ptr<kagome::common::WorkerThreadPool> thread_pool;
+};
+
+auto createRandomTrie(trie::PolkadotTrieFactory &factory,
+                      size_t values_num,
+                      size_t max_value_len) {
+  std::mt19937_64 random;
+
+  auto trie = factory.createEmpty(trie::PolkadotTrie::RetrieveFunctions{});
+  for (size_t i = 0; i < values_num; i++) {
+    storage::Buffer key;
+    key.resize(random() % 128);
+    for (auto &byte : key) {
+      byte = random() % 256;
+    }
+    storage::Buffer value;
+    value.resize(random() % max_value_len);
+    for (auto &byte : value) {
+      byte = random() % 256;
+    }
+    trie->put(key, std::move(value)).value();
+  }
+  return trie;
+}
+
+static void registerStateBenchmark(benchmark::State &state) {
+  TriePrunerBenchmark benchmark;
+
+  auto trie = createRandomTrie(*benchmark.trie_factory, 10000, 70);
+  auto logger = kagome::log::createLogger("Benchmark");
+
+  for (auto _ : state) {
+    auto pruner = benchmark.createPruner();
+    pruner->addNewState(*trie, kagome::storage::trie::StateVersion::V1).value();
+  }
+}
+
+static void pruneStateBenchmark(benchmark::State &state) {
+  TriePrunerBenchmark benchmark;
+
+  auto trie = createRandomTrie(*benchmark.trie_factory, 10000, 70);
+
+  for (auto _ : state) {
+    auto pruner = benchmark.createPruner();
+    pruner->addNewState(*trie, trie::StateVersion::V1).value();
+    auto [root, batch] =
+        benchmark.serializer
+            ->storeTrie(*trie, kagome::storage::trie::StateVersion::V1)
+            .value();
+    pruner
+        ->pruneFinalized(
+            root,
+            kagome::primitives::BlockInfo{kagome::primitives::BlockHash{}, 0})
+        .value();
+    batch->commit().value();
+  }
+}
+
+BENCHMARK(registerStateBenchmark)
+    ->Unit(benchmark::TimeUnit::kMillisecond)
+    ->Iterations(10);
+
+BENCHMARK(pruneStateBenchmark)
+    ->Unit(benchmark::TimeUnit::kMillisecond)
+    ->Iterations(10);
+
+BENCHMARK_MAIN();

cmake/Hunter/config.cmake

@@ -111,3 +111,10 @@ hunter_config(
     KEEP_PACKAGE_SOURCES
 )
 
+hunter_config(
+    wabt
+    URL https://github.com/qdrvm/wabt/archive/d1abde6977d4d3f81024a7eb099228107f27377e.tar.gz
+    SHA1 0917dd9e8300c888ef0005b027008efe001abda4
+    KEEP_PACKAGE_SOURCES
+)
+

cmake/dependencies.cmake

@@ -1,7 +1,7 @@
 # hunter dependencies
 # https://docs.hunter.sh/en/latest/packages/
 
-if (TESTING)
+if (TESTING OR KAGOME_BENCHMARK)
   # https://docs.hunter.sh/en/latest/packages/pkg/GTest.html
   hunter_add_package(GTest)
   find_package(GTest CONFIG REQUIRED)

cmake/toolchain/flags/sanitize_undefined.cmake

@@ -11,11 +11,17 @@ set(FLAGS
     -fno-omit-frame-pointer
     -g
     -O0
-    -fsanitize-ignorelist="${CMAKE_CURRENT_LIST_DIR}/ubsan_ignore.txt"
     )
+if(CMAKE_CXX_COMPILER_ID MATCHES ".*Clang")
+  list(APPEND FLAGS -fsanitize-ignorelist="${CMAKE_CURRENT_LIST_DIR}/ubsan_ignore.txt")
+else()
+  message(WARNING "Non-Clang compilers do not support -fsanitize-ignorelist flag, some known false positives are expected.")
+endif()
+
 if (UBSAN_ABORT) 
   list(APPEND FLAGS -fno-sanitize-recover=undefined)
 endif()
+
 if (UBSAN_TRAP)
   list(APPEND FLAGS -fsanitize-trap=undefined)
 endif()

core/application/app_configuration.hpp

@@ -218,8 +218,8 @@ namespace kagome::application {
      * List of telemetry endpoints specified via CLI argument or config file
      * @return a vector of parsed telemetry endpoints
      */
-    virtual const std::vector<telemetry::TelemetryEndpoint>
-        &telemetryEndpoints() const = 0;
+    virtual const std::vector<telemetry::TelemetryEndpoint> &
+    telemetryEndpoints() const = 0;
 
     /**
      * @return enum constant of the chosen sync method

core/benchmark/block_execution_benchmark.cpp

@@ -16,6 +16,7 @@
 #include "runtime/module_repository.hpp"
 #include "runtime/runtime_api/core.hpp"
 #include "storage/trie/trie_storage.hpp"
+#include "utils/pretty_duration.hpp"
 
 OUTCOME_CPP_DEFINE_CATEGORY(kagome::benchmark,
                             BlockExecutionBenchmark::Error,
@@ -48,55 +49,6 @@ OUTCOME_CPP_DEFINE_CATEGORY(kagome::benchmark,
 #define OUTCOME_TRY_MSG_VOID(expr, ...) \
   _OUTCOME_TRY_MSG_VOID(OUTCOME_UNIQUE, expr, __VA_ARGS__)
 
-namespace {
-
-  template <typename Rep, typename Period>
-  struct pretty_duration {
-    std::chrono::duration<Rep, Period> dur;
-  };
-
-  template <typename Rep, typename Period>
-  pretty_duration(std::chrono::duration<Rep, Period>)
-      -> pretty_duration<Rep, Period>;
-
-  const char *suffix(unsigned denominator) {
-    switch (denominator) {
-      case 1:
-        return "s";
-      case 1000:
-        return "ms";
-      case 1'000'000:
-        return "us";
-      case 1'000'000'000:
-        return "ns";
-      default:
-        return "??";
-    }
-  }
-
-}  // namespace
-
-template <typename Rep, typename Period>
-struct fmt::formatter<pretty_duration<Rep, Period>> {
-  constexpr auto parse(format_parse_context &ctx)
-      -> format_parse_context::iterator {
-    return ctx.end();
-  }
-
-  auto format(const pretty_duration<Rep, Period> &p, format_context &ctx) const
-      -> format_context::iterator {
-    auto denominator = 1;
-    static_assert(Period::num == 1);
-    while (p.dur.count() / denominator > 1000 && denominator < Period::den) {
-      denominator *= 1000;
-    }
-    return fmt::format_to(ctx.out(),
-                          "{:.2f} {}",
-                          static_cast<double>(p.dur.count()) / denominator,
-                          suffix(Period::den / denominator));
-  }
-};
-
 namespace kagome::benchmark {
 
   using common::literals::operator""_hex2buf;

core/blockchain/impl/block_tree_impl.cpp

@@ -8,6 +8,7 @@
 
 #include <algorithm>
 #include <set>
+#include <soralog/macro.hpp>
 #include <stack>
 
 #include "blockchain/block_tree_error.hpp"
@@ -1310,7 +1311,10 @@ namespace kagome::blockchain {
           }
           extrinsics.emplace_back(std::move(ext));
         }
-        OUTCOME_TRY(p.state_pruner_->pruneDiscarded(block_header));
+        p.state_pruner_->schedulePrune(
+            block_header.state_root,
+            block_header.blockInfo(),
+            storage::trie_pruner::PruneReason::Discarded);
       }
       retired_hashes.emplace_back(
           primitives::events::RemoveAfterFinalizationParams::HeaderInfo{
@@ -1350,7 +1354,10 @@ namespace kagome::blockchain {
       BOOST_ASSERT(next_hash_opt.has_value());
       auto &next_hash = *next_hash_opt;
       OUTCOME_TRY(header, getBlockHeader(hash));
-      OUTCOME_TRY(block_tree_data.state_pruner_->pruneFinalized(header));
+      block_tree_data.state_pruner_->schedulePrune(
+          header.state_root,
+          header.blockInfo(),
+          storage::trie_pruner::PruneReason::Finalized);
       hash = next_hash;
     }
 

core/common/buffer_view.hpp

@@ -126,8 +126,8 @@ struct fmt::formatter<kagome::common::BufferView> {
   // Formats the Blob using the parsed format specification (presentation)
   // stored in this formatter.
   template <typename FormatContext>
-  auto format(const kagome::common::BufferView &view, FormatContext &ctx) const
-      -> decltype(ctx.out()) {
+  auto format(const kagome::common::BufferView &view,
+              FormatContext &ctx) const -> decltype(ctx.out()) {
     // ctx.out() is an output iterator to write to.
 
     if (view.empty()) {