validate-block: Fix TrieCache implementation (#2214)

The trie cache implementation was ignoring the `storage_root` when
setting up the value cache. The problem with this is that the value
cache works using `storage_keys` and these keys are not unique across
different tries. A block can actually have different tries (main trie
and multiple child tries). This pull request fixes the issue by not
ignoring the `storage_root` and returning an unique `value_cache` per
`storage_root`. It also adds a test for the seen bug and improves
documentation that this doesn't happen again.

cumulus/pallets/parachain-system/src/validate_block/tests.rs

@@ -21,12 +21,13 @@ use cumulus_test_client::{
 	runtime::{
 		self as test_runtime, Block, Hash, Header, TestPalletCall, UncheckedExtrinsic, WASM_BINARY,
 	},
-	transfer, BlockData, BuildParachainBlockData, Client, DefaultTestClientBuilderExt, HeadData,
-	InitBlockBuilder, TestClientBuilder, TestClientBuilderExt, ValidationParams,
+	transfer, BlockData, BlockOrigin, BuildParachainBlockData, Client, ClientBlockImportExt,
+	DefaultTestClientBuilderExt, HeadData, InitBlockBuilder, TestClientBuilder,
+	TestClientBuilderExt, ValidationParams,
 };
 use cumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
 use sp_keyring::AccountKeyring::*;
-use sp_runtime::traits::Header as HeaderT;
+use sp_runtime::traits::{Block as BlockT, Header as HeaderT};
 use std::{env, process::Command};
 
 use crate::validate_block::MemoryOptimizedValidationParams;
@@ -100,7 +101,7 @@ fn build_block_with_witness(
 }
 
 #[test]
-fn validate_block_no_extra_extrinsics() {
+fn validate_block_works() {
 	sp_tracing::try_init_simple();
 
 	let (client, parent_head) = create_test_client();
@@ -290,3 +291,42 @@ fn validation_params_and_memory_optimized_validation_params_encode_and_decode()
 	let decoded = ValidationParams::decode_all(&mut &encoded[..]).unwrap();
 	assert_eq!(decoded, validation_params);
 }
+
+/// Test for ensuring that we are differentiating in the `validation::trie_cache` between different
+/// child tries.
+///
+/// This is achieved by first building a block using `read_and_write_child_tries` that should set
+/// the values in the child tries. In the second step we are building a second block with the same
+/// extrinsic that reads the values from the child tries and it asserts that we read the correct
+/// data from the state.
+#[test]
+fn validate_block_works_with_child_tries() {
+	sp_tracing::try_init_simple();
+
+	let (mut client, parent_head) = create_test_client();
+	let TestBlockData { block, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Charlie, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+	);
+	let block = block.into_block();
+
+	futures::executor::block_on(client.import(BlockOrigin::Own, block.clone())).unwrap();
+
+	let parent_head = block.header().clone();
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Alice, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+	);
+
+	let header = block.header().clone();
+
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}

cumulus/pallets/parachain-system/src/validate_block/trie_cache.rs

@@ -67,7 +67,13 @@ impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
 /// Provider of [`TrieCache`] instances.
 pub(crate) struct CacheProvider<H: Hasher> {
 	node_cache: RefCell<BTreeMap<H::Out, NodeOwned<H::Out>>>,
-	value_cache: RefCell<BTreeMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>,
+	/// Cache: `storage_root` => `storage_key` => `value`.
+	///
+	/// One `block` can for example use multiple tries (child tries) and we need to distinguish the
+	/// cached (`storage_key`, `value`) between them. For this we are using the `storage_root` to
+	/// distinguish them (even if the storage root is the same for two child tries, it just means
+	/// that both are exactly the same trie and there would happen no collision).
+	value_cache: RefCell<BTreeMap<H::Out, BTreeMap<Box<[u8]>, trie_db::CachedValue<H::Out>>>>,
 }
 
 impl<H: Hasher> CacheProvider<H> {
@@ -81,22 +87,26 @@ impl<H: Hasher> CacheProvider<H> {
 impl<H: Hasher> TrieCacheProvider<H> for CacheProvider<H> {
 	type Cache<'a> = TrieCache<'a, H> where H: 'a;
 
-	fn as_trie_db_cache(&self, _storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
 		TrieCache {
-			value_cache: Some(self.value_cache.borrow_mut()),
+			value_cache: Some(RefMut::map(self.value_cache.borrow_mut(), |c| {
+				c.entry(storage_root).or_default()
+			})),
 			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.
+		// We are not interested in caching new values (as we would throw them away directly after a
+		// block is validated) and thus, we don't pass any `value_cache`.
 		TrieCache { value_cache: None, node_cache: self.node_cache.borrow_mut() }
 	}
 
-	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {}
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {
+		// This is called to merge the `value_cache` from `as_trie_db_mut_cache`, which is not
+		// activated, so we don't need to do anything here.
+	}
 }
 
 // This is safe here since we are single-threaded in WASM