diff --git a/cumulus/Cargo.lock b/cumulus/Cargo.lock index 86cda7a662..c320c9d7dc 100644 --- a/cumulus/Cargo.lock +++ b/cumulus/Cargo.lock @@ -2907,6 +2907,7 @@ dependencies = [ "sp-tracing", "sp-trie", "sp-version", + "trie-db", "xcm", ] diff --git a/cumulus/pallets/parachain-system/Cargo.toml b/cumulus/pallets/parachain-system/Cargo.toml index f3d8a3bee4..61660d4de8 100644 --- a/cumulus/pallets/parachain-system/Cargo.toml +++ b/cumulus/pallets/parachain-system/Cargo.toml @@ -11,6 +11,7 @@ codec = { package = "parity-scale-codec", version = "3.0.0", default-features = environmental = { version = "1.1.4", default-features = false } impl-trait-for-tuples = "0.2.1" log = { version = "0.4.19", default-features = false } +trie-db = { version = "0.27.1", default-features = false } scale-info = { version = "2.9.0", default-features = false, features = ["derive"] } # Substrate @@ -69,6 +70,7 @@ std = [ "sp-std/std", "sp-trie/std", "xcm/std", + "trie-db/std", ] runtime-benchmarks = [ diff --git a/cumulus/pallets/parachain-system/src/validate_block/implementation.rs b/cumulus/pallets/parachain-system/src/validate_block/implementation.rs index f953dfc77c..71a6338dcd 100644 --- a/cumulus/pallets/parachain-system/src/validate_block/implementation.rs +++ b/cumulus/pallets/parachain-system/src/validate_block/implementation.rs @@ -16,7 +16,7 @@ //! The actual implementation of the validate block functionality. -use super::MemoryOptimizedValidationParams; +use super::{trie_cache, MemoryOptimizedValidationParams}; use cumulus_primitives_core::{ relay_chain::Hash as RHash, ParachainBlockData, PersistedValidationData, }; @@ -34,7 +34,11 @@ use sp_runtime::traits::{Block as BlockT, Extrinsic, HashFor, Header as HeaderT} use sp_std::prelude::*; use sp_trie::MemoryDB; -type TrieBackend = sp_state_machine::TrieBackend>, HashFor>; +type TrieBackend = sp_state_machine::TrieBackend< + MemoryDB>, + HashFor, + trie_cache::CacheProvider>, +>; type Ext<'a, B> = sp_state_machine::Ext<'a, HashFor, TrieBackend>; @@ -114,10 +118,15 @@ where sp_std::mem::drop(storage_proof); + let cache_provider = trie_cache::CacheProvider::new(); // We use the storage root of the `parent_head` to ensure that it is the correct root. // This is already being done above while creating the in-memory db, but let's be paranoid!! - let backend = - sp_state_machine::TrieBackendBuilder::new(db, *parent_header.state_root()).build(); + let backend = sp_state_machine::TrieBackendBuilder::new_with_cache( + db, + *parent_header.state_root(), + cache_provider, + ) + .build(); let _guard = ( // Replace storage calls with our own implementations diff --git a/cumulus/pallets/parachain-system/src/validate_block/mod.rs b/cumulus/pallets/parachain-system/src/validate_block/mod.rs index e641c77f89..4e387bf849 100644 --- a/cumulus/pallets/parachain-system/src/validate_block/mod.rs +++ b/cumulus/pallets/parachain-system/src/validate_block/mod.rs @@ -22,6 +22,10 @@ pub mod implementation; #[cfg(test)] mod tests; +#[cfg(not(feature = "std"))] +#[doc(hidden)] +mod trie_cache; + #[cfg(not(feature = "std"))] #[doc(hidden)] pub use bytes; diff --git a/cumulus/pallets/parachain-system/src/validate_block/trie_cache.rs b/cumulus/pallets/parachain-system/src/validate_block/trie_cache.rs new file mode 100644 index 0000000000..57876b8787 --- /dev/null +++ b/cumulus/pallets/parachain-system/src/validate_block/trie_cache.rs @@ -0,0 +1,104 @@ +// This file is part of Substrate. + +// Copyright (C) Parity Technologies (UK) Ltd. +// SPDX-License-Identifier: Apache-2.0 + +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use sp_state_machine::TrieCacheProvider; +use sp_std::{ + boxed::Box, + cell::{RefCell, RefMut}, + collections::btree_map::{BTreeMap, Entry}, +}; +use sp_trie::NodeCodec; +use trie_db::{node::NodeOwned, Hasher}; + +/// Special purpose trie cache implementation that is able to cache an unlimited number +/// of values. To be used in `validate_block` to serve values and nodes that +/// have already been loaded and decoded from the storage proof. +pub(crate) struct TrieCache<'a, H: Hasher> { + node_cache: RefMut<'a, BTreeMap>>, + value_cache: Option, trie_db::CachedValue>>>, +} + +impl<'a, H: Hasher> trie_db::TrieCache> for TrieCache<'a, H> { + fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&trie_db::CachedValue> { + self.value_cache.as_ref().and_then(|cache| cache.get(key)) + } + + fn cache_value_for_key(&mut self, key: &[u8], value: trie_db::CachedValue) { + self.value_cache.as_mut().and_then(|cache| cache.insert(key.into(), value)); + } + + fn get_or_insert_node( + &mut self, + hash: as trie_db::NodeCodec>::HashOut, + fetch_node: &mut dyn FnMut() -> trie_db::Result< + NodeOwned, + H::Out, + as trie_db::NodeCodec>::Error, + >, + ) -> trie_db::Result<&NodeOwned, H::Out, as trie_db::NodeCodec>::Error> { + match self.node_cache.entry(hash) { + Entry::Occupied(entry) => Ok(entry.into_mut()), + Entry::Vacant(entry) => Ok(entry.insert(fetch_node()?)), + } + } + + fn get_node( + &mut self, + hash: &H::Out, + ) -> Option<&NodeOwned< as trie_db::NodeCodec>::HashOut>> { + self.node_cache.get(hash) + } +} + +/// Provider of [`TrieCache`] instances. +pub(crate) struct CacheProvider { + node_cache: RefCell>>, + value_cache: RefCell, trie_db::CachedValue>>, +} + +impl CacheProvider { + /// Constructs a new instance of [`CacheProvider`] with an uninitialized state + /// and empty node and value caches. + pub fn new() -> Self { + CacheProvider { node_cache: Default::default(), value_cache: Default::default() } + } +} + +impl TrieCacheProvider for CacheProvider { + type Cache<'a> = TrieCache<'a, H> where H: 'a; + + fn as_trie_db_cache(&self, _storage_root: ::Out) -> Self::Cache<'_> { + TrieCache { + value_cache: Some(self.value_cache.borrow_mut()), + node_cache: self.node_cache.borrow_mut(), + } + } + + fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> { + // This method is called when we calculate the storage root. + // Since we are using a simplified cache architecture, + // we do not have separate key spaces for different storage roots. + // The value cache is therefore disabled here. + TrieCache { value_cache: None, node_cache: self.node_cache.borrow_mut() } + } + + fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: ::Out) {} +} + +// This is safe here since we are single-threaded in WASM +unsafe impl Send for CacheProvider {} +unsafe impl Sync for CacheProvider {}