diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index 20bb2c592e925..da4250b6ba3e1 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -1117,4 +1117,95 @@ pub mod tests { ); } } + + /// Test to ensure that recording the same `key` for different tries works as expected. + /// + /// Each trie stores a different value under the same key. The values are big enough to + /// be not inlined with `StateVersion::V1`, this is important to test the expected behavior. The + /// trie recorder is expected to differentiate key access based on the different storage roots + /// of the tries. + #[test] + fn recording_same_key_access_in_different_tries() { + recording_same_key_access_in_different_tries_inner(StateVersion::V0); + recording_same_key_access_in_different_tries_inner(StateVersion::V1); + } + fn recording_same_key_access_in_different_tries_inner(state_version: StateVersion) { + let key = b"test_key".to_vec(); + // Use some big values to ensure that we don't keep them inline + let top_trie_val = vec![1; 1024]; + let child_trie_1_val = vec![2; 1024]; + let child_trie_2_val = vec![3; 1024]; + + let child_info_1 = ChildInfo::new_default(b"sub1"); + let child_info_2 = ChildInfo::new_default(b"sub2"); + let child_info_1 = &child_info_1; + let child_info_2 = &child_info_2; + let contents = vec![ + (None, vec![(key.clone(), Some(top_trie_val.clone()))]), + (Some(child_info_1.clone()), vec![(key.clone(), Some(child_trie_1_val.clone()))]), + (Some(child_info_2.clone()), vec![(key.clone(), Some(child_trie_2_val.clone()))]), + ]; + let in_memory = new_in_mem::>(); + let in_memory = in_memory.update(contents, state_version); + let child_storage_keys = vec![child_info_1.to_owned(), child_info_2.to_owned()]; + let in_memory_root = in_memory + .full_storage_root( + std::iter::empty(), + child_storage_keys.iter().map(|k| (k, std::iter::empty())), + state_version, + ) + .0; + assert_eq!(in_memory.storage(&key).unwrap().unwrap(), top_trie_val); + assert_eq!(in_memory.child_storage(child_info_1, &key).unwrap().unwrap(), child_trie_1_val); + assert_eq!(in_memory.child_storage(child_info_2, &key).unwrap().unwrap(), child_trie_2_val); + + for cache in [Some(SharedTrieCache::new(CacheSize::Unlimited)), None] { + // Run multiple times to have a different cache conditions. + for i in 0..5 { + eprintln!("Running with cache {}, iteration {}", cache.is_some(), i); + + if let Some(cache) = &cache { + if i == 2 { + cache.reset_node_cache(); + } else if i == 3 { + cache.reset_value_cache(); + } + } + + let trie = in_memory.as_trie_backend(); + let trie_root = trie.storage_root(std::iter::empty(), state_version).0; + assert_eq!(in_memory_root, trie_root); + + let proving = TrieBackendBuilder::wrap(&trie) + .with_recorder(Recorder::default()) + .with_optional_cache(cache.as_ref().map(|c| c.local_cache())) + .build(); + assert_eq!(proving.storage(&key).unwrap().unwrap(), top_trie_val); + assert_eq!( + proving.child_storage(child_info_1, &key).unwrap().unwrap(), + child_trie_1_val + ); + assert_eq!( + proving.child_storage(child_info_2, &key).unwrap().unwrap(), + child_trie_2_val + ); + + let proof = proving.extract_proof().unwrap(); + + let proof_check = + create_proof_check_backend::(in_memory_root.into(), proof) + .unwrap(); + + assert_eq!(proof_check.storage(&key).unwrap().unwrap(), top_trie_val); + assert_eq!( + proof_check.child_storage(child_info_1, &key).unwrap().unwrap(), + child_trie_1_val + ); + assert_eq!( + proof_check.child_storage(child_info_2, &key).unwrap().unwrap(), + child_trie_2_val + ); + } + } + } } diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index cd2a71163e2ee..cdd1bb0bba055 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -177,7 +177,7 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss ) -> R, ) -> R { let storage_root = storage_root.unwrap_or_else(|| self.root); - let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder()); + let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root)); let recorder = match recorder.as_mut() { Some(recorder) => Some(recorder as &mut dyn TrieRecorder), None => None, @@ -209,16 +209,19 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss /// This function must only be used when the operation in `callback` is /// calculating a `storage_root`. It is expected that `callback` returns /// the new storage root. This is required to register the changes in the cache - /// for the correct storage root. + /// for the correct storage root. The given `storage_root` corresponds to the root of the "old" + /// trie. If the value is not given, `self.root` is used. #[cfg(feature = "std")] fn with_recorder_and_cache_for_storage_root( &self, + storage_root: Option, callback: impl FnOnce( Option<&mut dyn TrieRecorder>, Option<&mut dyn TrieCache>>, ) -> (Option, R), ) -> R { - let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder()); + let storage_root = storage_root.unwrap_or_else(|| self.root); + let mut recorder = self.recorder.as_ref().map(|r| r.as_trie_recorder(storage_root)); let recorder = match recorder.as_mut() { Some(recorder) => Some(recorder as &mut dyn TrieRecorder), None => None, @@ -244,6 +247,7 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss #[cfg(not(feature = "std"))] fn with_recorder_and_cache_for_storage_root( &self, + _: Option, callback: impl FnOnce( Option<&mut dyn TrieRecorder>, Option<&mut dyn TrieCache>>, @@ -675,7 +679,7 @@ where ) -> (H::Out, S::Overlay) { let mut write_overlay = S::Overlay::default(); - let root = self.with_recorder_and_cache_for_storage_root(|recorder, cache| { + let root = self.with_recorder_and_cache_for_storage_root(None, |recorder, cache| { let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay); let res = match state_version { StateVersion::V0 => delta_trie_root::, _, _, _, _, _>( @@ -719,35 +723,36 @@ where }, }; - let new_child_root = self.with_recorder_and_cache_for_storage_root(|recorder, cache| { - let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay); - match match state_version { - StateVersion::V0 => - child_delta_trie_root::, _, _, _, _, _, _>( - child_info.keyspace(), - &mut eph, - child_root, - delta, - recorder, - cache, - ), - StateVersion::V1 => - child_delta_trie_root::, _, _, _, _, _, _>( - child_info.keyspace(), - &mut eph, - child_root, - delta, - recorder, - cache, - ), - } { - Ok(ret) => (Some(ret), ret), - Err(e) => { - warn!(target: "trie", "Failed to write to trie: {}", e); - (None, child_root) - }, - } - }); + let new_child_root = + self.with_recorder_and_cache_for_storage_root(Some(child_root), |recorder, cache| { + let mut eph = Ephemeral::new(self.backend_storage(), &mut write_overlay); + match match state_version { + StateVersion::V0 => + child_delta_trie_root::, _, _, _, _, _, _>( + child_info.keyspace(), + &mut eph, + child_root, + delta, + recorder, + cache, + ), + StateVersion::V1 => + child_delta_trie_root::, _, _, _, _, _, _>( + child_info.keyspace(), + &mut eph, + child_root, + delta, + recorder, + cache, + ), + } { + Ok(ret) => (Some(ret), ret), + Err(e) => { + warn!(target: "trie", "Failed to write to trie: {}", e); + (None, child_root) + }, + } + }); let is_default = new_child_root == default_root; diff --git a/primitives/trie/src/cache/mod.rs b/primitives/trie/src/cache/mod.rs index f6a447763a80e..85539cf626857 100644 --- a/primitives/trie/src/cache/mod.rs +++ b/primitives/trie/src/cache/mod.rs @@ -490,7 +490,7 @@ mod tests { { let mut cache = local_cache.as_trie_db_cache(root); - let mut recorder = recorder.as_trie_recorder(); + let mut recorder = recorder.as_trie_recorder(root); let trie = TrieDBBuilder::::new(&db, &root) .with_cache(&mut cache) .with_recorder(&mut recorder) @@ -538,7 +538,7 @@ mod tests { { let mut db = db.clone(); let mut cache = local_cache.as_trie_db_cache(root); - let mut recorder = recorder.as_trie_recorder(); + let mut recorder = recorder.as_trie_recorder(root); let mut trie = TrieDBMutBuilder::::from_existing(&mut db, &mut new_root) .with_cache(&mut cache) .with_recorder(&mut recorder) diff --git a/primitives/trie/src/recorder.rs b/primitives/trie/src/recorder.rs index 6fbf698592a31..bc67cfc287942 100644 --- a/primitives/trie/src/recorder.rs +++ b/primitives/trie/src/recorder.rs @@ -41,7 +41,7 @@ const LOG_TARGET: &str = "trie-recorder"; /// The internals of [`Recorder`]. struct RecorderInner { /// The keys for that we have recorded the trie nodes and if we have recorded up to the value. - recorded_keys: HashMap, RecordedForKey>, + recorded_keys: HashMap, RecordedForKey>>, /// The encoded nodes we accessed while recording. accessed_nodes: HashMap>, } @@ -80,9 +80,16 @@ impl Clone for Recorder { impl Recorder { /// Returns the recorder as [`TrieRecorder`](trie_db::TrieRecorder) compatible type. - pub fn as_trie_recorder(&self) -> impl trie_db::TrieRecorder + '_ { + /// + /// - `storage_root`: The storage root of the trie for which accesses are recorded. This is + /// important when recording access to different tries at once (like top and child tries). + pub fn as_trie_recorder( + &self, + storage_root: H::Out, + ) -> impl trie_db::TrieRecorder + '_ { TrieRecorder:: { inner: self.inner.lock(), + storage_root, encoded_size_estimation: self.encoded_size_estimation.clone(), _phantom: PhantomData, } @@ -132,6 +139,7 @@ impl Recorder { /// The [`TrieRecorder`](trie_db::TrieRecorder) implementation. struct TrieRecorder { inner: I, + storage_root: H::Out, encoded_size_estimation: Arc, _phantom: PhantomData, } @@ -191,6 +199,8 @@ impl>> trie_db::TrieRecord self.inner .recorded_keys + .entry(self.storage_root) + .or_default() .entry(full_key.to_vec()) .and_modify(|e| *e = RecordedForKey::Value) .or_insert(RecordedForKey::Value); @@ -206,6 +216,8 @@ impl>> trie_db::TrieRecord // accounted for by the recorded node that holds the hash. self.inner .recorded_keys + .entry(self.storage_root) + .or_default() .entry(full_key.to_vec()) .or_insert(RecordedForKey::Hash); }, @@ -221,6 +233,8 @@ impl>> trie_db::TrieRecord // that the value doesn't exist in the trie. self.inner .recorded_keys + .entry(self.storage_root) + .or_default() .entry(full_key.to_vec()) .and_modify(|e| *e = RecordedForKey::Value) .or_insert(RecordedForKey::Value); @@ -231,7 +245,11 @@ impl>> trie_db::TrieRecord } fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey { - self.inner.recorded_keys.get(key).copied().unwrap_or(RecordedForKey::None) + self.inner + .recorded_keys + .get(&self.storage_root) + .and_then(|k| k.get(key).copied()) + .unwrap_or(RecordedForKey::None) } } @@ -267,7 +285,7 @@ mod tests { let recorder = Recorder::default(); { - let mut trie_recorder = recorder.as_trie_recorder(); + let mut trie_recorder = recorder.as_trie_recorder(root); let trie = TrieDBBuilder::::new(&db, &root) .with_recorder(&mut trie_recorder) .build();