diff --git a/primitives/io/src/lib.rs b/primitives/io/src/lib.rs index 8f62b03b62ea9..90c43ac517b2a 100644 --- a/primitives/io/src/lib.rs +++ b/primitives/io/src/lib.rs @@ -1766,6 +1766,11 @@ mod tests { assert!(storage::get(b":abdd").is_some()); assert!(storage::get(b":abcd").is_none()); assert!(storage::get(b":abc").is_none()); + + assert!(matches!( + storage::clear_prefix(b":abc", None), + KillStorageResult::AllRemoved(0) + )); }); } diff --git a/primitives/state-machine/src/ext.rs b/primitives/state-machine/src/ext.rs index e33569e2a1f67..6a9635bd26638 100644 --- a/primitives/state-machine/src/ext.rs +++ b/primitives/state-machine/src/ext.rs @@ -464,8 +464,10 @@ where } self.mark_dirty(); - self.overlay.clear_prefix(prefix); - self.limit_remove_from_backend(None, Some(prefix), limit) + // In the future, we should report this back to the caller. + let _overlay_count = self.overlay.clear_prefix(prefix); + let (all, count) = self.limit_remove_from_backend(None, Some(prefix), limit); + (all, count) } fn clear_child_prefix( @@ -484,8 +486,10 @@ where let _guard = guard(); self.mark_dirty(); - self.overlay.clear_child_prefix(child_info, prefix); - self.limit_remove_from_backend(Some(child_info), Some(prefix), limit) + // In the future, we should report this back to the caller. + let _overlay_count = self.overlay.clear_child_prefix(child_info, prefix); + let (all, count) = self.limit_remove_from_backend(Some(child_info), Some(prefix), limit); + (all, count) } fn storage_append(&mut self, key: Vec, value: Vec) { @@ -741,6 +745,13 @@ where all_deleted = false; return false } + if let Some(None) = match child_info { + Some(child_info) => self.overlay.child_storage(child_info, key), + None => self.overlay.storage(key), + } { + // already deleted. + return true + } if let Some(num) = num_deleted.checked_add(1) { num_deleted = num; } else { @@ -757,6 +768,13 @@ where (all_deleted, num_deleted) } else { self.backend.apply_to_keys_while(child_info, prefix, |key| { + if let Some(None) = match child_info { + Some(child_info) => self.overlay.child_storage(child_info, key), + None => self.overlay.storage(key), + } { + // already deleted. + return true + } num_deleted = num_deleted.saturating_add(1); if let Some(child_info) = child_info { self.overlay.set_child_storage(child_info, key.to_vec(), None); diff --git a/primitives/state-machine/src/lib.rs b/primitives/state-machine/src/lib.rs index 97a9ae8d88cd2..89764888146b3 100644 --- a/primitives/state-machine/src/lib.rs +++ b/primitives/state-machine/src/lib.rs @@ -1675,11 +1675,25 @@ mod tests { let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None); assert_eq!(ext.kill_child_storage(&child_info, Some(0)), (false, 0)); assert_eq!(ext.kill_child_storage(&child_info, Some(1)), (false, 1)); - assert_eq!(ext.kill_child_storage(&child_info, Some(2)), (false, 2)); - assert_eq!(ext.kill_child_storage(&child_info, Some(3)), (false, 3)); - assert_eq!(ext.kill_child_storage(&child_info, Some(4)), (true, 4)); - // Only 4 items to remove - assert_eq!(ext.kill_child_storage(&child_info, Some(5)), (true, 4)); + // Only 3 items remaining to remove + assert_eq!(ext.kill_child_storage(&child_info, Some(4)), (true, 3)); + } + + #[test] + fn limited_child_kill_off_by_one_works_without_limit() { + let child_info = ChildInfo::new_default(b"sub1"); + let initial: HashMap<_, BTreeMap<_, _>> = map![ + Some(child_info.clone()) => map![ + b"a".to_vec() => b"0".to_vec(), + b"b".to_vec() => b"1".to_vec(), + b"c".to_vec() => b"2".to_vec(), + b"d".to_vec() => b"3".to_vec() + ], + ]; + let backend = InMemoryBackend::::from((initial, StateVersion::default())); + let mut overlay = OverlayedChanges::default(); + let mut cache = StorageTransactionCache::default(); + let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None); assert_eq!(ext.kill_child_storage(&child_info, None), (true, 4)); } diff --git a/primitives/state-machine/src/overlayed_changes/changeset.rs b/primitives/state-machine/src/overlayed_changes/changeset.rs index 9e7f6ffeddfd7..ae5d47f6bb943 100644 --- a/primitives/state-machine/src/overlayed_changes/changeset.rs +++ b/primitives/state-machine/src/overlayed_changes/changeset.rs @@ -410,10 +410,15 @@ impl OverlayedChangeSet { &mut self, predicate: impl Fn(&[u8], &OverlayedValue) -> bool, at_extrinsic: Option, - ) { + ) -> u32 { + let mut count = 0; for (key, val) in self.changes.iter_mut().filter(|(k, v)| predicate(k, v)) { + if val.value_ref().is_some() { + count += 1; + } val.set(None, insert_dirty(&mut self.dirty_keys, key.clone()), at_extrinsic); } + count } /// Get the iterator over all changes that follow the supplied `key`. diff --git a/primitives/state-machine/src/overlayed_changes/mod.rs b/primitives/state-machine/src/overlayed_changes/mod.rs index 2161b343711c9..fdaa5ce787861 100644 --- a/primitives/state-machine/src/overlayed_changes/mod.rs +++ b/primitives/state-machine/src/overlayed_changes/mod.rs @@ -315,14 +315,14 @@ impl OverlayedChanges { /// Removes all key-value pairs which keys share the given prefix. /// /// Can be rolled back or committed when called inside a transaction. - pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) { - self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index()); + pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) -> u32 { + self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index()) } /// Removes all key-value pairs which keys share the given prefix. /// /// Can be rolled back or committed when called inside a transaction - pub(crate) fn clear_child_prefix(&mut self, child_info: &ChildInfo, prefix: &[u8]) { + pub(crate) fn clear_child_prefix(&mut self, child_info: &ChildInfo, prefix: &[u8]) -> u32 { let extrinsic_index = self.extrinsic_index(); let storage_key = child_info.storage_key().to_vec(); let top = &self.top; @@ -332,7 +332,7 @@ impl OverlayedChanges { .or_insert_with(|| (top.spawn_child(), child_info.clone())); let updatable = info.try_update(child_info); debug_assert!(updatable); - changeset.clear_where(|key, _| key.starts_with(prefix), extrinsic_index); + changeset.clear_where(|key, _| key.starts_with(prefix), extrinsic_index) } /// Returns the current nesting depth of the transaction stack.