Update reward pool on provider_boost or unstake #1699

pallets/capacity/src/benchmarking.rs

@@ -37,13 +37,47 @@ pub fn set_up_epoch<T: Config>(current_block: BlockNumberFor<T>, current_epoch:
 	CurrentEpochInfo::<T>::set(EpochInfo { epoch_start });
 }
 
+// caller stakes the given amount to the given target
+pub fn setup_provider_stake<T: Config>(
+	caller: &T::AccountId,
+	target: &MessageSourceId,
+	staking_amount: BalanceOf<T>,
+) {
+	let capacity_amount: BalanceOf<T> = Capacity::<T>::capacity_generated(staking_amount);
+
+	let mut staking_account = StakingDetails::<T>::default();
+	let mut target_details = StakingTargetDetails::<BalanceOf<T>>::default();
+	let mut capacity_details =
+		CapacityDetails::<BalanceOf<T>, <T as Config>::EpochNumber>::default();
+
+	staking_account.deposit(staking_amount);
+	target_details.deposit(staking_amount, capacity_amount);
+	capacity_details.deposit(&staking_amount, &capacity_amount);
+
+	Capacity::<T>::set_staking_account_and_lock(caller, &staking_account)
+		.expect("Failed to set staking account");
+	Capacity::<T>::set_target_details_for(caller, *target, target_details);
+	Capacity::<T>::set_capacity_for(*target, capacity_details);
+}
+
+// fill up unlock chunks to max bound - 1
+fn fill_unlock_chunks<T: Config>(caller: &T::AccountId, count: u32) {
+	let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
+	for _i in 0..count {
+		let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> =
+			UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
+		assert_ok!(unlocking.try_push(unlock_chunk));
+	}
+	UnstakeUnlocks::<T>::set(caller, Some(unlocking));
+}
+
 benchmarks! {
 	stake {
 		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 105u32);
 		let amount: BalanceOf<T> = T::MinimumStakingAmount::get();
 		let capacity: BalanceOf<T> = Capacity::<T>::capacity_generated(amount);
 		let target = 1;
-		let staking_type = StakingType::MaximumCapacity;
+		let staking_type = MaximumCapacity;
 
 		register_provider::<T>(target, "Foo");
 
@@ -57,12 +91,7 @@ benchmarks! {
 
 	withdraw_unstaked {
 		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 5u32);
-		let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
-		for _i in 0..T::MaxUnlockingChunks::get() {
-			let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> = UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
-			assert_ok!(unlocking.try_push(unlock_chunk));
-		}
-		UnstakeUnlocks::<T>::set(&caller, Some(unlocking));
+		fill_unlock_chunks::<T>(&caller, T::MaxUnlockingChunks::get());
 
 		CurrentEpoch::<T>::set(T::EpochNumber::from(5u32));
 
@@ -90,31 +119,16 @@ benchmarks! {
 		let target = 1;
 		let block_number = 4u32;
 
-		let mut staking_account = StakingDetails::<T>::default();
-		let mut target_details = StakingTargetDetails::<BalanceOf<T>>::default();
-		let mut capacity_details = CapacityDetails::<BalanceOf<T>, <T as Config>::EpochNumber>::default();
-
-		staking_account.deposit(staking_amount);
-		target_details.deposit(staking_amount, capacity_amount);
-		capacity_details.deposit(&staking_amount, &capacity_amount);
-
-		Capacity::<T>::set_staking_account_and_lock(&caller.clone(), &staking_account).expect("Failed to set staking account");
-		Capacity::<T>::set_target_details_for(&caller.clone(), target, target_details);
-		Capacity::<T>::set_capacity_for(target, capacity_details);
-
-		// fill up unlock chunks to max bound - 1
-		let count = T::MaxUnlockingChunks::get()-1;
-		let mut unlocking: UnlockChunkList<T> = BoundedVec::default();
-		for _i in 0..count {
-			let unlock_chunk: UnlockChunk<BalanceOf<T>, T::EpochNumber> = UnlockChunk { value: 1u32.into(), thaw_at: 3u32.into() };
-			assert_ok!(unlocking.try_push(unlock_chunk));
-		}
-		UnstakeUnlocks::<T>::set(&caller, Some(unlocking));
-
-
+		setup_provider_stake::<T>(&caller, &target, staking_amount);
+		fill_unlock_chunks::<T>(&caller, T::MaxUnlockingChunks::get() - 1);
 	}: _ (RawOrigin::Signed(caller.clone()), target, unstaking_amount.into())
 	verify {
-		assert_last_event::<T>(Event::<T>::UnStaked {account: caller, target: target, amount: unstaking_amount.into(), capacity: Capacity::<T>::calculate_capacity_reduction(unstaking_amount.into(), staking_amount, capacity_amount) }.into());
+		assert_last_event::<T>(Event::<T>::UnStaked {
+			account: caller.clone(),
+			target,
+			amount: unstaking_amount.into(),
+			capacity: Capacity::<T>::calculate_capacity_reduction(unstaking_amount.into(), staking_amount,capacity_amount)
+		}.into());
 	}
 
 	set_epoch_length {
@@ -125,7 +139,47 @@ benchmarks! {
 		assert_last_event::<T>(Event::<T>::EpochLengthUpdated {blocks: epoch_length}.into());
 	}
 
+	change_staking_target {
+		let caller: T::AccountId = create_funded_account::<T>("account", SEED, 5u32);
+		let from_msa = 33;
+		let to_msa = 34;
+		// amount in addition to minimum
+		let from_msa_amount = 32u32;
+		let to_msa_amount = 1u32;
+
+		register_provider::<T>(from_msa, "frommsa");
+		register_provider::<T>(to_msa, "tomsa");
+		setup_provider_stake::<T>(&caller, &from_msa, from_msa_amount.into());
+		setup_provider_stake::<T>(&caller, &to_msa, to_msa_amount.into());
+		let restake_amount = 11u32;
+
+	}: _ (RawOrigin::Signed(caller.clone(), ), from_msa, to_msa, restake_amount.into())
+	verify {
+		assert_last_event::<T>(Event::<T>::StakingTargetChanged {
+			account: caller,
+			from_msa,
+			to_msa,
+			amount: restake_amount.into()
+		}.into());
+	}
+
+	provider_boost {
+		let caller: T::AccountId = create_funded_account::<T>("boostaccount", SEED, 260u32);
+		let boost_amount: BalanceOf<T> = 200u32.into(); // enough for 1 Cap boosted
+		let capacity: BalanceOf<T> = Capacity::<T>::capacity_generated(<T>::RewardsProvider::capacity_boost(boost_amount));
+		let target = 1;
+
+		register_provider::<T>(target, "Foo");
+
+	}: _ (RawOrigin::Signed(caller.clone()), target, boost_amount)
+	verify {
+		assert!(StakingAccountLedger::<T>::contains_key(&caller));
+		assert!(StakingTargetLedger::<T>::contains_key(&caller, target));
+		assert!(CapacityLedger::<T>::contains_key(target));
+		assert_last_event::<T>(Event::<T>::ProviderBoosted {account: caller, amount: boost_amount, target, capacity}.into());
+	}
+
 	impl_benchmark_test_suite!(Capacity,
-		crate::tests::mock::new_test_ext(),
-		crate::tests::mock::Test);
+		tests::mock::new_test_ext(),
+		tests::mock::Test);
 }

pallets/capacity/src/lib.rs

@@ -90,8 +90,8 @@ pub mod weights;
 pub(crate) type BalanceOf<T> =
 	<<T as Config>::Currency as InspectFungible<<T as frame_system::Config>::AccountId>>::Balance;
 
-use frame_system::pallet_prelude::*;
 use crate::StakingType::{MaximumCapacity, ProviderBoost};
+use frame_system::pallet_prelude::*;
 
 #[frame_support::pallet]
 pub mod pallet {
@@ -269,13 +269,13 @@ pub mod pallet {
 	#[pallet::whitelist_storage]
 	#[pallet::getter(fn get_current_era)]
 	pub type CurrentEraInfo<T: Config> =
-	StorageValue<_, RewardEraInfo<T::RewardEra, BlockNumberFor<T>>, ValueQuery>;
+		StorageValue<_, RewardEraInfo<T::RewardEra, BlockNumberFor<T>>, ValueQuery>;
 
 	/// Reward Pool history
 	#[pallet::storage]
 	#[pallet::getter(fn get_reward_pool_for_era)]
 	pub type StakingRewardPool<T: Config> =
-	CountedStorageMap<_, Twox64Concat, T::RewardEra, RewardPoolInfo<BalanceOf<T>>>;
+		CountedStorageMap<_, Twox64Concat, T::RewardEra, RewardPoolInfo<BalanceOf<T>>>;
 
 	// TODO: storage for staking history
 
@@ -405,16 +405,14 @@ pub mod pallet {
 		/// Too many change_staking_target calls made in this RewardEra. (20)
 		MaxRetargetsExceeded,
 		/// There are no unstaked token amounts that have passed their thaw period.
-		NoThawedTokenAvailable
+		NoThawedTokenAvailable,
 	}
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(current: BlockNumberFor<T>) -> Weight {
 			Self::start_new_epoch_if_needed(current)
-				.saturating_add(
-					Self::start_new_reward_era_if_needed(current)
-			)
+				.saturating_add(Self::start_new_reward_era_if_needed(current))
 		}
 	}
 
@@ -494,10 +492,12 @@ pub mod pallet {
 
 			ensure!(requested_amount > Zero::zero(), Error::<T>::UnstakedAmountIsZero);
 
-			let (actual_amount, staking_type) = Self::decrease_active_staking_balance(&unstaker, requested_amount)?;
+			let (actual_amount, staking_type) =
+				Self::decrease_active_staking_balance(&unstaker, requested_amount)?;
 			Self::add_unlock_chunk(&unstaker, actual_amount)?;
 
-			let capacity_reduction = Self::reduce_capacity(&unstaker, target, actual_amount, staking_type)?;
+			let capacity_reduction =
+				Self::reduce_capacity(&unstaker, target, actual_amount, staking_type)?;
 
 			Self::deposit_event(Event::UnStaked {
 				account: unstaker,
@@ -528,7 +528,7 @@ pub mod pallet {
 		}
 
 		/// Sets the target of the staking capacity to a new target.
-		/// This adds a chunk to `StakingAccountDetails.stake_change_unlocking chunks`, up to `T::MaxUnlockingChunks`.
+		/// This adds a chunk to `StakingDetails.stake_change_unlocking chunks`, up to `T::MaxUnlockingChunks`.
 		/// The staked amount and Capacity generated by `amount` originally targeted to the `from` MSA Id is reassigned to the `to` MSA Id.
 		/// Does not affect unstaking process or additional stake amounts.
 		/// Changing a staking target to a Provider when Origin has nothing staked them will retain the staking type.
@@ -603,7 +603,6 @@ pub mod pallet {
 
 			Ok(())
 		}
-
 	}
 }
 
@@ -627,7 +626,10 @@ impl<T: Config> Pallet<T> {
 
 		let staking_details = Self::get_staking_account_for(&staker).unwrap_or_default();
 		if !staking_details.active.is_zero() {
-			ensure!(staking_details.staking_type.eq(&staking_type), Error::<T>::CannotChangeStakingType);
+			ensure!(
+				staking_details.staking_type.eq(&staking_type),
+				Error::<T>::CannotChangeStakingType
+			);
 		}
 
 		let stakable_amount = Self::get_stakable_amount_for(&staker, amount);
@@ -653,7 +655,7 @@ impl<T: Config> Pallet<T> {
 		amount: &BalanceOf<T>,
 	) -> Result<(StakingDetails<T>, BalanceOf<T>), DispatchError> {
 		// FIXME: if one is boosting additional amounts, this will fail.
-		let (mut staking_details, stakable_amount) = 
+		let (mut staking_details, stakable_amount) =
 			Self::ensure_can_stake(staker, *target, *amount, ProviderBoost)?;
 		staking_details.staking_type = ProviderBoost;
 		// TODO: update when boost history is implemented fully
@@ -692,9 +694,7 @@ impl<T: Config> Pallet<T> {
 		target: &MessageSourceId,
 		amount: &BalanceOf<T>,
 	) -> Result<BalanceOf<T>, DispatchError> {
-		staking_details
-			.deposit(*amount)
-			.ok_or(ArithmeticError::Overflow)?;
+		staking_details.deposit(*amount).ok_or(ArithmeticError::Overflow)?;
 
 		// get the capacity generated by a Provider Boost
 		let capacity = Self::capacity_generated(T::RewardsProvider::capacity_boost(*amount));
@@ -706,12 +706,16 @@ impl<T: Config> Pallet<T> {
 		let mut capacity_details = Self::get_capacity_for(target).unwrap_or_default();
 		capacity_details.deposit(amount, &capacity).ok_or(ArithmeticError::Overflow)?;
 
+		let era = Self::get_current_era().era_index;
+		let mut reward_pool =
+			Self::get_reward_pool_for_era(era).ok_or(Error::<T>::EraOutOfRange)?;
+		reward_pool.total_staked_token = reward_pool.total_staked_token.saturating_add(*amount);
+
 		// TODO:  add boost history record for era when boost history is implemented
-		// let era = Self::get_current_era().era_index;
 		Self::set_staking_account_and_lock(staker, staking_details)?;
 		Self::set_target_details_for(staker, *target, target_details);
 		Self::set_capacity_for(*target, capacity_details);
-
+		Self::set_reward_pool(era, &reward_pool);
 		Ok(capacity)
 	}
 
@@ -767,6 +771,10 @@ impl<T: Config> Pallet<T> {
 		CapacityLedger::<T>::insert(target, capacity_details);
 	}
 
+	fn set_reward_pool(era: <T>::RewardEra, new_reward_pool: &RewardPoolInfo<BalanceOf<T>>) {
+		StakingRewardPool::<T>::set(era, Some(new_reward_pool.clone()));
+	}
+
 	/// Decrease a staking account's active token and reap if it goes below the minimum.
 	/// Returns: actual amount unstaked, plus the staking type + StakingDetails,
 	/// since StakingDetails may be reaped and staking type must be used to calculate the
@@ -782,6 +790,12 @@ impl<T: Config> Pallet<T> {
 		let actual_unstaked_amount = staking_account.withdraw(amount)?;
 		let staking_type = staking_account.staking_type;
 		Self::set_staking_account(unstaker, &staking_account);
+
+		let era = Self::get_current_era().era_index;
+		let mut reward_pool =
+			Self::get_reward_pool_for_era(era).ok_or(Error::<T>::EraOutOfRange)?;
+		reward_pool.total_staked_token = reward_pool.total_staked_token.saturating_sub(amount);
+		Self::set_reward_pool(era, &reward_pool.clone());
 		Ok((actual_unstaked_amount, staking_type))
 	}
 
@@ -843,6 +857,7 @@ impl<T: Config> Pallet<T> {
 		Ok(amount_withdrawn)
 	}
 
+	#[allow(unused)]
 	fn get_thaw_at_epoch() -> <T as Config>::EpochNumber {
 		let current_epoch: T::EpochNumber = Self::get_current_epoch();
 		let thaw_period = T::UnstakingThawPeriod::get();
@@ -854,7 +869,7 @@ impl<T: Config> Pallet<T> {
 		unstaker: &T::AccountId,
 		target: MessageSourceId,
 		amount: BalanceOf<T>,
-		staking_type: StakingType
+		staking_type: StakingType,
 	) -> Result<BalanceOf<T>, DispatchError> {
 		let mut staking_target_details = Self::get_target_for(&unstaker, &target)
 			.ok_or(Error::<T>::StakerTargetRelationshipNotFound)?;
@@ -864,17 +879,16 @@ impl<T: Config> Pallet<T> {
 		let mut capacity_details =
 			Self::get_capacity_for(target).ok_or(Error::<T>::TargetCapacityNotFound)?;
 
-		let capacity_to_withdraw =
-			if staking_type.eq(&StakingType::ProviderBoost) {
-				Perbill::from_rational(amount, staking_target_details.amount)
-					.mul_ceil(staking_target_details.capacity)
-			} else {
-				Self::calculate_capacity_reduction(
-					amount,
-					capacity_details.total_tokens_staked,
-					capacity_details.total_capacity_issued,
-				)
-			};
+		let capacity_to_withdraw = if staking_type.eq(&StakingType::ProviderBoost) {
+			Perbill::from_rational(amount, staking_target_details.amount)
+				.mul_ceil(staking_target_details.capacity)
+		} else {
+			Self::calculate_capacity_reduction(
+				amount,
+				capacity_details.total_tokens_staked,
+				capacity_details.total_capacity_issued,
+			)
+		};
 		// this call will return an amount > than requested if the resulting StakingTargetDetails balance
 		// is below the minimum. This ensures we withdraw the same amounts as for staking_target_details.
 		// TODO: update this function
@@ -926,15 +940,16 @@ impl<T: Config> Pallet<T> {
 	}
 
 	fn start_new_reward_era_if_needed(current_block: BlockNumberFor<T>) -> Weight {
-		let current_era_info: RewardEraInfo<T::RewardEra, BlockNumberFor<T>> = Self::get_current_era(); // 1r
+		let current_era_info: RewardEraInfo<T::RewardEra, BlockNumberFor<T>> =
+			Self::get_current_era(); // 1r
 
 		if current_block.saturating_sub(current_era_info.started_at) >= T::EraLength::get().into() {
 			let new_era_info = RewardEraInfo {
 				era_index: current_era_info.era_index.saturating_add(One::one()),
 				started_at: current_block,
 			};
 
-			let current_reward_pool_info =
+			let current_reward_pool =
 				Self::get_reward_pool_for_era(current_era_info.era_index).unwrap_or_default(); // 1r
 
 			let past_eras_max = T::StakingRewardsPastErasMax::get();
@@ -948,9 +963,9 @@ impl<T: Config> Pallet<T> {
 			CurrentEraInfo::<T>::set(new_era_info); // 1w
 
 			let total_reward_pool =
-				T::RewardsProvider::reward_pool_size(current_reward_pool_info.total_staked_token);
+				T::RewardsProvider::reward_pool_size(current_reward_pool.total_staked_token);
 			let new_reward_pool = RewardPoolInfo {
-				total_staked_token: current_reward_pool_info.total_staked_token,
+				total_staked_token: current_reward_pool.total_staked_token,
 				total_reward_pool,
 				unclaimed_balance: total_reward_pool,
 			};
@@ -1039,7 +1054,11 @@ impl<T: Config> Nontransferable for Pallet<T> {
 	}
 
 	/// Increase all totals for the MSA's CapacityDetails.
-	fn deposit(msa_id: MessageSourceId, token_amount: Self::Balance, capacity_amount: Self::Balance) -> Result<(), DispatchError> {
+	fn deposit(
+		msa_id: MessageSourceId,
+		token_amount: Self::Balance,
+		capacity_amount: Self::Balance,
+	) -> Result<(), DispatchError> {
 		let mut capacity_details =
 			Self::get_capacity_for(msa_id).ok_or(Error::<T>::TargetCapacityNotFound)?;
 		capacity_details.deposit(&token_amount, &capacity_amount);