Frame remove_all with size limit. (#9106)

* remove prefixed content with limit.

* test match

* factor comment and factor ext limit removal.

* fix benchmark

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

Author: 2 people

client/db/src/bench.rs

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

client/db/src/lib.rs

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

client/db/src/storage_cache.rs

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

client/executor/runtime-test/src/lib.rs

@@ -63,7 +63,7 @@ sp_core::wasm_export_functions! {
 	}
 
 	fn test_clear_prefix(input: Vec<u8>) -> Vec<u8> {
-		storage::clear_prefix(&input);
+		storage::clear_prefix(&input, None);
 		b"all ok!".to_vec()
 	}
 

client/light/src/backend.rs

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

frame/contracts/src/storage.rs

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

frame/elections-phragmen/src/benchmarking.rs

@@ -142,7 +142,7 @@ fn clean<T: Config>() {
 	<Members<T>>::kill();
 	<Candidates<T>>::kill();
 	<RunnersUp<T>>::kill();
-	<Voting<T>>::remove_all();
+	<Voting<T>>::remove_all(None);
 }
 
 benchmarks! {

frame/im-online/src/lib.rs

@@ -809,8 +809,8 @@ impl<T: Config> OneSessionHandler<T::AccountId> for Pallet<T> {
 		// Remove all received heartbeats and number of authored blocks from the
 		// current session, they have already been processed and won't be needed
 		// anymore.
-		ReceivedHeartbeats::<T>::remove_prefix(&T::ValidatorSet::session_index());
-		AuthoredBlocks::<T>::remove_prefix(&T::ValidatorSet::session_index());
+		ReceivedHeartbeats::<T>::remove_prefix(&T::ValidatorSet::session_index(), None);
+		AuthoredBlocks::<T>::remove_prefix(&T::ValidatorSet::session_index(), None);
 
 		if offenders.is_empty() {
 			Self::deposit_event(Event::<T>::AllGood);

frame/society/src/lib.rs

@@ -871,7 +871,7 @@ decl_module! {
 			Founder::<T, I>::kill();
 			Rules::<T, I>::kill();
 			Candidates::<T, I>::kill();
-			SuspendedCandidates::<T, I>::remove_all();
+			SuspendedCandidates::<T, I>::remove_all(None);
 			Self::deposit_event(RawEvent::Unfounded(founder));
 		}
 
@@ -1402,7 +1402,7 @@ impl<T: Config<I>, I: Instance> Module<T, I> {
 			}).collect::<Vec<_>>();
 
 			// Clean up all votes.
-			<Votes<T, I>>::remove_all();
+			<Votes<T, I>>::remove_all(None);
 
 			// Reward one of the voters who voted the right way.
 			if !total_slash.is_zero() {
@@ -1570,7 +1570,7 @@ impl<T: Config<I>, I: Instance> Module<T, I> {
 				}
 
 				// Clean up all votes.
-				<DefenderVotes<T, I>>::remove_all();
+				<DefenderVotes<T, I>>::remove_all(None);
 			}
 
 			// Avoid challenging if there's only two members since we never challenge the Head or

frame/staking/src/lib.rs

@@ -2846,9 +2846,9 @@ impl<T: Config> Pallet<T> {
 
 	/// Clear all era information for given era.
 	fn clear_era_information(era_index: EraIndex) {
-		<ErasStakers<T>>::remove_prefix(era_index);
-		<ErasStakersClipped<T>>::remove_prefix(era_index);
-		<ErasValidatorPrefs<T>>::remove_prefix(era_index);
+		<ErasStakers<T>>::remove_prefix(era_index, None);
+		<ErasStakersClipped<T>>::remove_prefix(era_index, None);
+		<ErasValidatorPrefs<T>>::remove_prefix(era_index, None);
 		<ErasValidatorReward<T>>::remove(era_index);
 		<ErasRewardPoints<T>>::remove(era_index);
 		<ErasTotalStake<T>>::remove(era_index);