Fix pallet bags list and doc (paritytech#10231)

* fix bags list

* improve doc

* doc

* inner doc

* fix test

* Update docs in frame/election-provider-support/src/lib.rs

* fix staking impl

* prepend unsafe to clear and regenerate

* fix test

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

frame/bags-list/remote-tests/src/migration.rs

@@ -47,7 +47,7 @@ pub async fn execute<Runtime: RuntimeT, Block: BlockT + DeserializeOwned>(
 		log::info!(target: LOG_TARGET, "Nominator count: {}", pre_migrate_nominator_count);
 
 		// run the actual migration,
-		let moved = <Runtime as pallet_staking::Config>::SortedListProvider::regenerate(
+		let moved = <Runtime as pallet_staking::Config>::SortedListProvider::unsafe_regenerate(
 			pallet_staking::Nominators::<Runtime>::iter().map(|(n, _)| n),
 			pallet_staking::Pallet::<Runtime>::weight_of_fn(),
 		);

frame/bags-list/src/benchmarks.rs

@@ -35,7 +35,8 @@ frame_benchmarking::benchmarks! {
 		//   node in the destination in addition to the work we do otherwise. (2 W/R)
 
 		// clear any pre-existing storage.
-		List::<T>::clear(None);
+		// NOTE: safe to call outside block production
+		List::<T>::unsafe_clear();
 
 		// define our origin and destination thresholds.
 		let origin_bag_thresh = T::BagThresholds::get()[0];
@@ -94,7 +95,8 @@ frame_benchmarking::benchmarks! {
 		//   node in the destination in addition to the work we do otherwise. (2 W/R)
 
 		// clear any pre-existing storage.
-		List::<T>::clear(None);
+		// NOTE: safe to call outside block production
+		List::<T>::unsafe_clear();
 
 		// define our origin and destination thresholds.
 		let origin_bag_thresh = T::BagThresholds::get()[0];

frame/bags-list/src/lib.rs

@@ -26,7 +26,7 @@
 //! the weights of accounts via [`frame_election_provider_support::VoteWeightProvider`].
 //!
 //! This pallet is not configurable at genesis. Whoever uses it should call appropriate functions of
-//! the `SortedListProvider` (e.g. `on_insert`, or `regenerate`) at their genesis.
+//! the `SortedListProvider` (e.g. `on_insert`, or `unsafe_regenerate`) at their genesis.
 //!
 //! # Goals
 //!
@@ -256,11 +256,14 @@ impl<T: Config> SortedListProvider<T::AccountId> for Pallet<T> {
 		List::<T>::remove(id)
 	}
 
-	fn regenerate(
+	fn unsafe_regenerate(
 		all: impl IntoIterator<Item = T::AccountId>,
 		weight_of: Box<dyn Fn(&T::AccountId) -> VoteWeight>,
 	) -> u32 {
-		List::<T>::regenerate(all, weight_of)
+		// NOTE: This call is unsafe for the same reason as SortedListProvider::unsafe_regenerate.
+		// I.e. because it can lead to many storage accesses.
+		// So it is ok to call it as caller must ensure the conditions.
+		List::<T>::unsafe_regenerate(all, weight_of)
 	}
 
 	#[cfg(feature = "std")]
@@ -273,8 +276,11 @@ impl<T: Config> SortedListProvider<T::AccountId> for Pallet<T> {
 		Ok(())
 	}
 
-	fn clear(maybe_count: Option<u32>) -> u32 {
-		List::<T>::clear(maybe_count)
+	fn unsafe_clear() {
+		// NOTE: This call is unsafe for the same reason as SortedListProvider::unsafe_clear.
+		// I.e. because it can lead to many storage accesses.
+		// So it is ok to call it as caller must ensure the conditions.
+		List::<T>::unsafe_clear()
 	}
 
 	#[cfg(feature = "runtime-benchmarks")]

frame/bags-list/src/list/mod.rs

@@ -76,19 +76,15 @@ pub fn notional_bag_for<T: Config>(weight: VoteWeight) -> VoteWeight {
 pub struct List<T: Config>(PhantomData<T>);
 
 impl<T: Config> List<T> {
-	/// Remove all data associated with the list from storage. Parameter `items` is the number of
-	/// items to clear from the list.
+	/// Remove all data associated with the list from storage.
 	///
 	/// ## WARNING
 	///
-	/// `None` will clear all items and should generally not be used in production as it could lead
-	/// to a very large number of storage accesses.
-	pub(crate) fn clear(maybe_count: Option<u32>) -> u32 {
-		crate::ListBags::<T>::remove_all(maybe_count);
-		let pre = crate::ListNodes::<T>::count();
-		crate::ListNodes::<T>::remove_all(maybe_count);
-		let post = crate::ListNodes::<T>::count();
-		pre.saturating_sub(post)
+	/// this function should generally not be used in production as it could lead to a very large
+	/// number of storage accesses.
+	pub(crate) fn unsafe_clear() {
+		crate::ListBags::<T>::remove_all(None);
+		crate::ListNodes::<T>::remove_all();
 	}
 
 	/// Regenerate all of the data from the given ids.
@@ -100,11 +96,14 @@ impl<T: Config> List<T> {
 	/// pallet using this `List`.
 	///
 	/// Returns the number of ids migrated.
-	pub fn regenerate(
+	pub fn unsafe_regenerate(
 		all: impl IntoIterator<Item = T::AccountId>,
 		weight_of: Box<dyn Fn(&T::AccountId) -> VoteWeight>,
 	) -> u32 {
-		Self::clear(None);
+		// NOTE: This call is unsafe for the same reason as SortedListProvider::unsafe_regenerate.
+		// I.e. because it can lead to many storage accesses.
+		// So it is ok to call it as caller must ensure the conditions.
+		Self::unsafe_clear();
 		Self::insert_many(all, weight_of)
 	}
 

frame/election-provider-support/src/lib.rs

@@ -333,15 +333,23 @@ pub trait SortedListProvider<AccountId> {
 	/// Regenerate this list from scratch. Returns the count of items inserted.
 	///
 	/// This should typically only be used at a runtime upgrade.
-	fn regenerate(
+	///
+	/// ## WARNING
+	///
+	/// This function should be called with care, regenerate will remove the current list write the
+	/// new list, which can lead to too many storage accesses, exhausting the block weight.
+	fn unsafe_regenerate(
 		all: impl IntoIterator<Item = AccountId>,
 		weight_of: Box<dyn Fn(&AccountId) -> VoteWeight>,
 	) -> u32;
 
-	/// Remove `maybe_count` number of items from the list. Returns the number of items actually
-	/// removed. WARNING: removes all items if `maybe_count` is `None`, which should never be done
-	/// in production settings because it can lead to an unbounded amount of storage accesses.
-	fn clear(maybe_count: Option<u32>) -> u32;
+	/// Remove all items from the list.
+	///
+	/// ## WARNING
+	///
+	/// This function should never be called in production settings because it can lead to an
+	/// unbounded amount of storage accesses.
+	fn unsafe_clear();
 
 	/// Sanity check internal state of list. Only meant for debug compilation.
 	fn sanity_check() -> Result<(), &'static str>;

frame/staking/src/migrations.rs

@@ -40,7 +40,7 @@ pub mod v8 {
 		if StorageVersion::<T>::get() == crate::Releases::V7_0_0 {
 			crate::log!(info, "migrating staking to Releases::V8_0_0");
 
-			let migrated = T::SortedListProvider::regenerate(
+			let migrated = T::SortedListProvider::unsafe_regenerate(
 				Nominators::<T>::iter().map(|(id, _)| id),
 				Pallet::<T>::weight_of_fn(),
 			);

frame/staking/src/pallet/impls.rs

@@ -971,7 +971,7 @@ impl<T: Config> ElectionDataProvider<T::AccountId, BlockNumberFor<T>> for Pallet
 		<Nominators<T>>::remove_all(None);
 		<CounterForNominators<T>>::kill();
 		<CounterForValidators<T>>::kill();
-		let _ = T::SortedListProvider::clear(None);
+		T::SortedListProvider::unsafe_clear();
 	}
 
 	#[cfg(feature = "runtime-benchmarks")]
@@ -1299,7 +1299,7 @@ impl<T: Config> SortedListProvider<T::AccountId> for UseNominatorsMap<T> {
 	fn on_remove(_: &T::AccountId) {
 		// nothing to do on remove.
 	}
-	fn regenerate(
+	fn unsafe_regenerate(
 		_: impl IntoIterator<Item = T::AccountId>,
 		_: Box<dyn Fn(&T::AccountId) -> VoteWeight>,
 	) -> u32 {
@@ -1309,13 +1309,10 @@ impl<T: Config> SortedListProvider<T::AccountId> for UseNominatorsMap<T> {
 	fn sanity_check() -> Result<(), &'static str> {
 		Ok(())
 	}
-	fn clear(maybe_count: Option<u32>) -> u32 {
-		Nominators::<T>::remove_all(maybe_count);
-		if let Some(count) = maybe_count {
-			CounterForNominators::<T>::mutate(|noms| *noms - count);
-			count
-		} else {
-			CounterForNominators::<T>::take()
-		}
+	fn unsafe_clear() {
+		// NOTE: Caller must ensure this doesn't lead to too many storage accesses. This is a
+		// condition of SortedListProvider::unsafe_clear.
+		Nominators::<T>::remove_all(None);
+		CounterForNominators::<T>::take();
 	}
 }

frame/staking/src/testing_utils.rs

@@ -42,7 +42,8 @@ pub fn clear_validators_and_nominators<T: Config>() {
 	// whenever we touch nominators counter we should update `T::SortedListProvider` as well.
 	Nominators::<T>::remove_all(None);
 	CounterForNominators::<T>::kill();
-	let _ = T::SortedListProvider::clear(None);
+	// NOTE: safe to call outside block production
+	T::SortedListProvider::unsafe_clear();
 }
 
 /// Grab a funded user.

frame/support/src/storage/mod.rs

@@ -1133,7 +1133,11 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 		crate::storage::storage_prefix(Self::module_prefix(), Self::storage_prefix())
 	}
 
-	/// Remove all value of the storage.
+	/// Remove all values of the storage in the overlay and up to `limit` in the backend.
+	///
+	/// All values in the client overlay will be deleted, if there is some `limit` then up to
+	/// `limit` values are deleted from the client backend, if `limit` is none then all values in
+	/// the client backend are deleted.
 	fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
 		sp_io::storage::clear_prefix(&Self::final_prefix(), limit)
 	}

frame/support/src/storage/types/counted_map.rs

@@ -31,7 +31,6 @@ use crate::{
 	Never,
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, MaxEncodedLen, Ref};
-use sp_arithmetic::traits::Bounded;
 use sp_runtime::traits::Saturating;
 use sp_std::prelude::*;
 
@@ -263,10 +262,12 @@ where
 	}
 
 	/// Remove all value of the storage.
-	pub fn remove_all(maybe_limit: Option<u32>) {
-		let leftover = Self::count().saturating_sub(maybe_limit.unwrap_or_else(Bounded::max_value));
-		CounterFor::<Prefix>::set(leftover);
-		<Self as MapWrapper>::Map::remove_all(maybe_limit);
+	pub fn remove_all() {
+		// NOTE: it is not possible to remove up to some limit because
+		// `sp_io::storage::clear_prefix` and `StorageMap::remove_all` don't give the number of
+		// value removed from the overlay.
+		CounterFor::<Prefix>::set(0u32);
+		<Self as MapWrapper>::Map::remove_all(None);
 	}
 
 	/// Iter over all value of the storage.
@@ -678,7 +679,7 @@ mod test {
 			assert_eq!(A::count(), 2);
 
 			// Remove all.
-			A::remove_all(None);
+			A::remove_all();
 
 			assert_eq!(A::count(), 0);
 			assert_eq!(A::initialize_counter(), 0);
@@ -909,7 +910,7 @@ mod test {
 			assert_eq!(B::count(), 2);
 
 			// Remove all.
-			B::remove_all(None);
+			B::remove_all();
 
 			assert_eq!(B::count(), 0);
 			assert_eq!(B::initialize_counter(), 0);

frame/support/src/storage/types/double_map.rs

@@ -335,7 +335,11 @@ where
 		>(key1, key2)
 	}
 
-	/// Remove all value of the storage.
+	/// Remove all values of the storage in the overlay and up to `limit` in the backend.
+	///
+	/// All values in the client overlay will be deleted, if there is some `limit` then up to
+	/// `limit` values are deleted from the client backend, if `limit` is none then all values in
+	/// the client backend are deleted.
 	pub fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
 		<Self as crate::storage::StoragePrefixedMap<Value>>::remove_all(limit)
 	}

frame/support/src/storage/types/map.rs

@@ -234,7 +234,11 @@ where
 		<Self as crate::storage::StorageMap<Key, Value>>::migrate_key::<OldHasher, _>(key)
 	}
 
-	/// Remove all value of the storage.
+	/// Remove all values of the storage in the overlay and up to `limit` in the backend.
+	///
+	/// All values in the client overlay will be deleted, if there is some `limit` then up to
+	/// `limit` values are deleted from the client backend, if `limit` is none then all values in
+	/// the client backend are deleted.
 	pub fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
 		<Self as crate::storage::StoragePrefixedMap<Value>>::remove_all(limit)
 	}

primitives/io/src/lib.rs

@@ -88,12 +88,12 @@ pub enum EcdsaVerifyError {
 }
 
 /// The outcome of calling `storage_kill`. Returned value is the number of storage items
-/// removed from the trie from making the `storage_kill` call.
+/// removed from the backend from making the `storage_kill` call.
 #[derive(PassByCodec, Encode, Decode)]
 pub enum KillStorageResult {
-	/// No key remains in the child trie.
+	/// All key to remove were removed, return number of key removed from backend.
 	AllRemoved(u32),
-	/// At least one key still resides in the child trie due to the supplied limit.
+	/// Not all key to remove were removed, return number of key removed from backend.
 	SomeRemaining(u32),
 }