Pending code enabled (#2427)

ErakhtinB · web-flow · commit ed144aed8888 · 2025-04-30T14:08:15.000+05:00
diff --git a/core/api/jrpc/value_converter.hpp b/core/api/jrpc/value_converter.hpp
@@ -204,6 +204,7 @@ namespace kagome::api {
     data["implVersion"] = makeValue(val.impl_version);
     data["transactionVersion"] = makeValue(val.transaction_version);
     data["stateVersion"] = makeValue(val.state_version);
+    data["systemVersion"] = makeValue(val.system_version);
 
     jArray apis;
     std::ranges::transform(
diff --git a/core/host_api/impl/host_api_factory_impl.cpp b/core/host_api/impl/host_api_factory_impl.cpp
@@ -22,7 +22,8 @@ namespace kagome::host_api {
       std::shared_ptr<crypto::KeyStore> key_store,
       std::shared_ptr<offchain::OffchainPersistentStorage>
           offchain_persistent_storage,
-      std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool)
+      std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool,
+      LazySPtr<api::StateApi> state_api)
       : offchain_config_(offchain_config),
         ecdsa_provider_(std::move(ecdsa_provider)),
         ed25519_provider_(std::move(ed25519_provider)),
@@ -35,7 +36,8 @@ namespace kagome::host_api {
         // because boost.di doesn't like optional<shared_ptr>
         key_store_(key_store ? std::optional(key_store) : std::nullopt),
         offchain_persistent_storage_(std::move(offchain_persistent_storage)),
-        offchain_worker_pool_(std::move(offchain_worker_pool)) {
+        offchain_worker_pool_(std::move(offchain_worker_pool)),
+        state_api_(std::move(state_api)) {
     BOOST_ASSERT(ecdsa_provider_ != nullptr);
     BOOST_ASSERT(ed25519_provider_ != nullptr);
     BOOST_ASSERT(sr25519_provider_ != nullptr);
@@ -62,7 +64,8 @@ namespace kagome::host_api {
                                          hasher_,
                                          key_store_,
                                          offchain_persistent_storage_,
-                                         offchain_worker_pool_);
+                                         offchain_worker_pool_,
+                                         state_api_);
   }
 
 }  // namespace kagome::host_api
diff --git a/core/host_api/impl/host_api_factory_impl.hpp b/core/host_api/impl/host_api_factory_impl.hpp
@@ -8,7 +8,9 @@
 
 #include "host_api/host_api_factory.hpp"
 
+#include "api/service/state/state_api.hpp"
 #include "host_api/impl/offchain_extension.hpp"
+#include "injector/lazy.hpp"
 
 namespace kagome::crypto {
   class EllipticCurves;
@@ -44,7 +46,8 @@ namespace kagome::host_api {
         std::shared_ptr<crypto::KeyStore> key_store,
         std::shared_ptr<offchain::OffchainPersistentStorage>
             offchain_persistent_storage,
-        std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool);
+        std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool,
+        LazySPtr<api::StateApi> state_api);
 
     std::unique_ptr<HostApi> make(
         std::shared_ptr<const runtime::CoreApiFactory> core_factory,
@@ -65,6 +68,7 @@ namespace kagome::host_api {
     std::shared_ptr<offchain::OffchainPersistentStorage>
         offchain_persistent_storage_;
     std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool_;
+    LazySPtr<api::StateApi> state_api_;
   };
 
 }  // namespace kagome::host_api
diff --git a/core/host_api/impl/host_api_impl.cpp b/core/host_api/impl/host_api_impl.cpp
@@ -84,7 +84,8 @@ namespace kagome::host_api {
       std::optional<std::shared_ptr<crypto::KeyStore>> key_store,
       std::shared_ptr<offchain::OffchainPersistentStorage>
           offchain_persistent_storage,
-      std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool)
+      std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool,
+      LazySPtr<api::StateApi> state_api)
       : memory_provider_([&] {
           BOOST_ASSERT(memory_provider);
           return std::move(memory_provider);
@@ -108,7 +109,8 @@ namespace kagome::host_api {
                   memory_provider_,
                   storage_provider_,
                   std::move(core_provider)},
-        storage_ext_(storage_provider_, memory_provider_, hasher),
+        storage_ext_(
+            storage_provider_, memory_provider_, hasher, std::move(state_api)),
         child_storage_ext_(storage_provider_, memory_provider_),
         offchain_ext_(offchain_config,
                       memory_provider_,
diff --git a/core/host_api/impl/host_api_impl.hpp b/core/host_api/impl/host_api_impl.hpp
@@ -6,8 +6,8 @@
 
 #pragma once
 
+#include "api/service/state/state_api.hpp"
 #include "host_api/host_api.hpp"
-
 #include "host_api/impl/child_storage_extension.hpp"
 #include "host_api/impl/crypto_extension.hpp"
 #include "host_api/impl/elliptic_curves_extension.hpp"
@@ -16,6 +16,7 @@
 #include "host_api/impl/misc_extension.hpp"
 #include "host_api/impl/offchain_extension.hpp"
 #include "host_api/impl/storage_extension.hpp"
+#include "injector/lazy.hpp"
 #include "offchain/impl/offchain_persistent_storage.hpp"
 
 namespace kagome::runtime {
@@ -48,7 +49,8 @@ namespace kagome::host_api {
         std::optional<std::shared_ptr<crypto::KeyStore>> key_store,
         std::shared_ptr<offchain::OffchainPersistentStorage>
             offchain_persistent_storage,
-        std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool);
+        std::shared_ptr<offchain::OffchainWorkerPool> offchain_worker_pool,
+        LazySPtr<api::StateApi> state_api);
 
     ~HostApiImpl() override = default;
 
diff --git a/core/host_api/impl/storage_extension.cpp b/core/host_api/impl/storage_extension.cpp
@@ -35,10 +35,12 @@ namespace kagome::host_api {
   StorageExtension::StorageExtension(
       std::shared_ptr<runtime::TrieStorageProvider> storage_provider,
       std::shared_ptr<const runtime::MemoryProvider> memory_provider,
-      std::shared_ptr<const crypto::Hasher> hasher)
+      std::shared_ptr<const crypto::Hasher> hasher,
+      LazySPtr<api::StateApi> state_api)
       : storage_provider_(std::move(storage_provider)),
         memory_provider_(std::move(memory_provider)),
         hasher_(std::move(hasher)),
+        state_api_(std::move(state_api)),
         logger_{log::createLogger("StorageExtension", "storage_extension")} {
     BOOST_ASSERT(storage_provider_ != nullptr);
     BOOST_ASSERT(memory_provider_ != nullptr);
@@ -126,6 +128,31 @@ namespace kagome::host_api {
     SL_TRACE_VOID_FUNC_CALL(logger_, key, value);
 
     auto batch = storage_provider_->getCurrentBatch();
+
+    // Special handling for runtime code updates based on system version
+    if (key == storage::kRuntimeCodeKey and state_api_.get()) {
+      uint8_t system_version = 0;
+
+      const auto runtime_version_res =
+          state_api_.get()->getRuntimeVersion(std::nullopt);
+      if (runtime_version_res.has_value()) {
+        system_version = runtime_version_res.value().system_version;
+      }
+
+      if (system_version >= 3) {
+        SL_INFO(logger_,
+                "Storing runtime code in :pending_code (system_version: {})",
+                system_version);
+        auto put_result = batch->put(storage::kPendingRuntimeCodeKey, value);
+        if (not put_result) {
+          logger_->error(
+              "ext_set_storage failed to store pending code, reason: {}",
+              put_result.error());
+        }
+        return;
+      }
+    }
+
     auto put_result = batch->put(key, value);
     if (not put_result) {
       logger_->error(
diff --git a/core/host_api/impl/storage_extension.hpp b/core/host_api/impl/storage_extension.hpp
@@ -8,13 +8,14 @@
 
 #include <cstdint>
 
+#include "api/service/state/state_api.hpp"
 #include "common/buffer_or_view.hpp"
+#include "injector/lazy.hpp"
 #include "log/logger.hpp"
 #include "primitives/kill_storage_result.hpp"
 #include "runtime/types.hpp"
 #include "storage/trie/serialization/polkadot_codec.hpp"
 #include "storage/trie/types.hpp"
-
 namespace kagome::crypto {
   class Hasher;
 }
@@ -33,7 +34,8 @@ namespace kagome::host_api {
     StorageExtension(
         std::shared_ptr<runtime::TrieStorageProvider> storage_provider,
         std::shared_ptr<const runtime::MemoryProvider> memory_provider,
-        std::shared_ptr<const crypto::Hasher> hasher);
+        std::shared_ptr<const crypto::Hasher> hasher,
+        LazySPtr<api::StateApi> state_api);
 
     void reset();
 
@@ -177,6 +179,7 @@ namespace kagome::host_api {
     std::shared_ptr<runtime::TrieStorageProvider> storage_provider_;
     std::shared_ptr<const runtime::MemoryProvider> memory_provider_;
     std::shared_ptr<const crypto::Hasher> hasher_;
+    LazySPtr<api::StateApi> state_api_;
     storage::trie::PolkadotCodec codec_;
     log::Logger logger_;
 
diff --git a/core/parachain/pvf/kagome_pvf_worker_injector.hpp b/core/parachain/pvf/kagome_pvf_worker_injector.hpp
@@ -95,6 +95,7 @@ namespace kagome::parachain {
         bind_null<crypto::KeyStore>(),
         bind_null<offchain::OffchainPersistentStorage>(),
         bind_null<offchain::OffchainWorkerPool>(),
+        bind_null<LazySPtr<api::StateApi>>(),
         di::bind<runtime::CoreApiFactory>.to<runtime::CoreApiFactoryImpl>(),
 
         // bound by lambda because direct binding is failing: ctor gives
@@ -112,7 +113,8 @@ namespace kagome::parachain {
               injector.template create<std::shared_ptr<crypto::Hasher>>(),
               injector.template create<std::shared_ptr<crypto::KeyStore>>(),
               injector.template create<std::shared_ptr<offchain::OffchainPersistentStorage>>(),
-              injector.template create<std::shared_ptr<offchain::OffchainWorkerPool>>()
+              injector.template create<std::shared_ptr<offchain::OffchainWorkerPool>>(),
+              injector.template create<LazySPtr<api::StateApi>>()
           );
         }),
 
diff --git a/core/primitives/version.hpp b/core/primitives/version.hpp
@@ -98,8 +98,16 @@ namespace kagome::primitives {
     uint32_t transaction_version = 1u;
 
     /// Version of the state implementation used by this runtime.
+    /// This field is kept for backward compatibility. It always has the same
+    /// value as system_version.
     uint8_t state_version = 0u;
 
+    /// Version of the system implementation used by this runtime.
+    /// Originally named state_version, but renamed to system_version in newer
+    /// versions. Both fields are maintained with identical values for
+    /// compatibility.
+    uint8_t system_version = 0u;
+
     bool operator==(const Version &rhs) const = default;
 
     /**
@@ -146,11 +154,13 @@ namespace kagome::primitives {
       }
       if (core_version and *core_version >= 4) {
         try {
-          decode(v.state_version, decoder);
+          decode(v.system_version, decoder);
+          v.state_version = v.system_version;
         } catch (std::system_error &e) {
           return outcome::failure(e.code());
         }
       } else {
+        v.system_version = 0;
         v.state_version = 0;
       }
       return v;
diff --git a/core/runtime/common/storage_code_provider.cpp b/core/runtime/common/storage_code_provider.cpp
@@ -19,11 +19,13 @@ namespace kagome::runtime {
       std::shared_ptr<RuntimeUpgradeTracker> runtime_upgrade_tracker,
       std::shared_ptr<const primitives::CodeSubstituteBlockIds>
           code_substitutes,
-      std::shared_ptr<application::ChainSpec> chain_spec)
+      std::shared_ptr<application::ChainSpec> chain_spec,
+      LazySPtr<api::StateApi> state_api)
       : storage_{std::move(storage)},
         runtime_upgrade_tracker_{std::move(runtime_upgrade_tracker)},
         known_code_substitutes_{std::move(code_substitutes)},
         chain_spec_{std::move(chain_spec)},
+        state_api_{std::move(state_api)},
         logger_{log::createLogger("StorageCodeProvider", "runtime")} {
     BOOST_ASSERT(storage_ != nullptr);
     BOOST_ASSERT(runtime_upgrade_tracker_ != nullptr);
@@ -50,4 +52,12 @@ namespace kagome::runtime {
     }
     return cached_code_;
   }
+
+  RuntimeCodeProvider::Result StorageCodeProvider::getPendingCodeAt(
+      const storage::trie::RootHash &state) const {
+    OUTCOME_TRY(batch, storage_->getEphemeralBatchAt(state));
+    OUTCOME_TRY(code, batch->get(storage::kPendingRuntimeCodeKey));
+    return std::make_shared<common::Buffer>(std::move(code));
+  }
+
 }  // namespace kagome::runtime
diff --git a/core/runtime/common/storage_code_provider.hpp b/core/runtime/common/storage_code_provider.hpp
@@ -6,10 +6,11 @@
 
 #pragma once
 
-#include "runtime/runtime_code_provider.hpp"
-
+#include "api/service/state/state_api.hpp"
 #include "application/chain_spec.hpp"
+#include "injector/lazy.hpp"
 #include "log/logger.hpp"
+#include "runtime/runtime_code_provider.hpp"
 
 namespace kagome::storage::trie {
   class TrieBatch;
@@ -29,15 +30,19 @@ namespace kagome::runtime {
         std::shared_ptr<const storage::trie::TrieStorage> storage,
         std::shared_ptr<RuntimeUpgradeTracker> runtime_upgrade_tracker,
         std::shared_ptr<const CodeSubstituteBlockIds> code_substitutes,
-        std::shared_ptr<application::ChainSpec> chain_spec);
+        std::shared_ptr<application::ChainSpec> chain_spec,
+        LazySPtr<api::StateApi> state_api);
 
     Result getCodeAt(const storage::trie::RootHash &state) const override;
+    Result getPendingCodeAt(
+        const storage::trie::RootHash &state) const override;
 
    private:
     std::shared_ptr<const storage::trie::TrieStorage> storage_;
     std::shared_ptr<RuntimeUpgradeTracker> runtime_upgrade_tracker_;
     std::shared_ptr<const CodeSubstituteBlockIds> known_code_substitutes_;
     std::shared_ptr<application::ChainSpec> chain_spec_;
+    LazySPtr<api::StateApi> state_api_;
     mutable Code cached_code_;
     mutable storage::trie::RootHash last_state_root_;
     log::Logger logger_;
diff --git a/core/runtime/runtime_code_provider.hpp b/core/runtime/runtime_code_provider.hpp
@@ -23,6 +23,8 @@ namespace kagome::runtime {
     using Result = outcome::result<Code>;
 
     virtual Result getCodeAt(const storage::trie::RootHash &state) const = 0;
+    virtual Result getPendingCodeAt(
+        const storage::trie::RootHash &state) const = 0;
   };
 
 }  // namespace kagome::runtime
diff --git a/core/storage/predefined_keys.hpp b/core/storage/predefined_keys.hpp
@@ -14,6 +14,7 @@ namespace kagome::storage {
   using common::literals::operator""_buf;
 
   inline const common::Buffer kRuntimeCodeKey = ":code"_buf;
+  inline const common::Buffer kPendingRuntimeCodeKey = ":pending_code"_buf;
   inline const common::Buffer kRuntimeHeappagesKey = ":heappages"_buf;
 
   inline const common::Buffer kExtrinsicIndexKey = ":extrinsic_index"_buf;
diff --git a/test/core/host_api/storage_extension_test.cpp b/test/core/host_api/storage_extension_test.cpp
@@ -15,6 +15,7 @@
 #include <scale/scale.hpp>
 
 #include "crypto/hasher/hasher_impl.hpp"
+#include "mock/core/api/service/state/state_api_mock.hpp"
 #include "mock/core/runtime/memory_provider_mock.hpp"
 #include "mock/core/runtime/trie_storage_provider_mock.hpp"
 #include "mock/core/storage/trie/polkadot_trie_cursor_mock.h"
@@ -23,6 +24,7 @@
 #include "scale/encode_append.hpp"
 #include "scale/kagome_scale.hpp"
 #include "storage/predefined_keys.hpp"
+#include "testutil/lazy.hpp"
 #include "testutil/literals.hpp"
 #include "testutil/outcome/dummy_error.hpp"
 #include "testutil/prepare_loggers.hpp"
@@ -31,6 +33,7 @@
 
 using kagome::ClearPrefixLimit;
 using kagome::KillStorageResult;
+using kagome::api::StateApiMock;
 using kagome::common::Buffer;
 using kagome::common::BufferView;
 using kagome::common::Hash256;
@@ -69,10 +72,18 @@ class StorageExtensionTest : public ::testing::Test {
     memory_provider_ = std::make_shared<MemoryProviderMock>();
     EXPECT_CALL(*memory_provider_, getCurrentMemory())
         .WillRepeatedly(Return(std::ref(memory_.memory)));
+    auto state_api_mock = std::make_shared<StateApiMock>();
+
+    EXPECT_CALL(
+        *state_api_mock,
+        getRuntimeVersion(std::optional<kagome::primitives::BlockHash>{}))
+        .WillRepeatedly(Return(kagome::primitives::Version{}));
+
     storage_extension_ = std::make_shared<StorageExtension>(
         storage_provider_,
         memory_provider_,
-        std::make_shared<kagome::crypto::HasherImpl>());
+        std::make_shared<kagome::crypto::HasherImpl>(),
+        testutil::sptr_to_lazy<kagome::api::StateApi>(state_api_mock));
   }
 
  protected:
diff --git a/test/core/runtime/runtime_test_base.hpp b/test/core/runtime/runtime_test_base.hpp
diff --git a/test/core/runtime/storage_code_provider_test.cpp b/test/core/runtime/storage_code_provider_test.cpp
diff --git a/test/testutil/runtime/common/basic_code_provider.cpp b/test/testutil/runtime/common/basic_code_provider.cpp
diff --git a/test/testutil/runtime/common/basic_code_provider.hpp b/test/testutil/runtime/common/basic_code_provider.hpp
