Allow capping the amount of work performed when deleting a child trie

client/db/src/bench.rs

@@ -350,13 +350,13 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		}
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
 	) {
 		if let Some(ref state) = *self.state.borrow() {
-			state.for_keys_in_child_storage(child_info, f)
+			state.apply_to_child_keys_while(child_info, f)
 		}
 	}
 

client/db/src/lib.rs

@@ -195,12 +195,12 @@ impl<B: BlockT> StateBackend<HashFor<B>> for RefTrackingState<B> {
 		self.state.for_key_values_with_prefix(prefix, f)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
 	) {
-		self.state.for_keys_in_child_storage(child_info, f)
+		self.state.apply_to_child_keys_while(child_info, f)
 	}
 
 	fn for_child_keys_with_prefix<F: FnMut(&[u8])>(

client/db/src/storage_cache.rs

@@ -584,12 +584,12 @@ impl<S: StateBackend<HashFor<B>>, B: BlockT> StateBackend<HashFor<B>> for Cachin
 		self.state.exists_child_storage(child_info, key)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
 	) {
-		self.state.for_keys_in_child_storage(child_info, f)
+		self.state.apply_to_child_keys_while(child_info, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
@@ -766,12 +766,12 @@ impl<S: StateBackend<HashFor<B>>, B: BlockT> StateBackend<HashFor<B>> for Syncin
 		self.caching_state().exists_child_storage(child_info, key)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
 	) {
-		self.caching_state().for_keys_in_child_storage(child_info, f)
+		self.caching_state().apply_to_child_keys_while(child_info, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {

client/light/src/backend.rs

@@ -441,14 +441,14 @@ impl<H: Hasher> StateBackend<H> for GenesisOrUnavailableState<H>
 		}
 	}
 
-	fn for_keys_in_child_storage<A: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<A: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		action: A,
 	) {
 		match *self {
 			GenesisOrUnavailableState::Genesis(ref state) =>
-				state.for_keys_in_child_storage(child_info, action),
+				state.apply_to_child_keys_while(child_info, action),
 			GenesisOrUnavailableState::Unavailable => (),
 		}
 	}

frame/contracts/src/rent.rs

@@ -240,6 +240,7 @@ where
 				<ContractInfoOf<T>>::remove(account);
 				child::kill_storage(
 					&alive_contract_info.child_trie_info(),
+					None,
 				);
 				<Module<T>>::deposit_event(RawEvent::Evicted(account.clone(), false));
 				None
@@ -263,6 +264,7 @@ where
 
 				child::kill_storage(
 					&alive_contract_info.child_trie_info(),
+					None,
 				);
 
 				<Module<T>>::deposit_event(RawEvent::Evicted(account.clone(), true));

frame/contracts/src/storage.rs

@@ -195,7 +195,7 @@ where
 	/// This function doesn't affect the account.
 	pub fn destroy_contract(address: &AccountIdOf<T>, trie_id: &TrieId) {
 		<ContractInfoOf<T>>::remove(address);
-		child::kill_storage(&crate::child_trie_info(&trie_id));
+		child::kill_storage(&crate::child_trie_info(&trie_id), None);
 	}
 
 	/// This generator uses inner counter for account id and applies the hash over `AccountId +

frame/support/src/storage/child.rs

@@ -25,6 +25,14 @@ 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,
+}
+
 /// Return the value of the item in storage under `key`, or `None` if there is no explicit entry.
 pub fn get<T: Decode + Sized>(
 	child_info: &ChildInfo,
@@ -148,13 +156,37 @@ pub fn exists(
 }
 
 /// Remove all `storage_key` key/values
+///
+/// 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).
+///
+/// # 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.
 pub fn kill_storage(
 	child_info: &ChildInfo,
-) {
-	match child_info.child_type() {
+	limit: Option<u32>,
+) -> KillOutcome {
+	let all_removed = 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,
 	}
 }
 

primitives/externalities/src/lib.rs

@@ -137,7 +137,17 @@ pub trait Externalities: ExtensionStore {
 	) -> Option<Vec<u8>>;
 
 	/// Clear an entire child storage.
-	fn kill_child_storage(&mut self, child_info: &ChildInfo);
+	///
+	/// 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
+	/// removed completely and `false` if there are remaining keys after the function
+	/// returns.
+	///
+	/// # 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;
 
 	/// Clear storage entries which keys are start with the given prefix.
 	fn clear_prefix(&mut self, prefix: &[u8]);

primitives/io/src/lib.rs

@@ -279,7 +279,35 @@ pub trait DefaultChildStorage {
 		storage_key: &[u8],
 	) {
 		let child_info = ChildInfo::new_default(storage_key);
-		self.kill_child_storage(&child_info);
+		self.kill_child_storage(&child_info, None);
+	}
+
+	/// 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 false iff some keys are remaining in
+	/// the child trie after the functions returns.
+	///
+	/// # 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(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)
 	}
 
 	/// Check a child storage key.

primitives/state-machine/src/backend.rs

@@ -94,7 +94,8 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 	) -> Result<Option<StorageKey>, Self::Error>;
 
 	/// Retrieve all entries keys of child storage and call `f` for each of those keys.
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	/// Aborts as soon as `f` returns false.
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
@@ -263,12 +264,12 @@ impl<'a, T: Backend<H>, H: Hasher> Backend<H> for &'a T {
 		(*self).child_storage(child_info, key)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
 	) {
-		(*self).for_keys_in_child_storage(child_info, f)
+		(*self).apply_to_child_keys_while(child_info, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<StorageKey>, Self::Error> {

primitives/state-machine/src/basic.rs

@@ -210,8 +210,10 @@ impl Externalities for BasicExternalities {
 	fn kill_child_storage(
 		&mut self,
 		child_info: &ChildInfo,
-	) {
+		_limit: Option<u32>,
+	) -> bool {
 		self.inner.children_default.remove(child_info.storage_key());
+		true
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {
@@ -407,7 +409,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);
+		ext.kill_child_storage(child_info, None);
 		assert_eq!(ext.child_storage(child_info, b"doe"), None);
 	}
 

primitives/state-machine/src/ext.rs

@@ -411,18 +411,41 @@ where
 	fn kill_child_storage(
 		&mut self,
 		child_info: &ChildInfo,
-	) {
+		limit: Option<u32>,
+	) -> bool {
 		trace!(target: "state", "{:04x}: KillChild({})",
 			self.id,
 			HexDisplay::from(&child_info.storage_key()),
 		);
 		let _guard = guard();
-
 		self.mark_dirty();
 		self.overlay.clear_child_storage(child_info);
-		self.backend.for_keys_in_child_storage(child_info, |key| {
-			self.overlay.set_child_storage(child_info, key.to_vec(), None);
-		});
+
+		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 {
+					all_deleted = false;
+					return false;
+				}
+				if let Some(num) = num_deleted.checked_add(1) {
+					num_deleted = num;
+				} else {
+					all_deleted = false;
+					return false;
+				}
+				self.overlay.set_child_storage(child_info, key.to_vec(), None);
+				true
+			});
+			all_deleted
+		} else {
+			self.backend.apply_to_child_keys_while(child_info, |key| {
+				self.overlay.set_child_storage(child_info, key.to_vec(), None);
+				true
+			});
+			true
+		}
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {