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>
paritytech · Jun 15, 2021 · cdc55fe · cdc55fe
1 parent f32aa22
commit cdc55fe
Show file tree

Hide file tree

Showing 36 changed files with 312 additions and 239 deletions.
diff --git a/client/db/src/bench.rs b/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)
 		}
 	}
 

diff --git a/client/db/src/lib.rs b/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])>(

diff --git a/client/db/src/storage_cache.rs b/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> {

diff --git a/client/executor/runtime-test/src/lib.rs b/client/executor/runtime-test/src/lib.rs
@@ -71,7 +71,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()
 	}
 

diff --git a/client/light/src/backend.rs b/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 => (),
 		}
 	}

diff --git a/frame/contracts/src/storage.rs b/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

diff --git a/frame/elections-phragmen/src/benchmarking.rs b/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! {

diff --git a/frame/im-online/src/lib.rs b/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);

diff --git a/frame/society/src/lib.rs b/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

diff --git a/frame/staking/src/lib.rs b/frame/staking/src/lib.rs
@@ -2655,9 +2655,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);

diff --git a/frame/staking/src/slashing.rs b/frame/staking/src/slashing.rs
@@ -543,8 +543,8 @@ impl<'a, T: 'a + Config> Drop for InspectingSpans<'a, T> {
 
 /// Clear slashing metadata for an obsolete era.
 pub(crate) fn clear_era_metadata<T: Config>(obsolete_era: EraIndex) {
-	<Pallet<T> as Store>::ValidatorSlashInEra::remove_prefix(&obsolete_era);
-	<Pallet<T> as Store>::NominatorSlashInEra::remove_prefix(&obsolete_era);
+	<Pallet<T> as Store>::ValidatorSlashInEra::remove_prefix(&obsolete_era, None);
+	<Pallet<T> as Store>::NominatorSlashInEra::remove_prefix(&obsolete_era, None);
 }
 
 /// Clear slashing metadata for a dead account.

diff --git a/frame/staking/src/testing_utils.rs b/frame/staking/src/testing_utils.rs
@@ -29,8 +29,8 @@ const SEED: u32 = 0;
 
 /// This function removes all validators and nominators from storage.
 pub fn clear_validators_and_nominators<T: Config>() {
-	Validators::<T>::remove_all();
-	Nominators::<T>::remove_all();
+	Validators::<T>::remove_all(None);
+	Nominators::<T>::remove_all(None);
 }
 
 /// Grab a funded user.

diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs
@@ -1005,7 +1005,10 @@ pub mod tests {
 			DoubleMap::insert(&key1, &(key2 + 1), &4u64);
 			DoubleMap::insert(&(key1 + 1), &key2, &4u64);
 			DoubleMap::insert(&(key1 + 1), &(key2 + 1), &4u64);
-			DoubleMap::remove_prefix(&key1);
+			assert!(matches!(
+				DoubleMap::remove_prefix(&key1, None),
+				sp_io::KillStorageResult::AllRemoved(0), // all in overlay
+			));
 			assert_eq!(DoubleMap::get(&key1, &key2), 0u64);
 			assert_eq!(DoubleMap::get(&key1, &(key2 + 1)), 0u64);
 			assert_eq!(DoubleMap::get(&(key1 + 1), &key2), 4u64);

diff --git a/frame/support/src/storage/child.rs b/frame/support/src/storage/child.rs
@@ -24,7 +24,7 @@
 use crate::sp_std::prelude::*;
 use codec::{Codec, Encode, Decode};
 pub use sp_core::storage::{ChildInfo, ChildType};
-pub use crate::sp_io::KillChildStorageResult;
+pub use crate::sp_io::KillStorageResult;
 
 /// Return the value of the item in storage under `key`, or `None` if there is no explicit entry.
 pub fn get<T: Decode + Sized>(
@@ -174,7 +174,7 @@ pub fn exists(
 pub fn kill_storage(
 	child_info: &ChildInfo,
 	limit: Option<u32>,
-) -> KillChildStorageResult {
+) -> KillStorageResult {
 	match child_info.child_type() {
 		ChildType::ParentKeyId => sp_io::default_child_storage::storage_kill(
 			child_info.storage_key(),

diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs
@@ -212,8 +212,9 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 		unhashed::kill(&Self::storage_double_map_final_key(k1, k2))
 	}
 
-	fn remove_prefix<KArg1>(k1: KArg1) where KArg1: EncodeLike<K1> {
-		unhashed::kill_prefix(Self::storage_double_map_final_key1(k1).as_ref())
+	fn remove_prefix<KArg1>(k1: KArg1, limit: Option<u32>) -> sp_io::KillStorageResult
+		where KArg1: EncodeLike<K1> {
+		unhashed::kill_prefix(Self::storage_double_map_final_key1(k1).as_ref(), limit)
 	}
 
 	fn iter_prefix_values<KArg1>(k1: KArg1) -> storage::PrefixIterator<V> where

diff --git a/frame/support/src/storage/generator/nmap.rs b/frame/support/src/storage/generator/nmap.rs
@@ -196,11 +196,11 @@ where
 		unhashed::kill(&Self::storage_n_map_final_key::<K, _>(key));
 	}
 
-	fn remove_prefix<KP>(partial_key: KP)
+	fn remove_prefix<KP>(partial_key: KP, limit: Option<u32>) -> sp_io::KillStorageResult
 	where
 		K: HasKeyPrefix<KP>,
 	{
-		unhashed::kill_prefix(&Self::storage_n_map_partial_key(partial_key));
+		unhashed::kill_prefix(&Self::storage_n_map_partial_key(partial_key), limit)
 	}
 
 	fn iter_prefix_values<KP>(partial_key: KP) -> PrefixIterator<V>

diff --git a/frame/support/src/storage/migration.rs b/frame/support/src/storage/migration.rs
@@ -244,7 +244,7 @@ pub fn remove_storage_prefix(module: &[u8], item: &[u8], hash: &[u8]) {
 	key[0..16].copy_from_slice(&Twox128::hash(module));
 	key[16..32].copy_from_slice(&Twox128::hash(item));
 	key[32..].copy_from_slice(hash);
-	frame_support::storage::unhashed::kill_prefix(&key)
+	frame_support::storage::unhashed::kill_prefix(&key, None);
 }
 
 /// Get a particular value in storage by the `module`, the map's `item` name and the key `hash`.

diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs
@@ -464,7 +464,8 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		KArg2: EncodeLike<K2>;
 
 	/// Remove all values under the first key.
-	fn remove_prefix<KArg1>(k1: KArg1) where KArg1: ?Sized + EncodeLike<K1>;
+	fn remove_prefix<KArg1>(k1: KArg1, limit: Option<u32>) -> sp_io::KillStorageResult
+		where KArg1: ?Sized + EncodeLike<K1>;
 
 	/// Iterate over values that share the first key.
 	fn iter_prefix_values<KArg1>(k1: KArg1) -> PrefixIterator<V>
@@ -589,7 +590,8 @@ pub trait StorageNMap<K: KeyGenerator, V: FullCodec> {
 	fn remove<KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter>(key: KArg);
 
 	/// Remove all values under the partial prefix key.
-	fn remove_prefix<KP>(partial_key: KP) where K: HasKeyPrefix<KP>;
+	fn remove_prefix<KP>(partial_key: KP, limit: Option<u32>) -> sp_io::KillStorageResult
+		where K: HasKeyPrefix<KP>;
 
 	/// Iterate over values that share the partial prefix key.
 	fn iter_prefix_values<KP>(partial_key: KP) -> PrefixIterator<V> where K: HasKeyPrefix<KP>;
@@ -880,8 +882,8 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	}
 
 	/// Remove all value of the storage.
-	fn remove_all() {
-		sp_io::storage::clear_prefix(&Self::final_prefix())
+	fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
+		sp_io::storage::clear_prefix(&Self::final_prefix(), limit)
 	}
 
 	/// Iter over all value of the storage.
@@ -1184,7 +1186,7 @@ mod test {
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2, 3, 4]);
 
 			// test removal
-			MyStorage::remove_all();
+			MyStorage::remove_all(None);
 			assert!(MyStorage::iter_values().collect::<Vec<_>>().is_empty());
 
 			// test migration
@@ -1194,7 +1196,7 @@ mod test {
 			assert!(MyStorage::iter_values().collect::<Vec<_>>().is_empty());
 			MyStorage::translate_values(|v: u32| Some(v as u64));
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2]);
-			MyStorage::remove_all();
+			MyStorage::remove_all(None);
 
 			// test migration 2
 			unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u128);
@@ -1206,7 +1208,7 @@ mod test {
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2, 3]);
 			MyStorage::translate_values(|v: u128| Some(v as u64));
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2, 3]);
-			MyStorage::remove_all();
+			MyStorage::remove_all(None);
 
 			// test that other values are not modified.
 			assert_eq!(unhashed::get(&key_before[..]), Some(32u64));

diff --git a/frame/support/src/storage/types/double_map.rs b/frame/support/src/storage/types/double_map.rs
@@ -205,8 +205,9 @@ where
 	}
 
 	/// Remove all values under the first key.
-	pub fn remove_prefix<KArg1>(k1: KArg1) where KArg1: ?Sized + EncodeLike<Key1> {
-		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::remove_prefix(k1)
+	pub fn remove_prefix<KArg1>(k1: KArg1, limit: Option<u32>) -> sp_io::KillStorageResult
+		where KArg1: ?Sized + EncodeLike<Key1> {
+		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::remove_prefix(k1, limit)
 	}
 
 	/// Iterate over values that share the first key.
@@ -316,8 +317,8 @@ where
 	}
 
 	/// Remove all value of the storage.
-	pub fn remove_all() {
-		<Self as crate::storage::StoragePrefixedMap<Value>>::remove_all()
+	pub fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
+		<Self as crate::storage::StoragePrefixedMap<Value>>::remove_all(limit)
 	}
 
 	/// Iter over all value of the storage.
@@ -615,7 +616,7 @@ mod test {
 
 			A::insert(3, 30, 10);
 			A::insert(4, 40, 10);
-			A::remove_all();
+			A::remove_all(None);
 			assert_eq!(A::contains_key(3, 30), false);
 			assert_eq!(A::contains_key(4, 40), false);
 
@@ -655,7 +656,7 @@ mod test {
 			assert_eq!(AValueQueryWithAnOnEmpty::DEFAULT.0.default_byte(), 97u32.encode());
 			assert_eq!(A::DEFAULT.0.default_byte(), Option::<u32>::None.encode());
 
-			WithLen::remove_all();
+			WithLen::remove_all(None);
 			assert_eq!(WithLen::decode_len(3, 30), None);
 			WithLen::append(0, 100, 10);
 			assert_eq!(WithLen::decode_len(0, 100), Some(1));
@@ -669,7 +670,7 @@ mod test {
 			assert_eq!(A::iter_prefix_values(4).collect::<Vec<_>>(), vec![13, 14]);
 			assert_eq!(A::iter_prefix(4).collect::<Vec<_>>(), vec![(40, 13), (41, 14)]);
 
-			A::remove_prefix(3);
+			A::remove_prefix(3, None);
 			assert_eq!(A::iter_prefix(3).collect::<Vec<_>>(), vec![]);
 			assert_eq!(A::iter_prefix(4).collect::<Vec<_>>(), vec![(40, 13), (41, 14)]);