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/externalities/src/lib.rs

@@ -52,6 +52,30 @@ pub enum Error {
 	StorageUpdateFailed(&'static str),
 }
 
+/// Results concerning an operation to remove many keys.
+#[derive(codec::Encode, codec::Decode)]
+#[must_use]
+pub struct MultiRemovalResults {
+	/// A continuation cursor which, if `Some` must be provided to the subsequent removal call.
+	/// If `None` then all removals are complete and no further calls are needed.
+	pub maybe_cursor: Option<Vec<u8>>,
+	/// The number of items removed from the backend database.
+	pub backend: u32,
+	/// The number of unique keys removed, taking into account both the backend and the overlay.
+	pub unique: u32,
+	/// The number of iterations (each requiring a storage seek/read) which were done.
+	pub loops: u32,
+}
+
+impl MultiRemovalResults {
+	/// Deconstruct into the internal components.
+	///
+	/// Returns `(maybe_cursor, backend, unique, loops)`.
+	pub fn deconstruct(self) -> (Option<Vec<u8>>, u32, u32, u32) {
+		(self.maybe_cursor, self.backend, self.unique, self.loops)
+	}
+}
+
 /// The Substrate externalities.
 ///
 /// Provides access to the storage and to other registered extensions.
@@ -118,32 +142,47 @@ pub trait Externalities: ExtensionStore {
 
 	/// Clear an entire child storage.
 	///
-	/// Deletes all keys from the overlay and up to `limit` keys from the backend. No
-	/// limit is applied if `limit` is `None`. Returned boolean is `true` if the child trie was
-	/// removed completely and `false` if there are remaining keys after the function
-	/// returns. Returned `u32` is the number of keys that was removed at the end of the
-	/// operation.
+	/// Deletes all keys from the overlay and up to `maybe_limit` keys from the backend. No
+	/// limit is applied if `maybe_limit` is `None`. Returns the cursor for the next call as `Some`
+	/// if the child trie deletion operation is incomplete. In this case, it should be passed into
+	/// the next call to avoid unaccounted iterations on the backend. Returns also the the number
+	/// of keys that were removed from the backend, the number of unique keys removed in total
+	/// (including from the overlay) and the number of backend iterations done.
+	///
+	/// As long as `maybe_cursor` is passed from the result of the previous call, then the number of
+	/// iterations done will only ever be one more than the number of keys removed.
 	///
 	/// # Note
 	///
 	/// An implementation is free to delete more keys than the specified limit as long as
 	/// it is able to do that in constant time.
-	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;
 
 	/// Clear storage entries which keys are start with the given prefix.
 	///
-	/// `limit` and result works as for `kill_child_storage`.
-	fn clear_prefix(&mut self, prefix: &[u8], limit: Option<u32>) -> (bool, u32);
+	/// `maybe_limit`, `maybe_cursor` and result works as for `kill_child_storage`.
+	fn clear_prefix(
+		&mut self,
+		prefix: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults;
 
 	/// Clear child storage entries which keys are start with the given prefix.
 	///
-	/// `limit` and result works as for `kill_child_storage`.
+	/// `maybe_limit`, `maybe_cursor` and result works as for `kill_child_storage`.
 	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;
 
 	/// Set or clear a storage entry (`key`) of current contract being called (effective
 	/// immediately).

substrate/primitives/io/Cargo.toml

@@ -25,7 +25,7 @@ sp-state-machine = { version = "0.12.0", optional = true, path = "../state-machi
 sp-wasm-interface = { version = "6.0.0", path = "../wasm-interface", default-features = false }
 sp-runtime-interface = { version = "6.0.0", default-features = false, path = "../runtime-interface" }
 sp-trie = { version = "6.0.0", optional = true, path = "../trie" }
-sp-externalities = { version = "0.12.0", optional = true, path = "../externalities" }
+sp-externalities = { version = "0.12.0", default-features = false, path = "../externalities" }
 sp-tracing = { version = "5.0.0", default-features = false, path = "../tracing" }
 log = { version = "0.4.17", optional = true }
 futures = { version = "0.3.21", features = ["thread-pool"], optional = true }
@@ -37,6 +37,7 @@ tracing-core = { version = "0.1.26", default-features = false}
 [features]
 default = ["std"]
 std = [
+	"sp-externalities/std",
 	"sp-core/std",
 	"sp-keystore",
 	"codec/std",
@@ -47,7 +48,6 @@ std = [
 	"libsecp256k1",
 	"secp256k1",
 	"sp-runtime-interface/std",
-	"sp-externalities",
 	"sp-wasm-interface/std",
 	"sp-tracing/std",
 	"tracing/std",

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),
+			));
 		});
 	}
 

substrate/primitives/runtime-interface/Cargo.toml

@@ -18,7 +18,7 @@ sp-wasm-interface = { version = "6.0.0", path = "../wasm-interface", default-fea
 sp-std = { version = "4.0.0", default-features = false, path = "../std" }
 sp-tracing = { version = "5.0.0", default-features = false, path = "../tracing" }
 sp-runtime-interface-proc-macro = { version = "5.0.0", path = "proc-macro" }
-sp-externalities = { version = "0.12.0", optional = true, path = "../externalities" }
+sp-externalities = { version = "0.12.0", default-features = false, path = "../externalities" }
 codec = { package = "parity-scale-codec", version = "3.0.0", default-features = false }
 static_assertions = "1.0.0"
 primitive-types = { version = "0.11.1", default-features = false }
@@ -40,7 +40,7 @@ std = [
 	"sp-std/std",
 	"sp-tracing/std",
 	"codec/std",
-	"sp-externalities",
+	"sp-externalities/std",
 	"primitive-types/std",
 ]
 

substrate/primitives/runtime-interface/src/impls.rs

@@ -548,3 +548,7 @@ impl PassBy for sp_storage::TrackedStorageKey {
 impl PassBy for sp_storage::StateVersion {
 	type PassBy = Enum<Self>;
 }
+
+impl PassBy for sp_externalities::MultiRemovalResults {
+	type PassBy = Codec<Self>;
+}

substrate/primitives/state-machine/Cargo.toml

@@ -35,6 +35,7 @@ hex-literal = "0.3.4"
 pretty_assertions = "1.2.1"
 rand = "0.7.2"
 sp-runtime = { version = "6.0.0", path = "../runtime" }
+assert_matches = "1.5"
 
 [features]
 default = ["std"]

substrate/primitives/state-machine/src/backend.rs

@@ -112,6 +112,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	);
 

substrate/primitives/state-machine/src/basic.rs

@@ -29,7 +29,7 @@ use sp_core::{
 	traits::Externalities,
 	Blake2Hasher,
 };
-use sp_externalities::{Extension, Extensions};
+use sp_externalities::{Extension, Extensions, MultiRemovalResults};
 use sp_trie::{empty_child_trie_root, HashKey, LayoutV0, LayoutV1, TrieConfiguration};
 use std::{
 	any::{Any, TypeId},
@@ -209,23 +209,34 @@ impl Externalities for BasicExternalities {
 		}
 	}
 
-	fn kill_child_storage(&mut self, child_info: &ChildInfo, _limit: Option<u32>) -> (bool, u32) {
-		let num_removed = self
+	fn kill_child_storage(
+		&mut self,
+		child_info: &ChildInfo,
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
+		let count = self
 			.inner
 			.children_default
 			.remove(child_info.storage_key())
 			.map(|c| c.data.len())
-			.unwrap_or(0);
-		(true, num_removed as u32)
+			.unwrap_or(0) as u32;
+		MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 	}
 
-	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 {
 		if is_child_storage_key(prefix) {
 			warn!(
 				target: "trie",
 				"Refuse to clear prefix that is part of child storage key via main storage"
 			);
-			return (false, 0)
+			let maybe_cursor = Some(prefix.to_vec());
+			return MultiRemovalResults { maybe_cursor, backend: 0, unique: 0, loops: 0 }
 		}
 
 		let to_remove = self
@@ -237,19 +248,20 @@ impl Externalities for BasicExternalities {
 			.cloned()
 			.collect::<Vec<_>>();
 
-		let num_removed = to_remove.len();
+		let count = to_remove.len() as u32;
 		for key in to_remove {
 			self.inner.top.remove(&key);
 		}
-		(true, num_removed as u32)
+		MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 	}
 
 	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 {
 		if let Some(child) = self.inner.children_default.get_mut(child_info.storage_key()) {
 			let to_remove = child
 				.data
@@ -259,13 +271,13 @@ impl Externalities for BasicExternalities {
 				.cloned()
 				.collect::<Vec<_>>();
 
-			let num_removed = to_remove.len();
+			let count = to_remove.len() as u32;
 			for key in to_remove {
 				child.data.remove(&key);
 			}
-			(true, num_removed as u32)
+			MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 		} else {
-			(true, 0)
+			MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
 		}
 	}
 
@@ -434,7 +446,7 @@ mod tests {
 		ext.clear_child_storage(child_info, b"dog");
 		assert_eq!(ext.child_storage(child_info, b"dog"), None);
 
-		ext.kill_child_storage(child_info, None);
+		let _ = ext.kill_child_storage(child_info, None, None);
 		assert_eq!(ext.child_storage(child_info, b"doe"), None);
 	}
 
@@ -456,8 +468,8 @@ mod tests {
 			],
 		});
 
-		let res = ext.kill_child_storage(child_info, None);
-		assert_eq!(res, (true, 3));
+		let res = ext.kill_child_storage(child_info, None, None);
+		assert_eq!(res.deconstruct(), (None, 3, 3, 3));
 	}
 
 	#[test]

substrate/primitives/state-machine/src/ext.rs

@@ -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);
 	}
 

substrate/primitives/state-machine/src/lib.rs

@@ -1348,6 +1348,7 @@ mod execution {
 mod tests {
 	use super::{ext::Ext, *};
 	use crate::{execution::CallResult, in_memory_backend::new_in_mem_hash_key};
+	use assert_matches::assert_matches;
 	use codec::{Decode, Encode};
 	use sp_core::{
 		map,
@@ -1572,7 +1573,7 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
-			ext.clear_prefix(b"ab", None);
+			let _ = ext.clear_prefix(b"ab", None, None);
 		}
 		overlay.commit_transaction().unwrap();
 
@@ -1596,7 +1597,10 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
-			assert_eq!((false, 1), ext.clear_prefix(b"ab", Some(1)));
+			assert_matches!(
+				ext.clear_prefix(b"ab", Some(1), None).deconstruct(),
+				(Some(_), 1, 3, 1)
+			);
 		}
 		overlay.commit_transaction().unwrap();
 
@@ -1638,7 +1642,8 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
-			assert_eq!(ext.kill_child_storage(&child_info, Some(2)), (false, 2));
+			let r = ext.kill_child_storage(&child_info, Some(2), None);
+			assert_matches!(r.deconstruct(), (Some(_), 2, 6, 2));
 		}
 
 		assert_eq!(
@@ -1673,14 +1678,37 @@ mod tests {
 		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, 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));
-		assert_eq!(ext.kill_child_storage(&child_info, None), (true, 4));
+		let r = ext.kill_child_storage(&child_info, Some(0), None).deconstruct();
+		assert_matches!(r, (Some(_), 0, 0, 0));
+		let r = ext
+			.kill_child_storage(&child_info, Some(1), r.0.as_ref().map(|x| &x[..]))
+			.deconstruct();
+		assert_matches!(r, (Some(_), 1, 1, 1));
+		let r = ext
+			.kill_child_storage(&child_info, Some(4), r.0.as_ref().map(|x| &x[..]))
+			.deconstruct();
+		// Only 3 items remaining to remove
+		assert_matches!(r, (None, 3, 3, 3));
+		let r = ext.kill_child_storage(&child_info, Some(1), None).deconstruct();
+		assert_matches!(r, (Some(_), 0, 0, 1));
+	}
+
+	#[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::<BlakeTwo256>::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, None).deconstruct(), (None, 4, 4, 4));
 	}
 
 	#[test]
@@ -1695,7 +1723,7 @@ mod tests {
 
 		ext.set_child_storage(child_info, b"abc".to_vec(), b"def".to_vec());
 		assert_eq!(ext.child_storage(child_info, b"abc"), Some(b"def".to_vec()));
-		ext.kill_child_storage(child_info, None);
+		let _ = ext.kill_child_storage(child_info, None, None);
 		assert_eq!(ext.child_storage(child_info, b"abc"), None);
 	}
 

substrate/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>,
-	) {
+	) -> 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`.

substrate/primitives/state-machine/src/overlayed_changes/mod.rs

@@ -299,7 +299,7 @@ impl OverlayedChanges {
 	/// Clear child storage of given storage key.
 	///
 	/// Can be rolled back or committed when called inside a transaction.
-	pub(crate) fn clear_child_storage(&mut self, child_info: &ChildInfo) {
+	pub(crate) fn clear_child_storage(&mut self, child_info: &ChildInfo) -> u32 {
 		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
 		let top = &self.top;
@@ -309,20 +309,20 @@ impl OverlayedChanges {
 			.or_insert_with(|| (top.spawn_child(), child_info.clone()));
 		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
-		changeset.clear_where(|_, _| true, extrinsic_index);
+		changeset.clear_where(|_, _| true, 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_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.

substrate/primitives/state-machine/src/proving_backend.rs

@@ -288,9 +288,10 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	) {
-		self.0.apply_to_keys_while(child_info, prefix, f)
+		self.0.apply_to_keys_while(child_info, prefix, start_at, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {

substrate/primitives/state-machine/src/read_only.rs

@@ -25,6 +25,7 @@ use sp_core::{
 	traits::Externalities,
 	Blake2Hasher,
 };
+use sp_externalities::MultiRemovalResults;
 use std::{
 	any::{Any, TypeId},
 	marker::PhantomData,
@@ -124,11 +125,21 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		unimplemented!("place_child_storage not supported in ReadOnlyExternalities")
 	}
 
-	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 {
 		unimplemented!("kill_child_storage is not supported in ReadOnlyExternalities")
 	}
 
-	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 {
 		unimplemented!("clear_prefix is not supported in ReadOnlyExternalities")
 	}
 
@@ -136,8 +147,9 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		&mut self,
 		_child_info: &ChildInfo,
 		_prefix: &[u8],
-		_limit: Option<u32>,
-	) -> (bool, u32) {
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		unimplemented!("clear_child_prefix is not supported in ReadOnlyExternalities")
 	}
 

substrate/primitives/state-machine/src/testing.rs

@@ -381,7 +381,10 @@ mod tests {
 		{
 			let mut ext = ext.ext();
 
-			assert!(!ext.kill_child_storage(&child_info, Some(2)).0, "Should not delete all keys");
+			assert!(
+				ext.kill_child_storage(&child_info, Some(2), None).maybe_cursor.is_some(),
+				"Should not delete all keys"
+			);
 
 			assert!(ext.child_storage(&child_info, &b"doe"[..]).is_none());
 			assert!(ext.child_storage(&child_info, &b"dog"[..]).is_none());

substrate/primitives/state-machine/src/trie_backend.rs

@@ -123,9 +123,10 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	) {
-		self.essence.apply_to_keys_while(child_info, prefix, f)
+		self.essence.apply_to_keys_while(child_info, prefix, start_at, f)
 	}
 
 	fn for_child_keys_with_prefix<F: FnMut(&[u8])>(

substrate/primitives/state-machine/src/trie_backend_essence.rs

@@ -267,7 +267,8 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
-		mut f: F,
+		start_at: Option<&[u8]>,
+		f: F,
 	) {
 		let root = if let Some(child_info) = child_info.as_ref() {
 			match self.child_root(child_info) {
@@ -283,7 +284,7 @@ where
 			self.root
 		};
 
-		self.trie_iter_key_inner(&root, prefix, |k| f(k), child_info)
+		self.trie_iter_key_inner(&root, prefix, f, child_info, start_at)
 	}
 
 	/// Execute given closure for all keys starting with prefix.
@@ -311,6 +312,7 @@ where
 				true
 			},
 			Some(child_info),
+			None,
 		)
 	}
 
@@ -324,28 +326,31 @@ where
 				true
 			},
 			None,
+			None,
 		)
 	}
 
 	fn trie_iter_key_inner<F: FnMut(&[u8]) -> bool>(
 		&self,
 		root: &H::Out,
-		prefix: Option<&[u8]>,
+		maybe_prefix: Option<&[u8]>,
 		mut f: F,
 		child_info: Option<&ChildInfo>,
+		maybe_start_at: Option<&[u8]>,
 	) {
 		let mut iter = move |db| -> sp_std::result::Result<(), Box<TrieError<H::Out>>> {
 			let trie = TrieDB::<H>::new(db, root)?;
-			let iter = if let Some(prefix) = prefix.as_ref() {
-				TrieDBKeyIterator::new_prefixed(&trie, prefix)?
-			} else {
-				TrieDBKeyIterator::new(&trie)?
-			};
+			let prefix = maybe_prefix.unwrap_or(&[]);
+			let iter = match maybe_start_at {
+				Some(start_at) =>
+					TrieDBKeyIterator::new_prefixed_then_seek(&trie, prefix, start_at),
+				None => TrieDBKeyIterator::new_prefixed(&trie, prefix),
+			}?;
 
 			for x in iter {
 				let key = x?;
 
-				debug_assert!(prefix
+				debug_assert!(maybe_prefix
 					.as_ref()
 					.map(|prefix| key.starts_with(prefix))
 					.unwrap_or(true));

substrate/primitives/tasks/src/async_externalities.rs

@@ -22,7 +22,7 @@ use sp_core::{
 	storage::{ChildInfo, StateVersion, TrackedStorageKey},
 	traits::{Externalities, RuntimeSpawn, RuntimeSpawnExt, SpawnNamed, TaskExecutorExt},
 };
-use sp_externalities::{Extensions, ExternalitiesExt as _};
+use sp_externalities::{Extensions, ExternalitiesExt as _, MultiRemovalResults};
 use std::any::{Any, TypeId};
 
 /// Simple state-less externalities for use in async context.
@@ -105,11 +105,21 @@ impl Externalities for AsyncExternalities {
 		panic!("`place_child_storage`: should not be used in async externalities!")
 	}
 
-	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 {
 		panic!("`kill_child_storage`: should not be used in async externalities!")
 	}
 
-	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 {
 		panic!("`clear_prefix`: should not be used in async externalities!")
 	}
 
@@ -117,8 +127,9 @@ impl Externalities for AsyncExternalities {
 		&mut self,
 		_child_info: &ChildInfo,
 		_prefix: &[u8],
-		_limit: Option<u32>,
-	) -> (bool, u32) {
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		panic!("`clear_child_prefix`: should not be used in async externalities!")
 	}