address PR comments

pallets/capacity/src/benchmarking.rs

@@ -57,15 +57,17 @@ benchmarks! {
 	withdraw_unstaked {
 		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 5u32);
 		let mut unlocking: UnlockChunks<T> = UnlockChunks::default();
-		assert_ok!(unlocking.add(50u32.into(), 3u32.into()));
-		assert_ok!(unlocking.add(50u32.into(), 5u32.into()));
+		for _i in 0..T::MaxUnlockingChunks::get() {
+			assert_ok!(unlocking.add(1u32.into(), 3u32.into()));
+		}
 		UnstakeUnlocks::<T>::set(&caller, Some(unlocking));
 
 		CurrentEpoch::<T>::set(T::EpochNumber::from(5u32));
 
 	}: _ (RawOrigin::Signed(caller.clone()))
 	verify {
-		assert_last_event::<T>(Event::<T>::StakeWithdrawn {account: caller, amount: 100u32.into() }.into());
+		let total = T::MaxUnlockingChunks::get();
+		assert_last_event::<T>(Event::<T>::StakeWithdrawn {account: caller, amount: total.into() }.into());
 	}
 
 	on_initialize {

pallets/capacity/src/lib.rs

@@ -352,7 +352,6 @@ pub mod pallet {
 		/// in the caller's token account.
 		///
 		/// ### Errors
-		///   - Returns `Error::NotAStakingAccount` if no StakingDetails are found for `origin`.
 		///   - Returns `Error::NoUnstakedTokensAvailable` if the account has no unstaking chunks or none are thawed.
 		#[pallet::call_index(1)]
 		#[pallet::weight(T::WeightInfo::withdraw_unstaked())]
@@ -512,7 +511,7 @@ impl<T: Config> Pallet<T> {
 		CapacityLedger::<T>::insert(target, capacity_details);
 	}
 
-	/// Decrease a staking account's active token and create an unlocking chunk to be thawed at some future block.
+	/// Decrease a staking account's active token.
 	fn decrease_active_staking_balance(
 		unstaker: &T::AccountId,
 		amount: BalanceOf<T>,

pallets/capacity/src/migration/v2.rs

@@ -4,9 +4,9 @@ use crate::{
 };
 use frame_support::{
 	pallet_prelude::{GetStorageVersion, Weight},
-	traits::{OnRuntimeUpgrade, StorageVersion},
-	weights::constants::RocksDbWeight,
+	traits::{OnRuntimeUpgrade, StorageVersion, Get},
 };
+
 const LOG_TARGET: &str = "runtime::capacity";
 
 #[cfg(feature = "try-runtime")]
@@ -45,7 +45,7 @@ pub mod v1 {
 
 /// migrate StakingAccountLedger to use new StakingDetails
 pub fn migrate_to_v2<T: Config>() -> Weight {
-	let on_chain_version = Pallet::<T>::on_chain_storage_version();
+	let on_chain_version = Pallet::<T>::on_chain_storage_version(); // 1r
 
 	if on_chain_version.lt(&2) {
 		log::info!(target: LOG_TARGET, "🔄 StakingAccountLedger migration started");
@@ -65,13 +65,17 @@ pub fn migrate_to_v2<T: Config>() -> Weight {
 				Some(new_account)
 			},
 		);
-		StorageVersion::new(2).put::<Pallet<T>>();
+		StorageVersion::new(2).put::<Pallet<T>>(); // 1 w
+		let reads = (maybe_count + 1) as u64;
+		let writes = (maybe_count *2 + 1) as u64;
 		log::info!(target: LOG_TARGET, "🔄 migration finished");
-		let count = (StakingAccountLedger::<T>::iter().count() + 1) as u64;
-		RocksDbWeight::get().reads_writes(count, count)
+		let weight = T::DbWeight::get().reads_writes(reads, writes);
+		log::info!(target: LOG_TARGET, "Migration calculated weight = {:?}", weight);
+		weight
 	} else {
 		// storage was already migrated.
-		Weight::zero()
+		log::info!(target: LOG_TARGET, "Old StorageAccountLedger migration attempted to run. Please remove");
+		T::DbWeight::get().reads(1)
 	}
 }
 /// The OnRuntimeUpgrade implementation for this storage migration
@@ -103,7 +107,7 @@ impl<T: Config> OnRuntimeUpgrade for MigrateToV2<T> {
 		let on_chain_version = Pallet::<T>::on_chain_storage_version();
 
 		assert_eq!(on_chain_version, crate::pallet::STORAGE_VERSION);
-		assert_eq!(pre_upgrade_count as usize, v2::StakingAccountLedger::<T>::iter().count());
+		assert_eq!(pre_upgrade_count as usize, StakingAccountLedger::<T>::iter().count());
 		assert_eq!(pre_upgrade_count as usize, UnstakeUnlocks::<T>::iter().count());
 
 		log::info!(target: LOG_TARGET, "✅ migration post_upgrade checks passed");

pallets/capacity/src/types.rs

@@ -197,15 +197,15 @@ impl<T: Config> Default for UnlockChunks<T> {
 impl<T: Config> UnlockChunks<T> {
 	#[cfg(any(feature = "runtime-benchmarks", test))]
 	#[allow(clippy::unwrap_used)]
-	///  tmp fn for testing only
-	/// set unlock chunks with (balance, thaw_at).  does not check that the unlock chunks
-	/// don't exceed total.
+	/// set unlock chunks with (balance, thaw_at).  Does not check BoundedVec limit.
 	/// returns true on success, false on failure (?)
+	/// For testing and benchmarks ONLY, note possible panic via BoundedVec::try_from + unwrap
 	pub fn set_unlock_chunks(&mut self, chunks: &Vec<(u32, u32)>) -> bool {
 		let result: Vec<UnlockChunk<BalanceOf<T>, <T>::EpochNumber>> = chunks
 			.into_iter()
 			.map(|chunk| UnlockChunk { value: chunk.0.into(), thaw_at: chunk.1.into() })
 			.collect();
+		// CAUTION
 		self.unlocking = BoundedVec::try_from(result).unwrap();
 		self.unlocking.len() == chunks.len()
 	}
@@ -215,16 +215,14 @@ impl<T: Config> UnlockChunks<T> {
 	/// Returns: the total amount reaped from `unlocking`
 	pub fn reap_thawed(&mut self, current_epoch: <T>::EpochNumber) -> BalanceOf<T> {
 		let mut total_reaped: BalanceOf<T> = 0u32.into();
-		if !self.unlocking.is_empty() {
-			self.unlocking.retain(|chunk| {
-				if current_epoch.ge(&chunk.thaw_at) {
-					total_reaped = total_reaped.saturating_add(chunk.value);
-					false
-				} else {
-					true
-				}
-			});
-		}
+		self.unlocking.retain(|chunk| {
+			if current_epoch.ge(&chunk.thaw_at) {
+				total_reaped = total_reaped.saturating_add(chunk.value);
+				false
+			} else {
+				true
+			}
+		});
 		total_reaped
 	}
 
@@ -245,9 +243,6 @@ impl<T: Config> UnlockChunks<T> {
 
 	/// Get the total balance of all unlock chunks
 	pub fn total(&self) -> BalanceOf<T> {
-		if self.unlocking.is_empty() {
-			return Zero::zero()
-		}
 		self.unlocking
 			.iter()
 			.fold(Zero::zero(), |acc: BalanceOf<T>, chunk| acc.saturating_add(chunk.value))