Contracts: Refactor API to use WeightMeter (#2943)

Update the Contracts API to use `WeightMeter`, as it simplifies the code
and makes it easier to reason about, rather than taking a mutable weight
or returning a tuple with the weight consumed

---------

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

substrate/frame/contracts/src/benchmarking/mod.rs

@@ -40,7 +40,7 @@ use frame_support::{
 	self,
 	pallet_prelude::StorageVersion,
 	traits::{fungible::InspectHold, Currency},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 };
 use frame_system::RawOrigin;
 use pallet_balances;
@@ -198,7 +198,7 @@ mod benchmarks {
 	fn on_process_deletion_queue_batch() {
 		#[block]
 		{
-			ContractInfo::<T>::process_deletion_queue_batch(Weight::MAX);
+			ContractInfo::<T>::process_deletion_queue_batch(&mut WeightMeter::new())
 		}
 	}
 
@@ -213,7 +213,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			ContractInfo::<T>::process_deletion_queue_batch(Weight::MAX);
+			ContractInfo::<T>::process_deletion_queue_batch(&mut WeightMeter::new())
 		}
 
 		Ok(())
@@ -226,7 +226,7 @@ mod benchmarks {
 		let mut m = v09::Migration::<T>::default();
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -244,7 +244,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 
 		Ok(())
@@ -259,7 +259,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -276,7 +276,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -291,7 +291,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 		Ok(())
 	}
@@ -307,7 +307,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -322,7 +322,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 
 		Ok(())
@@ -335,7 +335,7 @@ mod benchmarks {
 		StorageVersion::new(version).put::<Pallet<T>>();
 		#[block]
 		{
-			Migration::<T>::migrate(Weight::MAX);
+			Migration::<T>::migrate(&mut WeightMeter::new());
 		}
 		assert_eq!(StorageVersion::get::<Pallet<T>>(), version);
 	}
@@ -2776,7 +2776,8 @@ mod benchmarks {
 	#[benchmark(extra, pov_mode = Ignored)]
 	fn print_schedule() -> Result<(), BenchmarkError> {
 		let max_weight = <T as frame_system::Config>::BlockWeights::get().max_block;
-		let (weight_per_key, key_budget) = ContractInfo::<T>::deletion_budget(max_weight);
+		let (weight_per_key, key_budget) =
+			ContractInfo::<T>::deletion_budget(&mut WeightMeter::with_limit(max_weight));
 		let schedule = T::Schedule::get();
 		log::info!(target: LOG_TARGET, "
 		{schedule:#?}

substrate/frame/contracts/src/lib.rs

@@ -123,7 +123,7 @@ use frame_support::{
 		fungible::{Inspect, Mutate, MutateHold},
 		ConstU32, Contains, Get, Randomness, Time,
 	},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 	BoundedVec, DefaultNoBound, RuntimeDebugNoBound,
 };
 use frame_system::{
@@ -461,17 +461,15 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_idle(_block: BlockNumberFor<T>, mut remaining_weight: Weight) -> Weight {
+		fn on_idle(_block: BlockNumberFor<T>, limit: Weight) -> Weight {
 			use migration::MigrateResult::*;
+			let mut meter = WeightMeter::with_limit(limit);
 
 			loop {
-				let (result, weight) = Migration::<T>::migrate(remaining_weight);
-				remaining_weight.saturating_reduce(weight);
-
-				match result {
-					// There is not enough weight to perform a migration, or make any progress, we
-					// just return the remaining weight.
-					NoMigrationPerformed | InProgress { steps_done: 0 } => return remaining_weight,
+				match Migration::<T>::migrate(&mut meter) {
+					// There is not enough weight to perform a migration.
+					// We can't do anything more, so we return the used weight.
+					NoMigrationPerformed | InProgress { steps_done: 0 } => return meter.consumed(),
 					// Migration is still in progress, we can start the next step.
 					InProgress { .. } => continue,
 					// Either no migration is in progress, or we are done with all migrations, we
@@ -480,8 +478,8 @@ pub mod pallet {
 				}
 			}
 
-			ContractInfo::<T>::process_deletion_queue_batch(remaining_weight)
-				.saturating_add(T::WeightInfo::on_process_deletion_queue_batch())
+			ContractInfo::<T>::process_deletion_queue_batch(&mut meter);
+			meter.consumed()
 		}
 
 		fn integrity_test() {
@@ -924,18 +922,25 @@ pub mod pallet {
 			ensure_signed(origin)?;
 
 			let weight_limit = weight_limit.saturating_add(T::WeightInfo::migrate());
-			let (result, weight) = Migration::<T>::migrate(weight_limit);
+			let mut meter = WeightMeter::with_limit(weight_limit);
+			let result = Migration::<T>::migrate(&mut meter);
 
 			match result {
-				Completed =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::No }),
-				InProgress { steps_done, .. } if steps_done > 0 =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::No }),
-				InProgress { .. } =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::Yes }),
+				Completed => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::No,
+				}),
+				InProgress { steps_done, .. } if steps_done > 0 => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::No,
+				}),
+				InProgress { .. } => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::Yes,
+				}),
 				NoMigrationInProgress | NoMigrationPerformed => {
 					let err: DispatchError = <Error<T>>::NoMigrationPerformed.into();
-					Err(err.with_weight(T::WeightInfo::migrate()))
+					Err(err.with_weight(meter.consumed()))
 				},
 			}
 		}

substrate/frame/contracts/src/migration.rs

@@ -71,6 +71,7 @@ use codec::{Codec, Decode};
 use frame_support::{
 	pallet_prelude::*,
 	traits::{ConstU32, OnRuntimeUpgrade},
+	weights::WeightMeter,
 };
 use sp_runtime::Saturating;
 use sp_std::marker::PhantomData;
@@ -112,8 +113,8 @@ pub trait MigrationStep: Codec + MaxEncodedLen + Default {
 
 	/// Process one step of the migration.
 	///
-	/// Returns whether the migration is finished and the weight consumed.
-	fn step(&mut self) -> (IsFinished, Weight);
+	/// Returns whether the migration is finished.
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished;
 
 	/// Verify that the migration step fits into `Cursor`, and that `max_step_weight` is not greater
 	/// than `max_block_weight`.
@@ -161,9 +162,9 @@ impl<const N: u16> MigrationStep for NoopMigration<N> {
 	fn max_step_weight() -> Weight {
 		Weight::zero()
 	}
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, _meter: &mut WeightMeter) -> IsFinished {
 		log::debug!(target: LOG_TARGET, "Noop migration for version {}", N);
-		(IsFinished::Yes, Weight::zero())
+		IsFinished::Yes
 	}
 }
 
@@ -209,8 +210,8 @@ pub trait MigrateSequence: private::Sealed {
 		Ok(())
 	}
 
-	/// Execute the migration step until the weight limit is reached.
-	fn steps(version: StorageVersion, cursor: &[u8], weight_left: &mut Weight) -> StepResult;
+	/// Execute the migration step until the available weight is consumed.
+	fn steps(version: StorageVersion, cursor: &[u8], meter: &mut WeightMeter) -> StepResult;
 
 	/// Verify that the migration step fits into `Cursor`, and that `max_step_weight` is not greater
 	/// than `max_block_weight`.
@@ -235,18 +236,18 @@ pub struct Migration<T: Config, const TEST_ALL_STEPS: bool = true>(PhantomData<T
 #[cfg(feature = "try-runtime")]
 impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
 	fn run_all_steps() -> Result<(), TryRuntimeError> {
-		let mut weight = Weight::zero();
+		let mut meter = &mut WeightMeter::new();
 		let name = <Pallet<T>>::name();
 		loop {
 			let in_progress_version = <Pallet<T>>::on_chain_storage_version() + 1;
 			let state = T::Migrations::pre_upgrade_step(in_progress_version)?;
-			let (status, w) = Self::migrate(Weight::MAX);
-			weight.saturating_accrue(w);
+			let before = meter.consumed();
+			let status = Self::migrate(&mut meter);
 			log::info!(
 				target: LOG_TARGET,
 				"{name}: Migration step {:?} weight = {}",
 				in_progress_version,
-				weight
+				meter.consumed() - before
 			);
 			T::Migrations::post_upgrade_step(in_progress_version, state)?;
 			if matches!(status, MigrateResult::Completed) {
@@ -255,7 +256,7 @@ impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
 		}
 
 		let name = <Pallet<T>>::name();
-		log::info!(target: LOG_TARGET, "{name}: Migration steps weight = {}", weight);
+		log::info!(target: LOG_TARGET, "{name}: Migration steps weight = {}", meter.consumed());
 		Ok(())
 	}
 }
@@ -384,19 +385,19 @@ impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
 		T::Migrations::integrity_test(max_weight)
 	}
 
-	/// Migrate
-	/// Return the weight used and whether or not a migration is in progress
-	pub(crate) fn migrate(weight_limit: Weight) -> (MigrateResult, Weight) {
+	/// Execute the multi-step migration.
+	/// Returns whether or not a migration is in progress
+	pub(crate) fn migrate(mut meter: &mut WeightMeter) -> MigrateResult {
 		let name = <Pallet<T>>::name();
-		let mut weight_left = weight_limit;
 
-		if weight_left.checked_reduce(T::WeightInfo::migrate()).is_none() {
-			return (MigrateResult::NoMigrationPerformed, Weight::zero())
+		if meter.try_consume(T::WeightInfo::migrate()).is_err() {
+			return MigrateResult::NoMigrationPerformed
 		}
 
 		MigrationInProgress::<T>::mutate_exists(|progress| {
 			let Some(cursor_before) = progress.as_mut() else {
-				return (MigrateResult::NoMigrationInProgress, T::WeightInfo::migration_noop())
+				meter.consume(T::WeightInfo::migration_noop());
+				return MigrateResult::NoMigrationInProgress
 			};
 
 			// if a migration is running it is always upgrading to the next version
@@ -410,38 +411,36 @@ impl<T: Config, const TEST_ALL_STEPS: bool> Migration<T, TEST_ALL_STEPS> {
 				in_progress_version,
 			);
 
-			let result = match T::Migrations::steps(
-				in_progress_version,
-				cursor_before.as_ref(),
-				&mut weight_left,
-			) {
-				StepResult::InProgress { cursor, steps_done } => {
-					*progress = Some(cursor);
-					MigrateResult::InProgress { steps_done }
-				},
-				StepResult::Completed { steps_done } => {
-					in_progress_version.put::<Pallet<T>>();
-					if <Pallet<T>>::in_code_storage_version() != in_progress_version {
-						log::info!(
-							target: LOG_TARGET,
-							"{name}: Next migration is {:?},",
-							in_progress_version + 1
-						);
-						*progress = Some(T::Migrations::new(in_progress_version + 1));
+			let result =
+				match T::Migrations::steps(in_progress_version, cursor_before.as_ref(), &mut meter)
+				{
+					StepResult::InProgress { cursor, steps_done } => {
+						*progress = Some(cursor);
 						MigrateResult::InProgress { steps_done }
-					} else {
-						log::info!(
-							target: LOG_TARGET,
-							"{name}: All migrations done. At version {:?},",
-							in_progress_version
-						);
-						*progress = None;
-						MigrateResult::Completed
-					}
-				},
-			};
+					},
+					StepResult::Completed { steps_done } => {
+						in_progress_version.put::<Pallet<T>>();
+						if <Pallet<T>>::in_code_storage_version() != in_progress_version {
+							log::info!(
+								target: LOG_TARGET,
+								"{name}: Next migration is {:?},",
+								in_progress_version + 1
+							);
+							*progress = Some(T::Migrations::new(in_progress_version + 1));
+							MigrateResult::InProgress { steps_done }
+						} else {
+							log::info!(
+								target: LOG_TARGET,
+								"{name}: All migrations done. At version {:?},",
+								in_progress_version
+							);
+							*progress = None;
+							MigrateResult::Completed
+						}
+					},
+				};
 
-			(result, weight_limit.saturating_sub(weight_left))
+			result
 		})
 	}
 
@@ -516,7 +515,7 @@ impl MigrateSequence for Tuple {
 		invalid_version(version)
 	}
 
-	fn steps(version: StorageVersion, mut cursor: &[u8], weight_left: &mut Weight) -> StepResult {
+	fn steps(version: StorageVersion, mut cursor: &[u8], meter: &mut WeightMeter) -> StepResult {
 		for_tuples!(
 			#(
 				if version == Tuple::VERSION {
@@ -524,11 +523,9 @@ impl MigrateSequence for Tuple {
 						.expect(PROOF_DECODE);
 					let max_weight = Tuple::max_step_weight();
 					let mut steps_done = 0;
-					while weight_left.all_gt(max_weight) {
-						let (finished, weight) = migration.step();
+					while meter.can_consume(max_weight) {
 						steps_done.saturating_accrue(1);
-						weight_left.saturating_reduce(weight);
-						if matches!(finished, IsFinished::Yes) {
+						if matches!(migration.step(meter), IsFinished::Yes) {
 							return StepResult::Completed{ steps_done }
 						}
 					}
@@ -567,13 +564,14 @@ mod test {
 		fn max_step_weight() -> Weight {
 			Weight::from_all(1)
 		}
-		fn step(&mut self) -> (IsFinished, Weight) {
+		fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 			assert!(self.count != N);
 			self.count += 1;
+			meter.consume(Weight::from_all(1));
 			if self.count == N {
-				(IsFinished::Yes, Weight::from_all(1))
+				IsFinished::Yes
 			} else {
-				(IsFinished::No, Weight::from_all(1))
+				IsFinished::No
 			}
 		}
 	}
@@ -603,15 +601,15 @@ mod test {
 		let version = StorageVersion::new(2);
 		let mut cursor = Migrations::new(version);
 
-		let mut weight = Weight::from_all(2);
-		let result = Migrations::steps(version, &cursor, &mut weight);
+		let mut meter = WeightMeter::with_limit(Weight::from_all(1));
+		let result = Migrations::steps(version, &cursor, &mut meter);
 		cursor = vec![1u8, 0].try_into().unwrap();
 		assert_eq!(result, StepResult::InProgress { cursor: cursor.clone(), steps_done: 1 });
-		assert_eq!(weight, Weight::from_all(1));
+		assert_eq!(meter.consumed(), Weight::from_all(1));
 
-		let mut weight = Weight::from_all(2);
+		let mut meter = WeightMeter::with_limit(Weight::from_all(1));
 		assert_eq!(
-			Migrations::steps(version, &cursor, &mut weight),
+			Migrations::steps(version, &cursor, &mut meter),
 			StepResult::Completed { steps_done: 1 }
 		);
 	}
@@ -622,7 +620,10 @@ mod test {
 
 		ExtBuilder::default().build().execute_with(|| {
 			assert_eq!(StorageVersion::get::<Pallet<Test>>(), LATEST_MIGRATION_VERSION);
-			assert_eq!(TestMigration::migrate(Weight::MAX).0, MigrateResult::NoMigrationInProgress)
+			assert_eq!(
+				TestMigration::migrate(&mut WeightMeter::new()),
+				MigrateResult::NoMigrationInProgress
+			)
 		});
 	}
 
@@ -640,7 +641,7 @@ mod test {
 					(LATEST_MIGRATION_VERSION - 1, MigrateResult::InProgress { steps_done: 1 }),
 					(LATEST_MIGRATION_VERSION, MigrateResult::Completed),
 				] {
-					assert_eq!(TestMigration::migrate(Weight::MAX).0, status);
+					assert_eq!(TestMigration::migrate(&mut WeightMeter::new()), status);
 					assert_eq!(
 						<Pallet<Test>>::on_chain_storage_version(),
 						StorageVersion::new(version)
@@ -648,7 +649,7 @@ mod test {
 				}
 
 				assert_eq!(
-					TestMigration::migrate(Weight::MAX).0,
+					TestMigration::migrate(&mut WeightMeter::new()),
 					MigrateResult::NoMigrationInProgress
 				);
 				assert_eq!(StorageVersion::get::<Pallet<Test>>(), LATEST_MIGRATION_VERSION);

substrate/frame/contracts/src/migration/v09.rs

@@ -23,7 +23,9 @@ use crate::{
 	CodeHash, Config, Determinism, Pallet, Weight, LOG_TARGET,
 };
 use codec::{Decode, Encode};
-use frame_support::{pallet_prelude::*, storage_alias, DefaultNoBound, Identity};
+use frame_support::{
+	pallet_prelude::*, storage_alias, weights::WeightMeter, DefaultNoBound, Identity,
+};
 #[cfg(feature = "try-runtime")]
 use sp_runtime::TryRuntimeError;
 use sp_std::prelude::*;
@@ -87,7 +89,7 @@ impl<T: Config> MigrationStep for Migration<T> {
 		T::WeightInfo::v9_migration_step(T::MaxCodeLen::get())
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_key) = self.last_code_hash.take() {
 			v8::CodeStorage::<T>::iter_from(v8::CodeStorage::<T>::hashed_key_for(last_key))
 		} else {
@@ -106,10 +108,12 @@ impl<T: Config> MigrationStep for Migration<T> {
 			};
 			CodeStorage::<T>::insert(key, module);
 			self.last_code_hash = Some(key);
-			(IsFinished::No, T::WeightInfo::v9_migration_step(len))
+			meter.consume(T::WeightInfo::v9_migration_step(len));
+			IsFinished::No
 		} else {
 			log::debug!(target: LOG_TARGET, "No more contracts code to migrate");
-			(IsFinished::Yes, T::WeightInfo::v9_migration_step(0))
+			meter.consume(T::WeightInfo::v9_migration_step(0));
+			IsFinished::Yes
 		}
 	}
 

substrate/frame/contracts/src/migration/v10.rs

@@ -36,6 +36,7 @@ use frame_support::{
 		tokens::{fungible::Inspect, Fortitude::Polite, Preservation::Preserve},
 		ExistenceRequirement, ReservableCurrency,
 	},
+	weights::WeightMeter,
 	DefaultNoBound,
 };
 use sp_core::hexdisplay::HexDisplay;
@@ -160,7 +161,7 @@ where
 		T::WeightInfo::v10_migration_step()
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_account) = self.last_account.take() {
 			v9::ContractInfoOf::<T, OldCurrency>::iter_from(
 				v9::ContractInfoOf::<T, OldCurrency>::hashed_key_for(last_account),
@@ -267,10 +268,12 @@ where
 			// Store last key for next migration step
 			self.last_account = Some(account);
 
-			(IsFinished::No, T::WeightInfo::v10_migration_step())
+			meter.consume(T::WeightInfo::v10_migration_step());
+			IsFinished::No
 		} else {
 			log::debug!(target: LOG_TARGET, "Done Migrating contract info");
-			(IsFinished::Yes, T::WeightInfo::v10_migration_step())
+			meter.consume(T::WeightInfo::v10_migration_step());
+			IsFinished::Yes
 		}
 	}
 

substrate/frame/contracts/src/migration/v11.rs

@@ -23,11 +23,10 @@ use crate::{
 	weights::WeightInfo,
 	Config, Pallet, TrieId, Weight, LOG_TARGET,
 };
+use codec::{Decode, Encode};
+use frame_support::{pallet_prelude::*, storage_alias, weights::WeightMeter, DefaultNoBound};
 #[cfg(feature = "try-runtime")]
 use sp_runtime::TryRuntimeError;
-
-use codec::{Decode, Encode};
-use frame_support::{pallet_prelude::*, storage_alias, DefaultNoBound};
 use sp_std::{marker::PhantomData, prelude::*};
 mod v10 {
 	use super::*;
@@ -79,9 +78,10 @@ impl<T: Config> MigrationStep for Migration<T> {
 		T::WeightInfo::v11_migration_step(128)
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let Some(old_queue) = v10::DeletionQueue::<T>::take() else {
-			return (IsFinished::Yes, Weight::zero())
+			meter.consume(T::WeightInfo::v11_migration_step(0));
+			return IsFinished::Yes
 		};
 		let len = old_queue.len();
 
@@ -101,7 +101,8 @@ impl<T: Config> MigrationStep for Migration<T> {
 			<DeletionQueueCounter<T>>::set(queue);
 		}
 
-		(IsFinished::Yes, T::WeightInfo::v11_migration_step(len as u32))
+		meter.consume(T::WeightInfo::v11_migration_step(len as u32));
+		IsFinished::Yes
 	}
 
 	#[cfg(feature = "try-runtime")]

substrate/frame/contracts/src/migration/v12.rs

@@ -25,7 +25,8 @@ use crate::{
 };
 use codec::{Decode, Encode};
 use frame_support::{
-	pallet_prelude::*, storage_alias, traits::ReservableCurrency, DefaultNoBound, Identity,
+	pallet_prelude::*, storage_alias, traits::ReservableCurrency, weights::WeightMeter,
+	DefaultNoBound, Identity,
 };
 use scale_info::prelude::format;
 use sp_core::hexdisplay::HexDisplay;
@@ -146,7 +147,7 @@ where
 		T::WeightInfo::v12_migration_step(T::MaxCodeLen::get())
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_key) = self.last_code_hash.take() {
 			v11::OwnerInfoOf::<T, OldCurrency>::iter_from(
 				v11::OwnerInfoOf::<T, OldCurrency>::hashed_key_for(last_key),
@@ -230,10 +231,12 @@ where
 
 			self.last_code_hash = Some(hash);
 
-			(IsFinished::No, T::WeightInfo::v12_migration_step(code_len as u32))
+			meter.consume(T::WeightInfo::v12_migration_step(code_len as u32));
+			IsFinished::No
 		} else {
 			log::debug!(target: LOG_TARGET, "No more OwnerInfo to migrate");
-			(IsFinished::Yes, T::WeightInfo::v12_migration_step(0))
+			meter.consume(T::WeightInfo::v12_migration_step(0));
+			IsFinished::Yes
 		}
 	}
 

substrate/frame/contracts/src/migration/v13.rs

@@ -24,7 +24,7 @@ use crate::{
 	AccountIdOf, BalanceOf, CodeHash, Config, Pallet, TrieId, Weight, LOG_TARGET,
 };
 use codec::{Decode, Encode};
-use frame_support::{pallet_prelude::*, storage_alias, DefaultNoBound};
+use frame_support::{pallet_prelude::*, storage_alias, weights::WeightMeter, DefaultNoBound};
 use sp_runtime::BoundedBTreeMap;
 use sp_std::prelude::*;
 
@@ -102,7 +102,7 @@ impl<T: Config> MigrationStep for Migration<T> {
 		T::WeightInfo::v13_migration_step()
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_account) = self.last_account.take() {
 			v12::ContractInfoOf::<T>::iter_from(v12::ContractInfoOf::<T>::hashed_key_for(
 				last_account,
@@ -126,10 +126,12 @@ impl<T: Config> MigrationStep for Migration<T> {
 			};
 			ContractInfoOf::<T>::insert(key.clone(), info);
 			self.last_account = Some(key);
-			(IsFinished::No, T::WeightInfo::v13_migration_step())
+			meter.consume(T::WeightInfo::v13_migration_step());
+			IsFinished::No
 		} else {
 			log::debug!(target: LOG_TARGET, "No more contracts to migrate");
-			(IsFinished::Yes, T::WeightInfo::v13_migration_step())
+			meter.consume(T::WeightInfo::v13_migration_step());
+			IsFinished::Yes
 		}
 	}
 }

substrate/frame/contracts/src/migration/v14.rs

@@ -35,6 +35,7 @@ use frame_support::{
 	pallet_prelude::*,
 	storage_alias,
 	traits::{fungible::MutateHold, ReservableCurrency},
+	weights::WeightMeter,
 	DefaultNoBound,
 };
 use sp_core::hexdisplay::HexDisplay;
@@ -132,7 +133,7 @@ where
 		T::WeightInfo::v14_migration_step()
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_hash) = self.last_code_hash.take() {
 			v13::CodeInfoOf::<T, OldCurrency>::iter_from(
 				v13::CodeInfoOf::<T, OldCurrency>::hashed_key_for(last_hash),
@@ -185,10 +186,12 @@ where
 			});
 
 			self.last_code_hash = Some(hash);
-			(IsFinished::No, T::WeightInfo::v14_migration_step())
+			meter.consume(T::WeightInfo::v14_migration_step());
+			IsFinished::No
 		} else {
 			log::debug!(target: LOG_TARGET, "No more code upload deposit to migrate");
-			(IsFinished::Yes, T::WeightInfo::v14_migration_step())
+			meter.consume(T::WeightInfo::v14_migration_step());
+			IsFinished::Yes
 		}
 	}
 

substrate/frame/contracts/src/migration/v15.rs

@@ -36,6 +36,7 @@ use frame_support::{
 		fungible::{Mutate, MutateHold},
 		tokens::{fungible::Inspect, Fortitude, Preservation},
 	},
+	weights::WeightMeter,
 	BoundedBTreeMap, DefaultNoBound,
 };
 use frame_system::Pallet as System;
@@ -125,7 +126,7 @@ impl<T: Config> MigrationStep for Migration<T> {
 		T::WeightInfo::v15_migration_step()
 	}
 
-	fn step(&mut self) -> (IsFinished, Weight) {
+	fn step(&mut self, meter: &mut WeightMeter) -> IsFinished {
 		let mut iter = if let Some(last_account) = self.last_account.take() {
 			v14::ContractInfoOf::<T>::iter_from(v14::ContractInfoOf::<T>::hashed_key_for(
 				last_account,
@@ -234,10 +235,12 @@ impl<T: Config> MigrationStep for Migration<T> {
 			// Store last key for next migration step
 			self.last_account = Some(account);
 
-			(IsFinished::No, T::WeightInfo::v15_migration_step())
+			meter.consume(T::WeightInfo::v15_migration_step());
+			IsFinished::No
 		} else {
 			log::info!(target: LOG_TARGET, "Done Migrating Storage Deposits.");
-			(IsFinished::Yes, T::WeightInfo::v15_migration_step())
+			meter.consume(T::WeightInfo::v15_migration_step());
+			IsFinished::Yes
 		}
 	}
 

substrate/frame/contracts/src/storage.rs

@@ -28,7 +28,7 @@ use crate::{
 use codec::{Decode, Encode, MaxEncodedLen};
 use frame_support::{
 	storage::child::{self, ChildInfo},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 	CloneNoBound, DefaultNoBound,
 };
 use scale_info::TypeInfo;
@@ -279,14 +279,15 @@ impl<T: Config> ContractInfo<T> {
 
 	/// Calculates the weight that is necessary to remove one key from the trie and how many
 	/// of those keys can be deleted from the deletion queue given the supplied weight limit.
-	pub fn deletion_budget(weight_limit: Weight) -> (Weight, u32) {
+	pub fn deletion_budget(meter: &WeightMeter) -> (Weight, u32) {
 		let base_weight = T::WeightInfo::on_process_deletion_queue_batch();
 		let weight_per_key = T::WeightInfo::on_initialize_per_trie_key(1) -
 			T::WeightInfo::on_initialize_per_trie_key(0);
 
 		// `weight_per_key` being zero makes no sense and would constitute a failure to
 		// benchmark properly. We opt for not removing any keys at all in this case.
-		let key_budget = weight_limit
+		let key_budget = meter
+			.limit()
 			.saturating_sub(base_weight)
 			.checked_div_per_component(&weight_per_key)
 			.unwrap_or(0) as u32;
@@ -295,24 +296,18 @@ impl<T: Config> ContractInfo<T> {
 	}
 
 	/// Delete as many items from the deletion queue possible within the supplied weight limit.
-	///
-	/// It returns the amount of weight used for that task.
-	pub fn process_deletion_queue_batch(weight_limit: Weight) -> Weight {
+	pub fn process_deletion_queue_batch(meter: &mut WeightMeter) {
+		if meter.try_consume(T::WeightInfo::on_process_deletion_queue_batch()).is_err() {
+			return
+		};
+
 		let mut queue = <DeletionQueueManager<T>>::load();
-
 		if queue.is_empty() {
-			return Weight::zero()
-		}
-
-		let (weight_per_key, mut remaining_key_budget) = Self::deletion_budget(weight_limit);
-
-		// We want to check whether we have enough weight to decode the queue before
-		// proceeding. Too little weight for decoding might happen during runtime upgrades
-		// which consume the whole block before the other `on_initialize` blocks are called.
-		if remaining_key_budget == 0 {
-			return weight_limit
+			return;
 		}
 
+		let (weight_per_key, budget) = Self::deletion_budget(&meter);
+		let mut remaining_key_budget = budget;
 		while remaining_key_budget > 0 {
 			let Some(entry) = queue.next() else { break };
 
@@ -324,7 +319,10 @@ impl<T: Config> ContractInfo<T> {
 
 			match outcome {
 				// This happens when our budget wasn't large enough to remove all keys.
-				KillStorageResult::SomeRemaining(_) => return weight_limit,
+				KillStorageResult::SomeRemaining(keys_removed) => {
+					remaining_key_budget.saturating_reduce(keys_removed);
+					break
+				},
 				KillStorageResult::AllRemoved(keys_removed) => {
 					entry.remove();
 					// charge at least one key even if none were removed.
@@ -333,7 +331,7 @@ impl<T: Config> ContractInfo<T> {
 			};
 		}
 
-		weight_limit.saturating_sub(weight_per_key.saturating_mul(u64::from(remaining_key_budget)))
+		meter.consume(weight_per_key.saturating_mul(u64::from(budget - remaining_key_budget)))
 	}
 
 	/// Returns the code hash of the contract specified by `account` ID.

substrate/frame/contracts/src/tests.rs

@@ -54,7 +54,7 @@ use frame_support::{
 		tokens::Preservation,
 		ConstU32, ConstU64, Contains, OnIdle, OnInitialize, StorageVersion,
 	},
-	weights::{constants::WEIGHT_REF_TIME_PER_SECOND, Weight},
+	weights::{constants::WEIGHT_REF_TIME_PER_SECOND, Weight, WeightMeter},
 };
 use frame_system::{EventRecord, Phase};
 use pallet_contracts_fixtures::compile_module;
@@ -1732,8 +1732,8 @@ fn lazy_removal_partial_remove_works() {
 
 	// We create a contract with some extra keys above the weight limit
 	let extra_keys = 7u32;
-	let weight_limit = Weight::from_parts(5_000_000_000, 0);
-	let (_, max_keys) = ContractInfo::<Test>::deletion_budget(weight_limit);
+	let mut meter = WeightMeter::with_limit(Weight::from_parts(5_000_000_000, 100 * 1024));
+	let (weight_per_key, max_keys) = ContractInfo::<Test>::deletion_budget(&meter);
 	let vals: Vec<_> = (0..max_keys + extra_keys)
 		.map(|i| (blake2_256(&i.encode()), (i as u32), (i as u32).encode()))
 		.collect();
@@ -1778,10 +1778,10 @@ fn lazy_removal_partial_remove_works() {
 
 	ext.execute_with(|| {
 		// Run the lazy removal
-		let weight_used = ContractInfo::<Test>::process_deletion_queue_batch(weight_limit);
+		ContractInfo::<Test>::process_deletion_queue_batch(&mut meter);
 
 		// Weight should be exhausted because we could not even delete all keys
-		assert_eq!(weight_used, weight_limit);
+		assert!(!meter.can_consume(weight_per_key));
 
 		let mut num_deleted = 0u32;
 		let mut num_remaining = 0u32;
@@ -1855,7 +1855,7 @@ fn lazy_removal_does_no_run_on_low_remaining_weight() {
 fn lazy_removal_does_not_use_all_weight() {
 	let (code, _hash) = compile_module::<Test>("self_destruct").unwrap();
 
-	let weight_limit = Weight::from_parts(5_000_000_000, 100 * 1024);
+	let mut meter = WeightMeter::with_limit(Weight::from_parts(5_000_000_000, 100 * 1024));
 	let mut ext = ExtBuilder::default().existential_deposit(50).build();
 
 	let (trie, vals, weight_per_key) = ext.execute_with(|| {
@@ -1867,7 +1867,8 @@ fn lazy_removal_does_not_use_all_weight() {
 			.build_and_unwrap_account_id();
 
 		let info = get_contract(&addr);
-		let (weight_per_key, max_keys) = ContractInfo::<Test>::deletion_budget(weight_limit);
+		let (weight_per_key, max_keys) = ContractInfo::<Test>::deletion_budget(&meter);
+		assert!(max_keys > 0);
 
 		// We create a contract with one less storage item than we can remove within the limit
 		let vals: Vec<_> = (0..max_keys - 1)
@@ -1902,10 +1903,10 @@ fn lazy_removal_does_not_use_all_weight() {
 
 	ext.execute_with(|| {
 		// Run the lazy removal
-		let weight_used = ContractInfo::<Test>::process_deletion_queue_batch(weight_limit);
-
-		// We have one less key in our trie than our weight limit suffices for
-		assert_eq!(weight_used, weight_limit - weight_per_key);
+		ContractInfo::<Test>::process_deletion_queue_batch(&mut meter);
+		let base_weight =
+			<<Test as Config>::WeightInfo as WeightInfo>::on_process_deletion_queue_batch();
+		assert_eq!(meter.consumed(), weight_per_key.mul(vals.len() as _) + base_weight);
 
 		// All the keys are removed
 		for val in vals {