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>

substrate/primitives/io/src/lib.rs

@@ -81,6 +81,8 @@ mod batch_verifier;
 #[cfg(feature = "std")]
 use batch_verifier::BatchVerifier;
 
+pub use sp_externalities::MultiRemovalResults;
+
 #[cfg(feature = "std")]
 const LOG_TARGET: &str = "runtime::io";
 
@@ -99,12 +101,24 @@ pub enum EcdsaVerifyError {
 /// removed from the backend from making the `storage_kill` call.
 #[derive(PassByCodec, Encode, Decode)]
 pub enum KillStorageResult {
-	/// All key to remove were removed, return number of key removed from backend.
+	/// All keys to remove were removed, return number of iterations performed during the
+	/// operation.
 	AllRemoved(u32),
-	/// Not all key to remove were removed, return number of key removed from backend.
+	/// Not all key to remove were removed, return number of iterations performed during the
+	/// operation.
 	SomeRemaining(u32),
 }
 
+impl From<MultiRemovalResults> for KillStorageResult {
+	fn from(r: MultiRemovalResults) -> Self {
+		match r {
+			MultiRemovalResults { maybe_cursor: None, backend, .. } => Self::AllRemoved(backend),
+			MultiRemovalResults { maybe_cursor: Some(..), backend, .. } =>
+				Self::SomeRemaining(backend),
+		}
+	}
+}
+
 /// Interface for accessing the storage from within the runtime.
 #[runtime_interface]
 pub trait Storage {
@@ -145,7 +159,7 @@ pub trait Storage {
 
 	/// Clear the storage of each key-value pair where the key starts with the given `prefix`.
 	fn clear_prefix(&mut self, prefix: &[u8]) {
-		let _ = Externalities::clear_prefix(*self, prefix, None);
+		let _ = Externalities::clear_prefix(*self, prefix, None, None);
 	}
 
 	/// Clear the storage of each key-value pair where the key starts with the given `prefix`.
@@ -175,13 +189,60 @@ pub trait Storage {
 	/// backend.
 	#[version(2)]
 	fn clear_prefix(&mut self, prefix: &[u8], limit: Option<u32>) -> KillStorageResult {
-		let (all_removed, num_removed) = Externalities::clear_prefix(*self, prefix, limit);
-		match all_removed {
-			true => KillStorageResult::AllRemoved(num_removed),
-			false => KillStorageResult::SomeRemaining(num_removed),
+		let r = Externalities::clear_prefix(*self, prefix, limit, None);
+		match r.maybe_cursor {
+			None => KillStorageResult::AllRemoved(r.loops),
+			Some(_) => KillStorageResult::SomeRemaining(r.loops),
 		}
 	}
 
+	/// Partially clear the storage of each key-value pair where the key starts with the given
+	/// prefix.
+	///
+	/// # Limit
+	///
+	/// A *limit* should always be provided through `maybe_limit`. This is one fewer than the
+	/// maximum number of backend iterations which may be done by this operation and as such
+	/// represents the maximum number of backend deletions which may happen. A *limit* of zero
+	/// implies that no keys will be deleted, though there may be a single iteration done.
+	///
+	/// The limit can be used to partially delete a prefix storage in case it is too large or costly
+	/// to delete in a single operation.
+	///
+	/// # Cursor
+	///
+	/// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be
+	/// passed once (in the initial call) for any given `maybe_prefix` value. Subsequent calls
+	/// operating on the same prefix should always pass `Some`, and this should be equal to the
+	/// previous call result's `maybe_cursor` field.
+	///
+	/// Returns [`MultiRemovalResults`](sp_io::MultiRemovalResults) to inform about the result. Once
+	/// the resultant `maybe_cursor` field is `None`, then no further items remain to be deleted.
+	///
+	/// NOTE: After the initial call for any given prefix, it is important that no keys further
+	/// keys under the same prefix are inserted. If so, then they may or may not be deleted by
+	/// subsequent calls.
+	///
+	/// # Note
+	///
+	/// Please note that keys which are residing in the overlay for that prefix when
+	/// issuing this call are deleted without counting towards the `limit`.
+	#[version(3, register_only)]
+	fn clear_prefix(
+		&mut self,
+		maybe_prefix: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<Vec<u8>>, //< TODO Make work or just Option<Vec<u8>>?
+	) -> MultiRemovalResults {
+		Externalities::clear_prefix(
+			*self,
+			maybe_prefix,
+			maybe_limit,
+			maybe_cursor.as_ref().map(|x| &x[..]),
+		)
+		.into()
+	}
+
 	/// Append the encoded `value` to the storage item at `key`.
 	///
 	/// The storage item needs to implement [`EncodeAppend`](codec::EncodeAppend).
@@ -323,7 +384,7 @@ pub trait DefaultChildStorage {
 	/// is removed.
 	fn storage_kill(&mut self, storage_key: &[u8]) {
 		let child_info = ChildInfo::new_default(storage_key);
-		self.kill_child_storage(&child_info, None);
+		let _ = self.kill_child_storage(&child_info, None, None);
 	}
 
 	/// Clear a child storage key.
@@ -332,8 +393,8 @@ pub trait DefaultChildStorage {
 	#[version(2)]
 	fn storage_kill(&mut self, storage_key: &[u8], limit: Option<u32>) -> bool {
 		let child_info = ChildInfo::new_default(storage_key);
-		let (all_removed, _num_removed) = self.kill_child_storage(&child_info, limit);
-		all_removed
+		let r = self.kill_child_storage(&child_info, limit, None);
+		r.maybe_cursor.is_none()
 	}
 
 	/// Clear a child storage key.
@@ -342,13 +403,28 @@ pub trait DefaultChildStorage {
 	#[version(3)]
 	fn storage_kill(&mut self, storage_key: &[u8], limit: Option<u32>) -> KillStorageResult {
 		let child_info = ChildInfo::new_default(storage_key);
-		let (all_removed, num_removed) = self.kill_child_storage(&child_info, limit);
-		match all_removed {
-			true => KillStorageResult::AllRemoved(num_removed),
-			false => KillStorageResult::SomeRemaining(num_removed),
+		let r = self.kill_child_storage(&child_info, limit, None);
+		match r.maybe_cursor {
+			None => KillStorageResult::AllRemoved(r.loops),
+			Some(..) => KillStorageResult::SomeRemaining(r.loops),
 		}
 	}
 
+	/// Clear a child storage key.
+	///
+	/// See `Storage` module `clear_prefix` documentation for `limit` usage.
+	#[version(4, register_only)]
+	fn storage_kill(
+		&mut self,
+		storage_key: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<Vec<u8>>,
+	) -> MultiRemovalResults {
+		let child_info = ChildInfo::new_default(storage_key);
+		self.kill_child_storage(&child_info, maybe_limit, maybe_cursor.as_ref().map(|x| &x[..]))
+			.into()
+	}
+
 	/// Check a child storage key.
 	///
 	/// Check whether the given `key` exists in default child defined at `storage_key`.
@@ -362,7 +438,7 @@ pub trait DefaultChildStorage {
 	/// Clear the child storage of each key-value pair where the key starts with the given `prefix`.
 	fn clear_prefix(&mut self, storage_key: &[u8], prefix: &[u8]) {
 		let child_info = ChildInfo::new_default(storage_key);
-		let _ = self.clear_child_prefix(&child_info, prefix, None);
+		let _ = self.clear_child_prefix(&child_info, prefix, None, None);
 	}
 
 	/// Clear the child storage of each key-value pair where the key starts with the given `prefix`.
@@ -376,13 +452,34 @@ pub trait DefaultChildStorage {
 		limit: Option<u32>,
 	) -> KillStorageResult {
 		let child_info = ChildInfo::new_default(storage_key);
-		let (all_removed, num_removed) = self.clear_child_prefix(&child_info, prefix, limit);
-		match all_removed {
-			true => KillStorageResult::AllRemoved(num_removed),
-			false => KillStorageResult::SomeRemaining(num_removed),
+		let r = self.clear_child_prefix(&child_info, prefix, limit, None);
+		match r.maybe_cursor {
+			None => KillStorageResult::AllRemoved(r.loops),
+			Some(..) => KillStorageResult::SomeRemaining(r.loops),
 		}
 	}
 
+	/// Clear the child storage of each key-value pair where the key starts with the given `prefix`.
+	///
+	/// See `Storage` module `clear_prefix` documentation for `limit` usage.
+	#[version(3, register_only)]
+	fn clear_prefix(
+		&mut self,
+		storage_key: &[u8],
+		prefix: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<Vec<u8>>,
+	) -> MultiRemovalResults {
+		let child_info = ChildInfo::new_default(storage_key);
+		self.clear_child_prefix(
+			&child_info,
+			prefix,
+			maybe_limit,
+			maybe_cursor.as_ref().map(|x| &x[..]),
+		)
+		.into()
+	}
+
 	/// Default child root calculation.
 	///
 	/// "Commit" all existing operations and compute the resulting child storage root.
@@ -1757,15 +1854,30 @@ mod tests {
 		});
 
 		t.execute_with(|| {
+			// We can switch to this once we enable v3 of the `clear_prefix`.
+			//assert!(matches!(
+			//	storage::clear_prefix(b":abc", None),
+			//	MultiRemovalResults::NoneLeft { db: 2, total: 2 }
+			//));
 			assert!(matches!(
 				storage::clear_prefix(b":abc", None),
-				KillStorageResult::AllRemoved(2)
+				KillStorageResult::AllRemoved(2),
 			));
 
 			assert!(storage::get(b":a").is_some());
 			assert!(storage::get(b":abdd").is_some());
 			assert!(storage::get(b":abcd").is_none());
 			assert!(storage::get(b":abc").is_none());
+
+			// We can switch to this once we enable v3 of the `clear_prefix`.
+			//assert!(matches!(
+			//	storage::clear_prefix(b":abc", None),
+			//	MultiRemovalResults::NoneLeft { db: 0, total: 0 }
+			//));
+			assert!(matches!(
+				storage::clear_prefix(b":abc", None),
+				KillStorageResult::AllRemoved(0),
+			));
 		});
 	}