add vesting benchmarks (#1039)

Author: zjb0807

vesting/Cargo.toml

@@ -12,6 +12,7 @@ parity-scale-codec = { workspace = true, default-features = false, features = ["
 scale-info = { workspace = true }
 serde = { workspace = true, optional = true }
 
+frame-benchmarking = { workspace = true, optional = true }
 frame-support = { workspace = true }
 frame-system = { workspace = true }
 sp-io = { workspace = true }
@@ -25,6 +26,7 @@ sp-core = { workspace = true, features = ["std"] }
 [features]
 default = [ "std" ]
 std = [
+	"frame-benchmarking?/std",
 	"frame-support/std",
 	"frame-system/std",
 	"parity-scale-codec/std",
@@ -35,6 +37,7 @@ std = [
 	"sp-std/std",
 ]
 runtime-benchmarks = [
+	"frame-benchmarking/runtime-benchmarks",
 	"frame-support/runtime-benchmarks",
 	"frame-system/runtime-benchmarks",
 	"sp-runtime/runtime-benchmarks",

vesting/src/benchmarking.rs

@@ -0,0 +1,129 @@
+pub use crate::*;
+
+use frame_benchmarking::v2::*;
+use frame_support::assert_ok;
+use frame_system::RawOrigin;
+use sp_std::vec;
+
+/// Helper trait for benchmarking.
+pub trait BenchmarkHelper<AccountId, Balance> {
+	fn get_vesting_account_and_amount() -> Option<(AccountId, Balance)>;
+}
+
+impl<AccountId, Balance> BenchmarkHelper<AccountId, Balance> for () {
+	fn get_vesting_account_and_amount() -> Option<(AccountId, Balance)> {
+		None
+	}
+}
+
+fn set_balance<T: Config>(who: &T::AccountId, amount: BalanceOf<T>) {
+	let _ = <<T as Config>::Currency as Currency<_>>::deposit_creating(&who, amount);
+}
+
+fn total_balance<T: Config>(who: &T::AccountId) -> BalanceOf<T> {
+	<<T as Config>::Currency as Currency<_>>::total_balance(who)
+}
+
+fn free_balance<T: Config>(who: &T::AccountId) -> BalanceOf<T> {
+	<<T as Config>::Currency as Currency<_>>::free_balance(who)
+}
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn vested_transfer() {
+		let schedule = VestingScheduleOf::<T> {
+			start: 0u32.into(),
+			period: 2u32.into(),
+			period_count: 3u32.into(),
+			per_period: T::MinVestedTransfer::get(),
+		};
+
+		// extra 1 dollar to pay fees
+		let (from, amount) = T::BenchmarkHelper::get_vesting_account_and_amount().unwrap();
+		set_balance::<T>(&from, schedule.total_amount().unwrap() + amount);
+
+		let to: T::AccountId = account("to", 0, 0);
+		let to_lookup = <T as frame_system::Config>::Lookup::unlookup(to.clone());
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(from), to_lookup, schedule.clone());
+
+		assert_eq!(total_balance::<T>(&to), schedule.total_amount().unwrap());
+	}
+
+	#[benchmark]
+	fn claim(i: Linear<1, { T::MaxVestingSchedules::get() }>) {
+		let mut schedule = VestingScheduleOf::<T> {
+			start: 0u32.into(),
+			period: 2u32.into(),
+			period_count: 3u32.into(),
+			per_period: T::MinVestedTransfer::get(),
+		};
+
+		// extra 1 dollar to pay fees
+		let (from, amount) = T::BenchmarkHelper::get_vesting_account_and_amount().unwrap();
+		set_balance::<T>(
+			&from,
+			schedule.total_amount().unwrap().saturating_mul(i.into()) + amount,
+		);
+
+		let to: T::AccountId = account("to", 0, 0);
+		let to_lookup = <T as frame_system::Config>::Lookup::unlookup(to.clone());
+
+		for _ in 0..i {
+			schedule.start = i.into();
+			assert_ok!(Pallet::<T>::vested_transfer(
+				RawOrigin::Signed(from.clone()).into(),
+				to_lookup.clone(),
+				schedule.clone()
+			));
+		}
+		frame_system::Pallet::<T>::set_block_number(schedule.end().unwrap() + 1u32.into());
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(to.clone()));
+
+		assert_eq!(
+			free_balance::<T>(&to),
+			schedule.total_amount().unwrap().saturating_mul(i.into()),
+		);
+	}
+
+	#[benchmark]
+	fn update_vesting_schedules(i: Linear<1, { T::MaxVestingSchedules::get() }>) {
+		let mut schedule = VestingScheduleOf::<T> {
+			start: 0u32.into(),
+			period: 2u32.into(),
+			period_count: 3u32.into(),
+			per_period: T::MinVestedTransfer::get(),
+		};
+
+		let to: T::AccountId = account("to", 0, 0);
+		let to_lookup = <T as frame_system::Config>::Lookup::unlookup(to.clone());
+
+		set_balance::<T>(&to, schedule.total_amount().unwrap().saturating_mul(i.into()));
+
+		let mut schedules = vec![];
+		for _ in 0..i {
+			schedule.start = i.into();
+			schedules.push(schedule.clone());
+		}
+
+		#[extrinsic_call]
+		_(RawOrigin::Root, to_lookup, schedules);
+
+		assert_eq!(
+			free_balance::<T>(&to),
+			schedule.total_amount().unwrap().saturating_mul(i.into()),
+		);
+	}
+
+	impl_benchmark_test_suite! {
+		Pallet,
+		crate::mock::ExtBuilder::build(),
+		crate::mock::Runtime,
+	}
+}

vesting/src/lib.rs

@@ -45,10 +45,14 @@ use sp_std::{
 	vec::Vec,
 };
 
+#[cfg(feature = "runtime-benchmarks")]
+mod benchmarking;
 mod mock;
 mod tests;
 mod weights;
 
+#[cfg(feature = "runtime-benchmarks")]
+pub use benchmarking::BenchmarkHelper;
 pub use module::*;
 pub use weights::WeightInfo;
 
@@ -143,6 +147,9 @@ pub mod module {
 
 		// The block number provider
 		type BlockNumberProvider: BlockNumberProvider<BlockNumber = BlockNumberFor<Self>>;
+
+		#[cfg(feature = "runtime-benchmarks")]
+		type BenchmarkHelper: BenchmarkHelper<Self::AccountId, BalanceOf<Self>>;
 	}
 
 	#[pallet::error]

vesting/src/mock.rs

@@ -73,13 +73,24 @@ impl BlockNumberProvider for MockBlockNumberProvider {
 	}
 }
 
+#[cfg(feature = "runtime-benchmarks")]
+pub struct MockBenchmarkHelper;
+#[cfg(feature = "runtime-benchmarks")]
+impl BenchmarkHelper<AccountId, Balance> for MockBenchmarkHelper {
+	fn get_vesting_account_and_amount() -> Option<(AccountId, Balance)> {
+		Some((ALICE, 1000))
+	}
+}
+
 impl Config for Runtime {
 	type Currency = PalletBalances;
 	type MinVestedTransfer = ConstU64<5>;
 	type VestedTransferOrigin = EnsureAliceOrBob;
 	type WeightInfo = ();
 	type MaxVestingSchedules = ConstU32<2>;
 	type BlockNumberProvider = MockBlockNumberProvider;
+	#[cfg(feature = "runtime-benchmarks")]
+	type BenchmarkHelper = MockBenchmarkHelper;
 }
 
 type Block = frame_system::mocking::MockBlock<Runtime>;