Fix child batches and introduce callback for processing loaded trie n…

…odes (#1499)

* Introduce opportunity to process loaded nodes while performing operations over trie

* Make a common base for batch types

* Fix child trie storage transactions

core/blockchain/impl/justification_storage_policy.cpp

@@ -14,15 +14,17 @@ namespace kagome::blockchain {
 
   outcome::result<bool> JustificationStoragePolicyImpl::shouldStoreFor(
       primitives::BlockHeader const &block_header) const {
-    if (block_header.number == 0) return true;
+    if (block_header.number == 0) {
+      return true;
+    }
 
     BOOST_ASSERT_MSG(block_tree_ != nullptr,
                      "Block tree must have been initialized with "
                      "JustificationStoragePolicyImpl::initBlockchainInfo()");
 
-    auto last_finalized = block_tree_->getLastFinalized();
-    BOOST_ASSERT_MSG(last_finalized.number >= block_header.number,
-                     "Target block must be finalized");
+    BOOST_ASSERT_MSG(
+        block_tree_->getLastFinalized().number >= block_header.number,
+        "Target block must be finalized");
 
     if (consensus::grandpa::HasAuthoritySetChange{block_header}) {
       return true;

core/host_api/impl/child_storage_extension.cpp

@@ -38,14 +38,25 @@ namespace kagome::host_api {
   }
 
   template <typename R, typename F, typename... Args>
-  outcome::result<R> ChildStorageExtension::executeOnChildStorage(
+  outcome::result<R> ChildStorageExtension::executeOnConstChildStorage(
       const Buffer &child_storage_key, F func, Args &&...args) const {
     OUTCOME_TRY(prefixed_child_key,
                 make_prefixed_child_storage_key(child_storage_key));
     OUTCOME_TRY(child_batch,
                 storage_provider_->getChildBatchAt(prefixed_child_key));
 
-    return func(child_batch, std::forward<Args>(args)...);
+    return func(child_batch.get(), std::forward<Args>(args)...);
+  }
+
+  template <typename R, typename F, typename... Args>
+  outcome::result<R> ChildStorageExtension::executeOnMutChildStorage(
+      const Buffer &child_storage_key, F func, Args &&...args) const {
+    OUTCOME_TRY(prefixed_child_key,
+                make_prefixed_child_storage_key(child_storage_key));
+    OUTCOME_TRY(child_batch,
+                storage_provider_->getMutableChildBatchAt(prefixed_child_key));
+
+    return func(child_batch.get(), std::forward<Args>(args)...);
   }
 
   template <typename Arg>
@@ -71,9 +82,9 @@ namespace kagome::host_api {
     SL_TRACE_VOID_FUNC_CALL(
         logger_, child_key_buffer, key_buffer, value_buffer);
 
-    auto result = executeOnChildStorage<void>(
+    auto result = executeOnMutChildStorage<void>(
         child_key_buffer, [&](auto &child_batch) mutable {
-          return child_batch->put(key_buffer, std::move(value_buffer));
+          return child_batch.put(key_buffer, std::move(value_buffer));
         });
 
     if (not result) {
@@ -92,7 +103,7 @@ namespace kagome::host_api {
 
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer, key_buffer);
 
-    return executeOnChildStorage<runtime::WasmSpan>(
+    return executeOnConstChildStorage<runtime::WasmSpan>(
                child_key_buffer,
                [&memory, &child_key_buffer, &key_buffer, this](
                    auto &child_batch, auto &key) {
@@ -101,7 +112,7 @@ namespace kagome::host_api {
                      "value was not "
                      "obtained. Reason: {}";
 
-                 auto result = child_batch->tryGet(key);
+                 auto result = child_batch.tryGet(key);
                  if (result) {
                    SL_TRACE_FUNC_CALL(
                        logger_, result.value(), child_key_buffer, key_buffer);
@@ -130,9 +141,9 @@ namespace kagome::host_api {
 
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer, key_buffer);
 
-    auto result = executeOnChildStorage<void>(
+    auto result = executeOnMutChildStorage<void>(
         child_key_buffer,
-        [](auto &child_batch, auto &key) { return child_batch->remove(key); },
+        [](auto &child_batch, auto &key) { return child_batch.remove(key); },
         key_buffer);
 
     if (not result) {
@@ -157,16 +168,16 @@ namespace kagome::host_api {
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer, key_buffer);
 
     auto child_batch_outcome =
-        storage_provider_->getChildBatchAt(prefixed_child_key);
+        storage_provider_->getMutableChildBatchAt(prefixed_child_key);
     if (child_batch_outcome.has_error()) {
       logger_->error(
           "ext_default_child_storage_next_key_version_1 resulted with error: "
           "{}",
           child_batch_outcome.error());
       return kErrorSpan;
     }
-    auto child_batch = child_batch_outcome.value();
-    auto cursor = child_batch->trieCursor();
+    auto &child_batch = child_batch_outcome.value().get();
+    auto cursor = child_batch.trieCursor();
 
     auto seek_result = cursor->seekUpperBound(key_buffer);
     if (seek_result.has_error()) {
@@ -200,8 +211,11 @@ namespace kagome::host_api {
     auto child_key_buffer = loadBuffer(memory, child_storage_key);
     auto prefixed_child_key = make_prefixed_child_storage_key(child_key_buffer);
     auto child_batch =
-        storage_provider_->getChildBatchAt(prefixed_child_key.value()).value();
-    auto res = child_batch->commit(storage::trie::StateVersion::V0);
+        storage_provider_->getMutableChildBatchAt(prefixed_child_key.value())
+            .value();
+
+    auto res = child_batch.get().commit(storage::trie::StateVersion::V0);
+
     if (res.has_error()) {
       logger_->error(
           "ext_default_child_storage_root resulted with an error: {}",
@@ -220,10 +234,10 @@ namespace kagome::host_api {
 
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer, prefix);
 
-    auto result = executeOnChildStorage<std::tuple<bool, uint32_t>>(
+    auto result = executeOnMutChildStorage<std::tuple<bool, uint32_t>>(
         child_key_buffer,
         [](auto &child_batch, auto &prefix) {
-          return child_batch->clearPrefix(prefix, std::nullopt);
+          return child_batch.clearPrefix(prefix, std::nullopt);
         },
         prefix_buffer);
 
@@ -246,10 +260,13 @@ namespace kagome::host_api {
         loadBuffer(memory, child_storage_key, key);
     auto [value_ptr, value_size] = runtime::PtrSize(value_out);
 
-    auto value = executeOnChildStorage<std::optional<common::BufferOrView>>(
-        child_key_buffer,
-        [](auto &child_batch, auto &key) { return child_batch->tryGet(key); },
-        key_buffer);
+    auto value =
+        executeOnConstChildStorage<std::optional<common::BufferOrView>>(
+            child_key_buffer,
+            [](auto &child_batch, auto &key) {
+              return child_batch.tryGet(key);
+            },
+            key_buffer);
     std::optional<uint32_t> res{std::nullopt};
     if (value) {
       auto &data_opt = value.value();
@@ -290,9 +307,9 @@ namespace kagome::host_api {
 
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer, key_buffer);
 
-    auto res = executeOnChildStorage<bool>(
+    auto res = executeOnConstChildStorage<bool>(
         child_key_buffer,
-        [](auto &child_batch, auto &key) { return child_batch->contains(key); },
+        [](auto &child_batch, auto &key) { return child_batch.contains(key); },
         key_buffer);
 
     if (not res) {
@@ -312,9 +329,9 @@ namespace kagome::host_api {
 
     SL_TRACE_VOID_FUNC_CALL(logger_, child_key_buffer);
 
-    auto result = executeOnChildStorage<std::tuple<bool, uint32_t>>(
+    auto result = executeOnMutChildStorage<std::tuple<bool, uint32_t>>(
         child_key_buffer, [](auto &child_batch) {
-          return child_batch->clearPrefix({}, std::nullopt);
+          return child_batch.clearPrefix({}, std::nullopt);
         });
 
     if (not result) {

core/host_api/impl/child_storage_extension.hpp

@@ -98,7 +98,11 @@ namespace kagome::host_api {
         "WASM Runtime [ChildStorageExtension]";
 
     template <typename R, typename F, typename... Args>
-    outcome::result<R> executeOnChildStorage(
+    outcome::result<R> executeOnConstChildStorage(
+        const common::Buffer &child_storage_key, F func, Args &&...args) const;
+
+    template <typename R, typename F, typename... Args>
+    outcome::result<R> executeOnMutChildStorage(
         const common::Buffer &child_storage_key, F func, Args &&...args) const;
   };
 

core/host_api/impl/storage_extension.cpp

@@ -229,7 +229,7 @@ namespace kagome::host_api {
   runtime::WasmSpan StorageExtension::ext_storage_root_version_2(
       runtime::WasmI32 version) {
     auto state_version = toStateVersion(version);
-    auto res = storage_provider_->forceCommit(state_version);
+    auto res = storage_provider_->commit(state_version);
     if (res.has_error()) {
       logger_->error("ext_storage_root resulted with an error: {}",
                      res.error());

core/runtime/common/executor.hpp

@@ -77,7 +77,7 @@ namespace kagome::runtime {
                                    std::string_view name,
                                    Args &&...args) {
       OUTCOME_TRY(env, env_factory_->start(block_info, storage_state)->make());
-      return callMediateInternal<Result>(
+      return callWithCache<Result>(
           *env, name, std::forward<Args>(args)...);
     }
 
@@ -92,7 +92,7 @@ namespace kagome::runtime {
                                    Args &&...args) {
       OUTCOME_TRY(env_template, env_factory_->start(block_hash));
       OUTCOME_TRY(env, env_template->make());
-      return callMediateInternal<Result>(
+      return callWithCache<Result>(
           *env, name, std::forward<Args>(args)...);
     }
 
@@ -106,7 +106,7 @@ namespace kagome::runtime {
                                           Args &&...args) {
       OUTCOME_TRY(env_template, env_factory_->start());
       OUTCOME_TRY(env, env_template->make());
-      return callMediateInternal<Result>(
+      return callWithCache<Result>(
           *env, name, std::forward<Args>(args)...);
     }
 
@@ -133,36 +133,6 @@ namespace kagome::runtime {
       return result;
     }
 
-    /**
-     * Internal method for calling a Runtime API method
-     * Resets the runtime memory with the module's heap base,
-     * encodes the arguments with SCALE codec, calls the method from the
-     * provided module instance and  returns a result, decoded from SCALE.
-     * Changes, made to the Host API state, are reset after the call.
-     */
-    template <typename Result, typename... Args>
-    inline outcome::result<Result> callMediateInternal(RuntimeEnvironment &env,
-                                                       std::string_view name,
-                                                       Args &&...args) {
-      if constexpr (std::is_same_v<Result, primitives::Version>) {
-        if (likely(name == "Core_version")) {
-          return cache_->getVersion(env.module_instance->getCodeHash(), [&] {
-            return call<Result>(env, name, std::forward<Args>(args)...);
-          });
-        }
-      }
-
-      if constexpr (std::is_same_v<Result, primitives::OpaqueMetadata>) {
-        if (likely(name == "Metadata_metadata")) {
-          return cache_->getMetadata(env.module_instance->getCodeHash(), [&] {
-            return call<Result>(env, name, std::forward<Args>(args)...);
-          });
-        }
-      }
-
-      return call<Result>(env, name, std::forward<Args>(args)...);
-    }
-
     /**
      * Internal method for calling a Runtime API method
      * Resets the runtime memory with the module's heap base,
@@ -218,6 +188,30 @@ namespace kagome::runtime {
     }
 
    private:
+    // returns cached results for some common runtime calls
+    template <typename Result, typename... Args>
+    inline outcome::result<Result> callWithCache(RuntimeEnvironment &env,
+                                                       std::string_view name,
+                                                       Args &&...args) {
+      if constexpr (std::is_same_v<Result, primitives::Version>) {
+        if (likely(name == "Core_version")) {
+          return cache_->getVersion(env.module_instance->getCodeHash(), [&] {
+            return call<Result>(env, name, std::forward<Args>(args)...);
+          });
+        }
+      }
+
+      if constexpr (std::is_same_v<Result, primitives::OpaqueMetadata>) {
+        if (likely(name == "Metadata_metadata")) {
+          return cache_->getMetadata(env.module_instance->getCodeHash(), [&] {
+            return call<Result>(env, name, std::forward<Args>(args)...);
+          });
+        }
+      }
+
+      return call<Result>(env, name, std::forward<Args>(args)...);
+    }
+
     std::shared_ptr<RuntimeEnvironmentFactory> env_factory_;
     std::shared_ptr<RuntimePropertiesCache> cache_;
     log::Logger logger_;

core/runtime/common/runtime_environment_factory.cpp

@@ -141,6 +141,13 @@ namespace kagome::runtime {
     return *this;
   }
 
+  [[nodiscard]] RuntimeEnvironmentFactory::RuntimeEnvironmentTemplate &
+  RuntimeEnvironmentFactory::RuntimeEnvironmentTemplate::withStorageBatch(
+      std::shared_ptr<storage::trie::TrieBatch> batch) {
+    batch_ = batch;
+    return *this;
+  }
+
   outcome::result<std::unique_ptr<RuntimeEnvironment>>
   RuntimeEnvironmentFactory::RuntimeEnvironmentTemplate::make() {
     KAGOME_PROFILE_START(runtime_env_making);
@@ -166,7 +173,10 @@ namespace kagome::runtime {
                     header_res.value()));
 
     const auto &env = instance->getEnvironment();
-    if (persistent_) {
+    if (batch_) {
+      env.storage_provider->setTo(batch_);
+
+    } else if (persistent_) {
       if (auto res = env.storage_provider->setToPersistentAt(storage_state_);
           !res) {
         SL_DEBUG(

core/runtime/common/storage_code_provider.cpp

@@ -54,7 +54,7 @@ namespace kagome::runtime {
   }
 
   outcome::result<common::Buffer> StorageCodeProvider::setCodeFromBatch(
-      const storage::trie::EphemeralTrieBatch &batch) const {
+      const storage::trie::TrieBatch &batch) const {
     OUTCOME_TRY(code, batch.get(storage::kRuntimeCodeKey));
     common::Buffer uncompressed;
     OUTCOME_TRY(uncompressCodeIfNeeded(code, uncompressed));

core/runtime/common/storage_code_provider.hpp

@@ -12,7 +12,7 @@
 #include "log/logger.hpp"
 
 namespace kagome::storage::trie {
-  class EphemeralTrieBatch;
+  class TrieBatch;
   class TrieStorage;
 }  // namespace kagome::storage::trie
 
@@ -36,7 +36,7 @@ namespace kagome::runtime {
 
    private:
     outcome::result<common::Buffer> setCodeFromBatch(
-        const storage::trie::EphemeralTrieBatch &batch) const;
+        const storage::trie::TrieBatch &batch) const;
     std::shared_ptr<const storage::trie::TrieStorage> storage_;
     std::shared_ptr<RuntimeUpgradeTracker> runtime_upgrade_tracker_;
     std::shared_ptr<const CodeSubstituteBlockIds> known_code_substitutes_;