Safe and sane multi-item storage removal (#11490)

* Fix overlay prefix removal result * Second part of the overlay prefix removal fix. * Report only items deleted from storage in clear_prefix * Fix kill_prefix * Formatting * Remove unused code * Fixes * Fixes * Introduce clear_prefix host function v3 * Formatting * Use v2 for now * Fixes * Formatting * Docs * Child prefix removal should also hide v3 for now * Fixes * Fixes * Formatting * Fixes * apply_to_keys_whle takes start_at * apply_to_keys_whle takes start_at * apply_to_keys_whle takes start_at * Cursor API; force limits * Use unsafe deprecated functions * Formatting * Fixes * Grumbles * Fixes * Docs * Some nitpicks 🙈 * Update primitives/externalities/src/lib.rs Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> * Formatting * Fixes * cargo fmt * Fixes * Update primitives/io/src/lib.rs Co-authored-by: Keith Yeung <kungfukeith11@gmail.com> * Formatting * Fixes Co-authored-by: Bastian Köcher <info@kchr.de> Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com> Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>
2026-04-27 02:17:58 +00:00 · 2022-05-29 12:56:26 +01:00
parent 189a310e4c
commit ecbd65fb95
45 changed files with 968 additions and 206 deletions
@@ -27,7 +27,7 @@ use sp_core::hexdisplay::HexDisplay;
 use sp_core::storage::{
 	well_known_keys::is_child_storage_key, ChildInfo, StateVersion, TrackedStorageKey,
 };
-use sp_externalities::{Extension, ExtensionStore, Externalities};
+use sp_externalities::{Extension, ExtensionStore, Externalities, MultiRemovalResults};
 use sp_trie::{empty_child_trie_root, LayoutV1};

 use crate::{log_error, trace, warn, StorageTransactionCache};
@@ -433,7 +433,12 @@ where
 		self.overlay.set_child_storage(child_info, key, value);
 	}

-	fn kill_child_storage(&mut self, child_info: &ChildInfo, limit: Option<u32>) -> (bool, u32) {
+	fn kill_child_storage(
+		&mut self,
+		child_info: &ChildInfo,
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ChildKill",
@@ -442,11 +447,18 @@ where
 		);
 		let _guard = guard();
 		self.mark_dirty();
-		self.overlay.clear_child_storage(child_info);
-		self.limit_remove_from_backend(Some(child_info), None, limit)
+		let overlay = self.overlay.clear_child_storage(child_info);
+		let (maybe_cursor, backend, loops) =
+			self.limit_remove_from_backend(Some(child_info), None, maybe_limit, maybe_cursor);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}

-	fn clear_prefix(&mut self, prefix: &[u8], limit: Option<u32>) -> (bool, u32) {
+	fn clear_prefix(
+		&mut self,
+		prefix: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ClearPrefix",
@@ -460,20 +472,23 @@ where
 				target: "trie",
 				"Refuse to directly clear prefix that is part or contains of child storage key",
 			);
-			return (false, 0)
+			return MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
 		}

 		self.mark_dirty();
-		self.overlay.clear_prefix(prefix);
-		self.limit_remove_from_backend(None, Some(prefix), limit)
+		let overlay = self.overlay.clear_prefix(prefix);
+		let (maybe_cursor, backend, loops) =
+			self.limit_remove_from_backend(None, Some(prefix), maybe_limit, maybe_cursor);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}

 	fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
-		limit: Option<u32>,
-	) -> (bool, u32) {
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ChildClearPrefix",
@@ -484,8 +499,14 @@ 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)
+		let overlay = self.overlay.clear_child_prefix(child_info, prefix);
+		let (maybe_cursor, backend, loops) = self.limit_remove_from_backend(
+			Some(child_info),
+			Some(prefix),
+			maybe_limit,
+			maybe_cursor,
+		);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}

 	fn storage_append(&mut self, key: Vec<u8>, value: Vec<u8>) {
@@ -728,45 +749,37 @@ where
 {
 	fn limit_remove_from_backend(
 		&mut self,
-		child_info: Option<&ChildInfo>,
-		prefix: Option<&[u8]>,
-		limit: Option<u32>,
-	) -> (bool, u32) {
-		let mut num_deleted: u32 = 0;
-
-		if let Some(limit) = limit {
-			let mut all_deleted = true;
-			self.backend.apply_to_keys_while(child_info, prefix, |key| {
-				if num_deleted == limit {
-					all_deleted = false;
+		maybe_child: Option<&ChildInfo>,
+		maybe_prefix: Option<&[u8]>,
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> (Option<Vec<u8>>, u32, u32) {
+		let mut delete_count: u32 = 0;
+		let mut loop_count: u32 = 0;
+		let mut maybe_next_key = None;
+		self.backend
+			.apply_to_keys_while(maybe_child, maybe_prefix, maybe_cursor, |key| {
+				if maybe_limit.map_or(false, |limit| loop_count == limit) {
+					maybe_next_key = Some(key.to_vec());
 					return false
 				}
-				if let Some(num) = num_deleted.checked_add(1) {
-					num_deleted = num;
-				} else {
-					all_deleted = false;
-					return false
-				}
-				if let Some(child_info) = child_info {
-					self.overlay.set_child_storage(child_info, key.to_vec(), None);
-				} else {
-					self.overlay.set_storage(key.to_vec(), None);
+				let overlay = match maybe_child {
+					Some(child_info) => self.overlay.child_storage(child_info, key),
+					None => self.overlay.storage(key),
+				};
+				if !matches!(overlay, Some(None)) {
+					// not pending deletion from the backend - delete it.
+					if let Some(child_info) = maybe_child {
+						self.overlay.set_child_storage(child_info, key.to_vec(), None);
+					} else {
+						self.overlay.set_storage(key.to_vec(), None);
+					}
+					delete_count = delete_count.saturating_add(1);
 				}
+				loop_count = loop_count.saturating_add(1);
 				true
 			});
-			(all_deleted, num_deleted)
-		} else {
-			self.backend.apply_to_keys_while(child_info, prefix, |key| {
-				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);
-				} else {
-					self.overlay.set_storage(key.to_vec(), None);
-				}
-				true
-			});
-			(true, num_deleted)
-		}
+		(maybe_next_key, delete_count, loop_count)
 	}
 }

@@ -1085,14 +1098,14 @@ mod tests {
 		not_under_prefix.extend(b"path");
 		ext.set_storage(not_under_prefix.clone(), vec![10]);

-		ext.clear_prefix(&[], None);
-		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None);
+		let _ = ext.clear_prefix(&[], None, None);
+		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
 		let mut under_prefix = well_known_keys::CHILD_STORAGE_KEY_PREFIX.to_vec();
 		under_prefix.extend(b"path");
-		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None);
+		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
 		assert_eq!(ext.child_storage(child_info, &[30]), Some(vec![40]));
 		assert_eq!(ext.storage(not_under_prefix.as_slice()), Some(vec![10]));
-		ext.clear_prefix(&not_under_prefix[..5], None);
+		let _ = ext.clear_prefix(&not_under_prefix[..5], None, None);
 		assert_eq!(ext.storage(not_under_prefix.as_slice()), None);
 	}