Return number of keys removed when calling storage_kill on child trie (#8166)

* Initial piping of returning amount of keys killed

* One more test for `None` limit

* forgot to update

* fix return value

* use version 3

* Update to return `KillOutcome`

* Update name to KillChildStorageResult

substrate/frame/contracts/src/storage.rs

@@ -32,7 +32,7 @@ use sp_core::crypto::UncheckedFrom;
 use frame_support::{
 	dispatch::DispatchResult,
 	debug,
-	storage::child::{self, KillOutcome},
+	storage::child::{self, KillChildStorageResult},
 	traits::Get,
 	weights::Weight,
 };
@@ -269,13 +269,13 @@ where
 				let removed = queue.swap_remove(0);
 				match outcome {
 					// This should not happen as our budget was large enough to remove all keys.
-					KillOutcome::SomeRemaining => {
+					KillChildStorageResult::SomeRemaining(_) => {
 						debug::error!(
 							"After deletion keys are remaining in this child trie: {:?}",
 							removed.trie_id,
 						);
 					},
-					KillOutcome::AllRemoved => (),
+					KillChildStorageResult::AllRemoved(_) => (),
 				}
 			}
 			remaining_key_budget = remaining_key_budget

substrate/frame/support/src/storage/child.rs

@@ -24,14 +24,7 @@
 use crate::sp_std::prelude::*;
 use codec::{Codec, Encode, Decode};
 pub use sp_core::storage::{ChildInfo, ChildType};
-
-/// The outcome of calling [`kill_storage`].
-pub enum KillOutcome {
-	/// No key remains in the child trie.
-	AllRemoved,
-	/// At least one key still resides in the child trie due to the supplied limit.
-	SomeRemaining,
-}
+pub use crate::sp_io::KillChildStorageResult;
 
 /// Return the value of the item in storage under `key`, or `None` if there is no explicit entry.
 pub fn get<T: Decode + Sized>(
@@ -177,16 +170,12 @@ pub fn exists(
 pub fn kill_storage(
 	child_info: &ChildInfo,
 	limit: Option<u32>,
-) -> KillOutcome {
-	let all_removed = match child_info.child_type() {
+) -> KillChildStorageResult {
+	match child_info.child_type() {
 		ChildType::ParentKeyId => sp_io::default_child_storage::storage_kill(
 			child_info.storage_key(),
 			limit
 		),
-	};
-	match all_removed {
-		true => KillOutcome::AllRemoved,
-		false => KillOutcome::SomeRemaining,
 	}
 }
 

substrate/primitives/externalities/src/lib.rs

@@ -139,15 +139,16 @@ 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`. Returns `true` if the child trie was
+	/// 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.
+	/// returns. Returned `u32` is the number of keys that was removed at the end of the
+	/// operation.
 	///
 	/// # 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;
+	fn kill_child_storage(&mut self, child_info: &ChildInfo, limit: Option<u32>) -> (bool, u32);
 
 	/// Clear storage entries which keys are start with the given prefix.
 	fn clear_prefix(&mut self, prefix: &[u8]);

substrate/primitives/io/src/lib.rs

@@ -57,7 +57,7 @@ use sp_core::{
 use sp_trie::{TrieConfiguration, trie_types::Layout};
 
 use sp_runtime_interface::{runtime_interface, Pointer};
-use sp_runtime_interface::pass_by::PassBy;
+use sp_runtime_interface::pass_by::{PassBy, PassByCodec};
 
 use codec::{Encode, Decode};
 
@@ -81,6 +81,16 @@ pub enum EcdsaVerifyError {
 	BadSignature,
 }
 
+/// The outcome of calling [`kill_storage`]. Returned value is the number of storage items
+/// removed from the trie from making the `kill_storage` call.
+#[derive(PassByCodec, Encode, Decode)]
+pub enum KillChildStorageResult {
+	/// No key remains in the child trie.
+	AllRemoved(u32),
+	/// At least one key still resides in the child trie due to the supplied limit.
+	SomeRemaining(u32),
+}
+
 /// Interface for accessing the storage from within the runtime.
 #[runtime_interface]
 pub trait Storage {
@@ -290,7 +300,7 @@ pub trait DefaultChildStorage {
 	/// The limit can be used to partially delete a child trie in case it is too large
 	/// to delete in one go (block).
 	///
-	/// It returns false iff some keys are remaining in
+	/// It returns a boolean false iff some keys are remaining in
 	/// the child trie after the functions returns.
 	///
 	/// # Note
@@ -307,7 +317,41 @@ 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);
-		self.kill_child_storage(&child_info, limit)
+		let (all_removed, _num_removed) = self.kill_child_storage(&child_info, limit);
+		all_removed
+	}
+
+	/// Clear a child storage key.
+	///
+	/// Deletes all keys from the overlay and up to `limit` keys from the backend if
+	/// it is set to `Some`. No limit is applied when `limit` is set to `None`.
+	///
+	/// The limit can be used to partially delete a child trie in case it is too large
+	/// to delete in one go (block).
+	///
+	/// It returns a boolean false iff some keys are remaining in
+	/// the child trie after the functions returns. Also returns a `u32` with
+	/// the number of keys removed from the process.
+	///
+	/// # Note
+	///
+	/// Please note that keys that are residing in the overlay for that child trie when
+	/// issuing this call are all deleted without counting towards the `limit`. Only keys
+	/// written during the current block are part of the overlay. Deleting with a `limit`
+	/// mostly makes sense with an empty overlay for that child trie.
+	///
+	/// Calling this function multiple times per block for the same `storage_key` does
+	/// not make much sense because it is not cumulative when called inside the same block.
+	/// Use this function to distribute the deletion of a single child trie across multiple
+	/// blocks.
+	#[version(3)]
+	fn storage_kill(&mut self, storage_key: &[u8], limit: Option<u32>) -> KillChildStorageResult {
+		let child_info = ChildInfo::new_default(storage_key);
+		let (all_removed, num_removed) = self.kill_child_storage(&child_info, limit);
+		match all_removed {
+			true => KillChildStorageResult::AllRemoved(num_removed),
+			false => KillChildStorageResult::SomeRemaining(num_removed),
+		}
 	}
 
 	/// Check a child storage key.

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

@@ -211,9 +211,9 @@ impl Externalities for BasicExternalities {
 		&mut self,
 		child_info: &ChildInfo,
 		_limit: Option<u32>,
-	) -> bool {
-		self.inner.children_default.remove(child_info.storage_key());
-		true
+	) -> (bool, u32) {
+		let num_removed = self.inner.children_default.remove(child_info.storage_key()).map(|c| c.data.len()).unwrap_or(0);
+		(true, num_removed as u32)
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {
@@ -411,6 +411,29 @@ mod tests {
 		assert_eq!(ext.child_storage(child_info, b"doe"), None);
 	}
 
+	#[test]
+	fn kill_child_storage_returns_num_elements_removed() {
+		let child_info = ChildInfo::new_default(b"storage_key");
+		let child_info = &child_info;
+		let mut ext = BasicExternalities::new(Storage {
+			top: Default::default(),
+			children_default: map![
+				child_info.storage_key().to_vec() => StorageChild {
+					data: map![
+						b"doe".to_vec() => b"reindeer".to_vec(),
+						b"dog".to_vec() => b"puppy".to_vec(),
+						b"hello".to_vec() => b"world".to_vec(),
+					],
+					child_info: child_info.to_owned(),
+				}
+			]
+		});
+
+
+		let res = ext.kill_child_storage(child_info, None);
+		assert_eq!(res, (true, 3));
+	}
+
 	#[test]
 	fn basic_externalities_is_empty() {
 		// Make sure no values are set by default in `BasicExternalities`.

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

@@ -391,7 +391,7 @@ where
 		&mut self,
 		child_info: &ChildInfo,
 		limit: Option<u32>,
-	) -> bool {
+	) -> (bool, u32) {
 		trace!(target: "state", "{:04x}: KillChild({})",
 			self.id,
 			HexDisplay::from(&child_info.storage_key()),
@@ -399,9 +399,9 @@ where
 		let _guard = guard();
 		self.mark_dirty();
 		self.overlay.clear_child_storage(child_info);
+		let mut num_deleted: u32 = 0;
 
 		if let Some(limit) = limit {
-			let mut num_deleted: u32 = 0;
 			let mut all_deleted = true;
 			self.backend.apply_to_child_keys_while(child_info, |key| {
 				if num_deleted == limit {
@@ -417,13 +417,14 @@ where
 				self.overlay.set_child_storage(child_info, key.to_vec(), None);
 				true
 			});
-			all_deleted
+			(all_deleted, num_deleted)
 		} else {
 			self.backend.apply_to_child_keys_while(child_info, |key| {
+				num_deleted = num_deleted.saturating_add(1);
 				self.overlay.set_child_storage(child_info, key.to_vec(), None);
 				true
 			});
-			true
+			(true, num_deleted)
 		}
 	}
 

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

@@ -1159,7 +1159,7 @@ mod tests {
 				changes_trie::disabled_state::<_, u64>(),
 				None,
 			);
-			assert_eq!(ext.kill_child_storage(&child_info, Some(2)), false);
+			assert_eq!(ext.kill_child_storage(&child_info, Some(2)), (false, 2));
 		}
 
 		assert_eq!(
@@ -1199,12 +1199,14 @@ mod tests {
 			changes_trie::disabled_state::<_, u64>(),
 			None,
 		);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(0)), false);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(1)), false);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(2)), false);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(3)), false);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(4)), true);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(5)), true);
+		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));
 	}
 
 	#[test]

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

@@ -132,7 +132,7 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		&mut self,
 		_child_info: &ChildInfo,
 		_limit: Option<u32>,
-	) -> bool {
+	) -> (bool, u32) {
 		unimplemented!("kill_child_storage is not supported in ReadOnlyExternalities")
 	}
 

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

@@ -326,7 +326,7 @@ mod tests {
 		{
 			let mut ext = ext.ext();
 
-			assert!(!ext.kill_child_storage(&child_info, Some(2)), "Should not delete all keys");
+			assert!(!ext.kill_child_storage(&child_info, Some(2)).0, "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/tasks/src/async_externalities.rs

@@ -119,7 +119,7 @@ impl Externalities for AsyncExternalities {
 		&mut self,
 		_child_info: &ChildInfo,
 		_limit: Option<u32>,
-	) -> bool {
+	) -> (bool, u32) {
 		panic!("`kill_child_storage`: should not be used in async externalities!")
 	}