Migrating pallet-state-trie-migration to benchmarking V2 (paritytech#6617)

# Description
Migrated pallet-state-trie-migration benchmarking to the new
benchmarking syntax v2.
This is part of paritytech#6202

Co-authored-by: Shawn Tabrizi <[email protected]>
Co-authored-by: Giuseppe Re <[email protected]>

Author: 3 people

substrate/frame/state-trie-migration/src/lib.rs

@@ -249,13 +249,13 @@ pub mod pallet {
 			if limits.item.is_zero() || limits.size.is_zero() {
 				// handle this minor edge case, else we would call `migrate_tick` at least once.
 				log!(warn, "limits are zero. stopping");
-				return Ok(())
+				return Ok(());
 			}
 
 			while !self.exhausted(limits) && !self.finished() {
 				if let Err(e) = self.migrate_tick() {
 					log!(error, "migrate_until_exhaustion failed: {:?}", e);
-					return Err(e)
+					return Err(e);
 				}
 			}
 
@@ -332,7 +332,7 @@ pub mod pallet {
 				_ => {
 					// defensive: there must be an ongoing top migration.
 					frame_support::defensive!("cannot migrate child key.");
-					return Ok(())
+					return Ok(());
 				},
 			};
 
@@ -374,7 +374,7 @@ pub mod pallet {
 				Progress::Complete => {
 					// defensive: there must be an ongoing top migration.
 					frame_support::defensive!("cannot migrate top key.");
-					return Ok(())
+					return Ok(());
 				},
 			};
 
@@ -669,7 +669,7 @@ pub mod pallet {
 			// ensure that the migration witness data was correct.
 			if real_size_upper < task.dyn_size {
 				Self::slash(who, deposit)?;
-				return Ok(().into())
+				return Ok(().into());
 			}
 
 			Self::deposit_event(Event::<T>::Migrated {
@@ -957,6 +957,7 @@ pub mod pallet {
 mod benchmarks {
 	use super::{pallet::Pallet as StateTrieMigration, *};
 	use alloc::vec;
+	use frame_benchmarking::v2::*;
 	use frame_support::traits::fungible::{Inspect, Mutate};
 
 	// The size of the key seemingly makes no difference in the read/write time, so we make it
@@ -970,8 +971,12 @@ mod benchmarks {
 		stash
 	}
 
-	frame_benchmarking::benchmarks! {
-		continue_migrate {
+	#[benchmarks]
+	mod inner_benchmarks {
+		use super::*;
+
+		#[benchmark]
+		fn continue_migrate() -> Result<(), BenchmarkError> {
 			// note that this benchmark should migrate nothing, as we only want the overhead weight
 			// of the bookkeeping, and the migration cost itself is noted via the `dynamic_weight`
 			// function.
@@ -980,116 +985,151 @@ mod benchmarks {
 			let stash = set_balance_for_deposit::<T>(&caller, null.item);
 			// Allow signed migrations.
 			SignedMigrationMaxLimits::<T>::put(MigrationLimits { size: 1024, item: 5 });
-		}: _(frame_system::RawOrigin::Signed(caller.clone()), null, 0, StateTrieMigration::<T>::migration_process())
-		verify {
+
+			#[extrinsic_call]
+			_(
+				frame_system::RawOrigin::Signed(caller.clone()),
+				null,
+				0,
+				StateTrieMigration::<T>::migration_process(),
+			);
+
 			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
-			assert_eq!(T::Currency::balance(&caller), stash)
+			assert_eq!(T::Currency::balance(&caller), stash);
+
+			Ok(())
 		}
 
-		continue_migrate_wrong_witness {
+		#[benchmark]
+		fn continue_migrate_wrong_witness() -> Result<(), BenchmarkError> {
 			let null = MigrationLimits::default();
 			let caller = frame_benchmarking::whitelisted_caller();
-			let bad_witness = MigrationTask { progress_top: Progress::LastKey(vec![1u8].try_into().unwrap()), ..Default::default() };
-		}: {
-			assert!(
-				StateTrieMigration::<T>::continue_migrate(
+			let bad_witness = MigrationTask {
+				progress_top: Progress::LastKey(vec![1u8].try_into().unwrap()),
+				..Default::default()
+			};
+			#[block]
+			{
+				assert!(StateTrieMigration::<T>::continue_migrate(
 					frame_system::RawOrigin::Signed(caller).into(),
 					null,
 					0,
 					bad_witness,
 				)
-				.is_err()
-			)
-		}
-		verify {
-			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default())
+				.is_err());
+			}
+
+			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
+
+			Ok(())
 		}
 
-		migrate_custom_top_success {
+		#[benchmark]
+		fn migrate_custom_top_success() -> Result<(), BenchmarkError> {
 			let null = MigrationLimits::default();
 			let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
 			let stash = set_balance_for_deposit::<T>(&caller, null.item);
-		}: migrate_custom_top(frame_system::RawOrigin::Signed(caller.clone()), Default::default(), 0)
-		verify {
+			#[extrinsic_call]
+			migrate_custom_top(
+				frame_system::RawOrigin::Signed(caller.clone()),
+				Default::default(),
+				0,
+			);
+
 			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
-			assert_eq!(T::Currency::balance(&caller), stash)
+			assert_eq!(T::Currency::balance(&caller), stash);
+			Ok(())
 		}
 
-		migrate_custom_top_fail {
+		#[benchmark]
+		fn migrate_custom_top_fail() -> Result<(), BenchmarkError> {
 			let null = MigrationLimits::default();
 			let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
 			let stash = set_balance_for_deposit::<T>(&caller, null.item);
 			// for tests, we need to make sure there is _something_ in storage that is being
 			// migrated.
-			sp_io::storage::set(b"foo", vec![1u8;33].as_ref());
-		}: {
-			assert!(
-				StateTrieMigration::<T>::migrate_custom_top(
+			sp_io::storage::set(b"foo", vec![1u8; 33].as_ref());
+			#[block]
+			{
+				assert!(StateTrieMigration::<T>::migrate_custom_top(
 					frame_system::RawOrigin::Signed(caller.clone()).into(),
 					vec![b"foo".to_vec()],
 					1,
-				).is_ok()
-			);
+				)
+				.is_ok());
+
+				frame_system::Pallet::<T>::assert_last_event(
+					<T as Config>::RuntimeEvent::from(crate::Event::Slashed {
+						who: caller.clone(),
+						amount: StateTrieMigration::<T>::calculate_deposit_for(1u32),
+					})
+					.into(),
+				);
+			}
 
-			frame_system::Pallet::<T>::assert_last_event(
-				<T as Config>::RuntimeEvent::from(crate::Event::Slashed {
-					who: caller.clone(),
-					amount: StateTrieMigration::<T>::calculate_deposit_for(1u32),
-				}).into(),
-			);
-		}
-		verify {
 			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
 			// must have gotten slashed
-			assert!(T::Currency::balance(&caller) < stash)
+			assert!(T::Currency::balance(&caller) < stash);
+
+			Ok(())
 		}
 
-		migrate_custom_child_success {
+		#[benchmark]
+		fn migrate_custom_child_success() -> Result<(), BenchmarkError> {
 			let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
 			let stash = set_balance_for_deposit::<T>(&caller, 0);
-		}: migrate_custom_child(
-			frame_system::RawOrigin::Signed(caller.clone()),
-			StateTrieMigration::<T>::childify(Default::default()),
-			Default::default(),
-			0
-		)
-		verify {
+
+			#[extrinsic_call]
+			migrate_custom_child(
+				frame_system::RawOrigin::Signed(caller.clone()),
+				StateTrieMigration::<T>::childify(Default::default()),
+				Default::default(),
+				0,
+			);
+
 			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
 			assert_eq!(T::Currency::balance(&caller), stash);
+
+			Ok(())
 		}
 
-		migrate_custom_child_fail {
+		#[benchmark]
+		fn migrate_custom_child_fail() -> Result<(), BenchmarkError> {
 			let caller: T::AccountId = frame_benchmarking::whitelisted_caller();
 			let stash = set_balance_for_deposit::<T>(&caller, 1);
 			// for tests, we need to make sure there is _something_ in storage that is being
 			// migrated.
-			sp_io::default_child_storage::set(b"top", b"foo", vec![1u8;33].as_ref());
-		}: {
-			assert!(
-				StateTrieMigration::<T>::migrate_custom_child(
+			sp_io::default_child_storage::set(b"top", b"foo", vec![1u8; 33].as_ref());
+
+			#[block]
+			{
+				assert!(StateTrieMigration::<T>::migrate_custom_child(
 					frame_system::RawOrigin::Signed(caller.clone()).into(),
 					StateTrieMigration::<T>::childify("top"),
 					vec![b"foo".to_vec()],
 					1,
-				).is_ok()
-			)
-		}
-		verify {
+				)
+				.is_ok());
+			}
 			assert_eq!(StateTrieMigration::<T>::migration_process(), Default::default());
 			// must have gotten slashed
-			assert!(T::Currency::balance(&caller) < stash)
+			assert!(T::Currency::balance(&caller) < stash);
+			Ok(())
 		}
 
-		process_top_key {
-			let v in 1 .. (4 * 1024 * 1024);
-
+		#[benchmark]
+		fn process_top_key(v: Linear<1, { 4 * 1024 * 1024 }>) -> Result<(), BenchmarkError> {
 			let value = alloc::vec![1u8; v as usize];
 			sp_io::storage::set(KEY, &value);
-		}: {
-			let data = sp_io::storage::get(KEY).unwrap();
-			sp_io::storage::set(KEY, &data);
-			let _next = sp_io::storage::next_key(KEY);
-			assert_eq!(data, value);
+			#[block]
+			{
+				let data = sp_io::storage::get(KEY).unwrap();
+				sp_io::storage::set(KEY, &data);
+				let _next = sp_io::storage::next_key(KEY);
+				assert_eq!(data, value);
+			}
+
+			Ok(())
 		}
 
 		impl_benchmark_test_suite!(
@@ -1741,7 +1781,7 @@ pub(crate) mod remote_tests {
 			let ((finished, weight), proof) = ext.execute_and_prove(|| {
 				let weight = run_to_block::<Runtime>(now + One::one()).1;
 				if StateTrieMigration::<Runtime>::migration_process().finished() {
-					return (true, weight)
+					return (true, weight);
 				}
 				duration += One::one();
 				now += One::one();
@@ -1768,7 +1808,7 @@ pub(crate) mod remote_tests {
 			ext.commit_all().unwrap();
 
 			if finished {
-				break
+				break;
 			}
 		}