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

client/db/src/bench.rs

@@ -350,7 +350,7 @@ impl<B: BlockT> StateBackend<HashFor<B>> for BenchmarkingState<B> {
 		}
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

client/db/src/lib.rs

@@ -195,7 +195,7 @@ 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 for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

client/db/src/storage_cache.rs

@@ -584,7 +584,7 @@ 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 for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
@@ -766,7 +766,7 @@ 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 for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

client/light/src/backend.rs

@@ -441,7 +441,7 @@ impl<H: Hasher> StateBackend<H> for GenesisOrUnavailableState<H>
 		}
 	}
 
-	fn for_keys_in_child_storage<A: FnMut(&[u8])>(
+	fn for_keys_in_child_storage<A: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		action: A,

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,33 @@ 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 at most `limit` keys from the specified child storage.
+	///
+	/// Deletes all keys from the overlay and up to `limit` keys from the backend.
+	/// This 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.
+	///
+	/// Please note that keys that are residing in the overlay for that child tie when
+	/// issuing this call are all deleted without counting towards the `limit`. Only keys
+	/// written during the current block are part of the overlay.
+	///
+	/// # Note
+	///
+	/// 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.
+	fn storage_kill_limited(
+		&mut self,
+		storage_key: &[u8],
+		limit: u32,
+	) -> bool {
+		let child_info = ChildInfo::new_default(storage_key);
+		self.kill_child_storage(&child_info, Some(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 for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,
@@ -263,7 +264,7 @@ 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 for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

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,40 @@ 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.for_keys_in_child_storage(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 {
+					return false;
+				}
+				self.overlay.set_child_storage(child_info, key.to_vec(), None);
+				num_deleted <= limit
+			});
+			all_deleted
+		} else {
+			self.backend.for_keys_in_child_storage(child_info, |key| {
+				self.overlay.set_child_storage(child_info, key.to_vec(), None);
+				true
+			});
+			true
+		}
 	}
 
 	fn clear_prefix(&mut self, prefix: &[u8]) {

primitives/state-machine/src/lib.rs

@@ -1147,6 +1147,86 @@ mod tests {
 		);
 	}
 
+	#[test]
+	fn limited_child_kill_works() {
+		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);
+
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_child_storage(&child_info, b"1".to_vec(), Some(b"1312".to_vec()));
+		overlay.set_child_storage(&child_info, b"2".to_vec(), Some(b"1312".to_vec()));
+		overlay.set_child_storage(&child_info, b"3".to_vec(), Some(b"1312".to_vec()));
+		overlay.set_child_storage(&child_info, b"4".to_vec(), Some(b"1312".to_vec()));
+
+		{
+			let mut offchain_overlay = Default::default();
+			let mut cache = StorageTransactionCache::default();
+			let mut ext = Ext::new(
+				&mut overlay,
+				&mut offchain_overlay,
+				&mut cache,
+				&backend,
+				changes_trie::disabled_state::<_, u64>(),
+				None,
+			);
+			assert_eq!(ext.kill_child_storage(&child_info, Some(2)), false);
+		}
+
+		assert_eq!(
+			overlay.children()
+				.flat_map(|(iter, _child_info)| iter)
+				.map(|(k, v)| (k.clone(), v.value().clone()))
+				.collect::<BTreeMap<_, _>>(),
+			map![
+				b"1".to_vec() => None.into(),
+				b"2".to_vec() => None.into(),
+				b"3".to_vec() => None.into(),
+				b"4".to_vec() => None.into(),
+				b"a".to_vec() => None.into(),
+				b"b".to_vec() => None.into(),
+			],
+		);
+	}
+
+	#[test]
+	fn limited_child_kill_off_by_one_works() {
+		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);
+		let mut overlay = OverlayedChanges::default();
+		let mut offchain_overlay = Default::default();
+		let mut cache = StorageTransactionCache::default();
+		let mut ext = Ext::new(
+			&mut overlay,
+			&mut offchain_overlay,
+			&mut cache,
+			&backend,
+			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);
+	}
+
 	#[test]
 	fn set_child_storage_works() {
 		let child_info = ChildInfo::new_default(b"sub1");
@@ -1179,6 +1259,7 @@ mod tests {
 		);
 		ext.kill_child_storage(
 			child_info,
+			None,
 		);
 		assert_eq!(
 			ext.child_storage(

primitives/state-machine/src/proving_backend.rs

@@ -204,7 +204,7 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		self.0.child_storage(child_info, key)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

primitives/state-machine/src/read_only.rs

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

primitives/state-machine/src/trie_backend.rs

@@ -113,7 +113,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 		self.essence.for_key_values_with_prefix(prefix, f)
 	}
 
-	fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	fn for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

primitives/state-machine/src/trie_backend_essence.rs

@@ -190,7 +190,8 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 	}
 
 	/// Retrieve all entries keys of child storage and call `f` for each of those keys.
-	pub fn for_keys_in_child_storage<F: FnMut(&[u8])>(
+	/// Aborts as soon as `f` returns false.
+	pub fn for_keys_in_child_storage<F: FnMut(&[u8]) -> bool>(
 		&self,
 		child_info: &ChildInfo,
 		f: F,

primitives/tasks/src/async_externalities.rs

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

primitives/trie/src/lib.rs

@@ -271,7 +271,8 @@ pub fn child_delta_trie_root<L: TrieConfiguration, I, A, B, DB, RD, V>(
 }
 
 /// Call `f` for all keys in a child trie.
-pub fn for_keys_in_child_trie<L: TrieConfiguration, F: FnMut(&[u8]), DB>(
+/// Aborts as soon as `f` returns false.
+pub fn for_keys_in_child_trie<L: TrieConfiguration, F: FnMut(&[u8]) -> bool, DB>(
 	keyspace: &[u8],
 	db: &DB,
 	root_slice: &[u8],
@@ -290,7 +291,9 @@ pub fn for_keys_in_child_trie<L: TrieConfiguration, F: FnMut(&[u8]), DB>(
 
 	for x in iter {
 		let (key, _) = x?;
-		f(&key);
+		if !f(&key) {
+			break;
+		}
 	}
 
 	Ok(())