refactor aggregator_update_metadata

imstar15 · Mar 3, 2023 · 267b8cc · 267b8cc
1 parent 47c50cf
commit 267b8cc
Show file tree

Hide file tree

Showing 3 changed files with 109 additions and 98 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pallets/parachain-staking/src/lib.rs b/pallets/parachain-staking/src/lib.rs
@@ -110,6 +110,13 @@ impl<D: Decode> FromInfiniteZeros for D {
 	}
 }
 
+#[derive(Eq, PartialEq, Encode, Decode, TypeInfo, Debug, Clone)]
+pub enum MetadataUpdateAction {
+	ExtendApprovedCollators,
+	RemoveApprovedCollators,
+}
+
+
 #[pallet]
 pub mod pallet {
 	pub use super::*;
@@ -2523,91 +2530,36 @@ pub mod pallet {
 			Ok(().into())
 		}
 
-		#[pallet::weight(1_000_000_000)]
+		#[pallet::weight(T::DbWeight::get().reads_writes(20, 20))]
 		#[transactional]
 		pub fn aggregator_update_metadata(
 			origin: OriginFor<T>,
 			collator_candidates: Vec<T::AccountId>,
-			should_be_approved: bool,
+			action: MetadataUpdateAction,
 		) -> DispatchResultWithPostInfo {
 			let aggregator = ensure_signed(origin)?;
 
-			ensure!(
-				!Self::is_candidate(&aggregator),
-				Error::<T>::CandidateExists
-			);
-			ensure!(
-				!Self::is_delegator(&aggregator),
-				Error::<T>::DelegatorExists
-			);
+			ensure!( !Self::is_candidate(&aggregator), Error::<T>::CandidateExists);
+			ensure!( !Self::is_delegator(&aggregator), Error::<T>::DelegatorExists);
 
 			AggregatorMetadata::<T>::try_mutate_exists(
 				aggregator.clone(),
 				|maybe_aggregator_metadata| -> DispatchResult {
 					let mut aggregator_metadata = maybe_aggregator_metadata
 						.take()
-						.unwrap_or(AggregatorMetadataType::<T::AccountId>::default());
+						.unwrap_or_default();
 
-					match should_be_approved {
-						true => {
-							collator_candidates.iter().try_for_each(
-								|collator_candidate| -> DispatchResult {
-									ensure!(
-										Self::is_candidate(&collator_candidate),
-										Error::<T>::CandidateDNE
-									);
-									ensure!(
-										aggregator_metadata
-											.approved_candidates
-											.insert(collator_candidate.clone()),
-										Error::<T>::CandidateAlreadyApprovedByAggregator
-									);
-									Ok(())
-								},
+					match action {
+						MetadataUpdateAction::ExtendApprovedCollators => {
+							collator_candidates.iter().try_for_each(|collator| -> DispatchResult {
+								Self::add_approved_candidate_for_collator_metadata(collator, &mut aggregator_metadata)
+							}
 							)?;
 						}
-						false => {
+						MetadataUpdateAction::RemoveApprovedCollators => {
 							collator_candidates.iter().try_for_each(
-								|collator_candidate| -> DispatchResult {
-									// Do not propagate the error if there's an error here
-									// Then it means that the aggregator wasn't aggregating for this candidate
-									// Or that the target is no longer a candidate, which can happen if the candidate has left
-									// removing all his aggregation details except for approvals from various aggregators
-									let _ =
-										CandidateAggregator::<T>::try_mutate(
-											|candidate_aggregator_map| -> DispatchResult {
-												let collator_candidate_state =
-													<CandidateState<T>>::get(&collator_candidate)
-														.ok_or(Error::<T>::CandidateDNE)?;
-												ensure!(Some(aggregator.clone()) == candidate_aggregator_map.remove(collator_candidate),
-											Error::<T>::CandidateNotAggregatingUnderAggregator);
-
-												// If CandidateAggregator has the aggregator listed under this candidate then
-												// the aggregator metadata will have this candidate listed under its liquidity token
-												let res =
-													aggregator_metadata.token_collator_map.remove(
-														&collator_candidate_state.liquidity_token,
-													);
-												if res != Some(collator_candidate.clone()) {
-													log::error!(
-												"Inconsistent aggregator metadata: candidate - {:?}, aggregator - {:?}",
-												collator_candidate,
-												aggregator
-											);
-												}
-
-												Ok(())
-											},
-										);
-
-									ensure!(
-										aggregator_metadata
-											.approved_candidates
-											.remove(collator_candidate),
-										Error::<T>::CandidateNotApprovedByAggregator
-									);
-
-									Ok(())
+								|collator| -> DispatchResult {
+									Self::remove_approved_candidates_from_collator_metadata(collator, &aggregator, &mut aggregator_metadata)
 								},
 							)?;
 						}
@@ -2625,8 +2577,8 @@ pub mod pallet {
 			Ok(().into())
 		}
 
-		// Benchmark on the worst case max candidate count
-		#[pallet::weight(1_000_000_000)]
+		// TODO: use more precise benchmark
+		#[pallet::weight(T::DbWeight::get().reads_writes(20, 20))]
 		#[transactional]
 		pub fn update_candidate_aggregator(
 			origin: OriginFor<T>,
@@ -2643,7 +2595,8 @@ pub mod pallet {
 			Ok(().into())
 		}
 
-		#[pallet::weight(1_000_000_000)]
+		// TODO: use more precise benchmark
+		#[pallet::weight(T::DbWeight::get().reads_writes(20, 20))]
 		#[transactional]
 		pub fn payout_collator_rewards(
 			origin: OriginFor<T>,
@@ -2678,7 +2631,8 @@ pub mod pallet {
 			Ok(().into())
 		}
 
-		#[pallet::weight(1_000_000_000)]
+		// TODO: use more precise benchmark
+		#[pallet::weight(T::DbWeight::get().reads_writes(20, 20))]
 		#[transactional]
 		pub fn payout_delegator_reward(
 			origin: OriginFor<T>,
@@ -2709,6 +2663,62 @@ pub mod pallet {
 	}
 
 	impl<T: Config> Pallet<T> {
+		fn add_approved_candidate_for_collator_metadata(collator_candidate: &T::AccountId, aggregator_metadata: &mut AggregatorMetadataType<T::AccountId>) -> DispatchResult {
+			ensure!(
+				Self::is_candidate(&collator_candidate),
+				Error::<T>::CandidateDNE
+				);
+			ensure!(
+				aggregator_metadata
+				.approved_candidates
+				.insert(collator_candidate.clone()),
+				Error::<T>::CandidateAlreadyApprovedByAggregator
+				);
+			Ok(())
+		}
+
+		fn remove_approved_candidates_from_collator_metadata(collator_candidate: &T::AccountId, aggregator: &T::AccountId, aggregator_metadata: &mut AggregatorMetadataType<T::AccountId>) -> DispatchResult {
+									// Do not propagate the error if there's an error here
+									// Then it means that the aggregator wasn't aggregating for this candidate
+									// Or that the target is no longer a candidate, which can happen if the candidate has left
+									// removing all his aggregation details except for approvals from various aggregators
+									let _ =
+										CandidateAggregator::<T>::try_mutate(
+											|candidate_aggregator_map| -> DispatchResult {
+												let collator_candidate_state =
+													<CandidateState<T>>::get(&collator_candidate)
+														.ok_or(Error::<T>::CandidateDNE)?;
+												ensure!(Some(aggregator.clone()) == candidate_aggregator_map.remove(collator_candidate),
+											Error::<T>::CandidateNotAggregatingUnderAggregator);
+
+												// If CandidateAggregator has the aggregator listed under this candidate then
+												// the aggregator metadata will have this candidate listed under its liquidity token
+												let res =
+													aggregator_metadata.token_collator_map.remove(
+														&collator_candidate_state.liquidity_token,
+													);
+												if res != Some(collator_candidate.clone()) {
+													log::error!(
+												"Inconsistent aggregator metadata: candidate - {:?}, aggregator - {:?}",
+												collator_candidate,
+												aggregator
+											);
+												}
+
+												Ok(())
+											},
+										);
+
+									ensure!(
+										aggregator_metadata
+											.approved_candidates
+											.remove(collator_candidate),
+										Error::<T>::CandidateNotApprovedByAggregator
+									);
+
+			Ok(())
+		}
+
 		pub fn payout_reward(to: T::AccountId, amt: Balance) -> DispatchResult {
 			let _ = <T as pallet::Config>::Currency::transfer(
 				T::NativeTokenId::get().into(),

diff --git a/pallets/parachain-staking/src/tests.rs b/pallets/parachain-staking/src/tests.rs
@@ -29,6 +29,7 @@ use crate::mock::{
 
 
 use crate::{
+MetadataUpdateAction,
 	assert_eq_events, assert_event_emitted, assert_last_event, Balance, Bond, CandidateBondChange,
 	CandidateBondRequest, CollatorStatus, DelegationChange, DelegationRequest, DelegatorAdded,
 	Error, Event, PairedOrLiquidityToken,RoundAggregatorInfo, RoundCollatorRewardInfo,
@@ -5322,7 +5323,7 @@ fn token_valuations_works_with_aggregators() {
 			assert_ok!(Stake::aggregator_update_metadata(
 				Origin::signed(4),
 				vec![1, 2],
-				true
+				MetadataUpdateAction::ExtendApprovedCollators
 			));
 			assert_ok!(Stake::update_candidate_aggregator(
 				Origin::signed(1),
@@ -5434,7 +5435,7 @@ fn round_aggregator_info_is_updated() {
 			assert_ok!(Stake::aggregator_update_metadata(
 				Origin::signed(4),
 				vec![1, 2],
-				true
+				MetadataUpdateAction::ExtendApprovedCollators
 			));
 			assert_ok!(Stake::update_candidate_aggregator(
 				Origin::signed(1),
@@ -5518,7 +5519,7 @@ fn payouts_with_aggregators_work() {
 			assert_ok!(Stake::aggregator_update_metadata(
 				Origin::signed(4),
 				vec![1, 2],
-				true
+				MetadataUpdateAction::ExtendApprovedCollators
 			));
 			assert_ok!(Stake::update_candidate_aggregator(
 				Origin::signed(1),