From 2ec760f709b870c7ed21d919cf8361376b24d31d Mon Sep 17 00:00:00 2001
From: Carlos Lostao <hastro1990@gmail.com>
Date: Wed, 8 Nov 2023 13:02:45 +0100
Subject: [PATCH] fix: coherent extrinsic weights

---
 pallets/root-controller/src/lib.rs          | 10 ++++++-
 pallets/upgrade-runtime-proposal/src/lib.rs | 16 ++++++++---
 pallets/validator-set/src/lib.rs            | 32 ++++++++++++++++++---
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/pallets/root-controller/src/lib.rs b/pallets/root-controller/src/lib.rs
index f213d4b0..1e8b64ef 100644
--- a/pallets/root-controller/src/lib.rs
+++ b/pallets/root-controller/src/lib.rs
@@ -41,10 +41,18 @@ pub mod pallet {
 		DispatchAsRootOccurred { dispatch_result: DispatchResult },
 	}
 
+	impl<T: Config> Pallet<T> {
+		pub fn dispatch_as_root_weight() -> Weight {
+			Weight::from_parts(7_984_000, 0)
+				.saturating_add(Weight::from_parts(0, 1505))
+				.saturating_add(T::DbWeight::get().reads(1))
+		}
+	}
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {
 		#[pallet::call_index(0)]
-		#[pallet::weight(0)]
+		#[pallet::weight(Pallet::<T>::dispatch_as_root_weight())]
 		pub fn dispatch_as_root(
 			origin: OriginFor<T>,
 			call: Box<<T as Config>::RuntimeCall>,
diff --git a/pallets/upgrade-runtime-proposal/src/lib.rs b/pallets/upgrade-runtime-proposal/src/lib.rs
index a0a614a0..ea3d60ff 100644
--- a/pallets/upgrade-runtime-proposal/src/lib.rs
+++ b/pallets/upgrade-runtime-proposal/src/lib.rs
@@ -119,6 +119,14 @@ pub mod pallet {
 		}
 	}
 
+	impl<T: Config> Pallet<T> {
+		fn weight_for_read_writes(reads: u64, writes: u64) -> Weight {
+			Weight::from_parts(21_330_000, 1602)
+				.saturating_add(T::DbWeight::get().reads(reads))
+				.saturating_add(T::DbWeight::get().writes(writes))
+		}
+	}
+
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(n: T::BlockNumber) -> Weight {
@@ -126,17 +134,17 @@ pub mod pallet {
 			let application_block_number_option = <ApplicationBlockNumber<T>>::get();
 
 			if code_saved_option.is_none() {
-				return Weight::zero();
+				return Pallet::<T>::weight_for_read_writes(2, 0);
 			}
 			let code = code_saved_option.unwrap().to_vec();
 
 			if application_block_number_option.is_none() {
-				return Weight::zero();
+				return Pallet::<T>::weight_for_read_writes(2, 0);
 			}
 			let application_block_number = application_block_number_option.unwrap();
 
 			if n != application_block_number {
-				return Weight::zero();
+				return Pallet::<T>::weight_for_read_writes(2, 0);
 			}
 
 			let call = frame_system::Call::<T>::set_code { code: code.into() };
@@ -152,7 +160,7 @@ pub mod pallet {
 			Pallet::<T>::clear_proposed_code();
 			Pallet::<T>::clear_application_block_number();
 
-			Weight::zero()
+			return Pallet::<T>::weight_for_read_writes(2, 2);
 		}
 	}
 
diff --git a/pallets/validator-set/src/lib.rs b/pallets/validator-set/src/lib.rs
index 9b8d2023..4572600b 100644
--- a/pallets/validator-set/src/lib.rs
+++ b/pallets/validator-set/src/lib.rs
@@ -365,6 +365,30 @@ pub mod pallet {
 		pub authority_index: u32,
 	}
 
+	impl<T: Config> Pallet<T> {
+		fn add_validator_weight() -> Weight {
+			Weight::from_parts(21_330_000, 1602)
+				.saturating_add(T::DbWeight::get().reads(1_u64))
+				.saturating_add(T::DbWeight::get().writes(1_u64))
+		}
+
+		fn remove_validator_weight() -> Weight {
+			Weight::from_parts(19_840_000, 1602)
+				.saturating_add(T::DbWeight::get().reads(2_u64))
+				.saturating_add(T::DbWeight::get().writes(2_u64))
+		}
+
+		fn add_validator_again_weight() -> Weight {
+			Weight::from_parts(21_330_000, 1602)
+				.saturating_add(T::DbWeight::get().reads(2_u64))
+				.saturating_add(T::DbWeight::get().writes(1_u64))
+		}
+
+		fn update_max_missed_epochs_weight() -> Weight {
+			Weight::from_parts(21_330_000, 1602).saturating_add(T::DbWeight::get().writes(1_u64))
+		}
+	}
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {
 		/// Add a new validator.
@@ -375,7 +399,7 @@ pub mod pallet {
 		/// The origin can be configured using the `AddRemoveOrigin` type in the
 		/// host runtime. Can also be set to sudo/root.
 		#[pallet::call_index(0)]
-		#[pallet::weight(0)]
+		#[pallet::weight(Pallet::<T>::add_validator_weight())]
 		pub fn add_validator(origin: OriginFor<T>, validator_id: T::AccountId) -> DispatchResult {
 			T::AddRemoveOrigin::ensure_origin(origin)?;
 
@@ -389,7 +413,7 @@ pub mod pallet {
 		/// The origin can be configured using the `AddRemoveOrigin` type in the
 		/// host runtime. Can also be set to sudo/root.
 		#[pallet::call_index(1)]
-		#[pallet::weight(0)]
+		#[pallet::weight(Pallet::<T>::remove_validator_weight())]
 		pub fn remove_validator(
 			origin: OriginFor<T>,
 			validator_id: T::AccountId,
@@ -407,7 +431,7 @@ pub mod pallet {
 		/// The origin can be configured using the `AddRemoveOrigin` type in the
 		/// host runtime. Can also be set to sudo/root.
 		#[pallet::call_index(2)]
-		#[pallet::weight(0)]
+		#[pallet::weight(Pallet::<T>::update_max_missed_epochs_weight())]
 		pub fn update_max_missed_epochs(
 			origin: OriginFor<T>,
 			max_missed_epochs: U256,
@@ -423,7 +447,7 @@ pub mod pallet {
 		///
 		/// For this call, the dispatch origin must be the validator itself.
 		#[pallet::call_index(3)]
-		#[pallet::weight(0)]
+		#[pallet::weight(Pallet::<T>::add_validator_again_weight())]
 		pub fn add_validator_again(
 			origin: OriginFor<T>,
 			heartbeat: Heartbeat<T::BlockNumber, T::AuthorityId>,