From 8fcc6f2e9ff4da0193c309c8b16e59cbd7a96c0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastian=20K=C3=B6cher?= Date: Tue, 8 Nov 2022 10:04:55 +0100 Subject: [PATCH] `sp_trie::Recorder`: Fix recording the same key for different tries (#12636) With `StateVersion::V1` values over a certain size are not inlined and being put into the backend with their own hash. When accessing a value in the trie with a recorder, we check if the value is maybe already recorded and thus, we can check the cache. To check if a value is already recorded, we use the key of the value to differentiate them. The problem is when there are multiple tries, like multiple child tries that all have different values under the same key. Before this pull request we didn't have differentiated for which trie we already had recorded a (key, value) pair. This is now done by also taking the storage root into account in the recorder to differentiate the different (key, value) pair in the tries. --- .../state-machine/src/trie_backend.rs | 91 +++++++++++++++++++ .../state-machine/src/trie_backend_essence.rs | 71 ++++++++------- substrate/primitives/trie/src/cache/mod.rs | 4 +- substrate/primitives/trie/src/recorder.rs | 26 +++++- 4 files changed, 153 insertions(+), 39 deletions(-) diff --git a/substrate/primitives/state-machine/src/trie_backend.rs b/substrate/primitives/state-machine/src/trie_backend.rs index 20bb2c592e..da4250b6ba 100644 --- a/substrate/primitives/state-machine/src/trie_backend.rs +++ b/substrate/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/substrate/primitives/state-machine/src/trie_backend_essence.rs b/substrate/primitives/state-machine/src/trie_backend_essence.rs index cd2a71163e..cdd1bb0bba 100644 --- a/substrate/primitives/state-machine/src/trie_backend_essence.rs +++ b/substrate/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/substrate/primitives/trie/src/cache/mod.rs b/substrate/primitives/trie/src/cache/mod.rs index f6a447763a..85539cf626 100644 --- a/substrate/primitives/trie/src/cache/mod.rs +++ b/substrate/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/substrate/primitives/trie/src/recorder.rs b/substrate/primitives/trie/src/recorder.rs index 6fbf698592..bc67cfc287 100644 --- a/substrate/primitives/trie/src/recorder.rs +++ b/substrate/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();