Benchmarking fixes (#1219)

pallets/elections/Cargo.toml

@@ -20,7 +20,7 @@ scale-codec.workspace = true
 scale-info.workspace = true
 serde.workspace = true
 simple-mermaid.workspace = true
-
+pallet-members.workspace = true
 polkadot-sdk = { workspace = true, features = [ "frame-support", "frame-system", "sp-runtime", "sp-std" ] }
 
 time-primitives.workspace = true
@@ -40,7 +40,7 @@ std = [
 	"scale-codec/std",
 	"scale-info/std",
 	"serde/std",
-
+	"pallet-members/std",
 	"polkadot-sdk/std",
 
 	"time-primitives/std",

pallets/elections/src/benchmarking.rs

@@ -2,14 +2,50 @@ use super::*;
 use crate::Pallet;
 
 use polkadot_sdk::frame_benchmarking::benchmarks;
-use polkadot_sdk::frame_system;
+use polkadot_sdk::{frame_system, sp_runtime};
 
 use frame_system::RawOrigin;
+use pallet_members::{MemberOnline, MemberStake};
+use sp_runtime::Vec;
+use time_primitives::{AccountId, NetworkId};
+const ETHEREUM: NetworkId = 0;
 
 benchmarks! {
+	where_clause { where T: pallet_members::Config }
+
 	set_shard_config {
 	}: _(RawOrigin::Root, 3, 1)
 	verify { }
 
+	try_elect_shard {
+		let b in (ShardSize::<T>::get().into())..256;
+		let max_stake: u128 = 1_000_000_000;
+		let mut shard: Vec<AccountId> = Vec::new();
+		for i in 0..b {
+			let member = Into::<AccountId>::into([i as u8; 32]);
+			Unassigned::<T>::insert(ETHEREUM, member.clone(), ());
+			MemberOnline::<T>::insert(member.clone(), ());
+			let member_stake: u128 = max_stake - Into::<u128>::into(i);
+			MemberStake::<T>::insert(member.clone(), member_stake);
+			if (shard.len() as u16) < ShardSize::<T>::get() {
+				shard.push(member);
+			}
+		}
+		let pre_unassigned_count: u16 = Unassigned::<T>::iter().count().try_into().unwrap_or_default();
+	}: {
+		Pallet::<T>::try_elect_shard(ETHEREUM);
+	} verify {
+		let post_unassigned_count: u16 = Unassigned::<T>::iter().count().try_into().unwrap_or_default();
+		// ShardSize # of unassigned were elected to a shard
+		assert_eq!(
+			pre_unassigned_count - post_unassigned_count,
+			ShardSize::<T>::get(),
+		);
+		// New shard members were removed from Unassigned
+		for m in shard {
+			assert!(Unassigned::<T>::get(ETHEREUM, m).is_none());
+		}
+	}
+
 	impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test);
 }

pallets/elections/src/lib.rs

@@ -45,10 +45,11 @@ mod tests;
 
 #[polkadot_sdk::frame_support::pallet]
 pub mod pallet {
-	use polkadot_sdk::{frame_support, frame_system, sp_std};
+	use polkadot_sdk::{frame_support, frame_system, sp_runtime, sp_std};
 
 	use frame_support::pallet_prelude::*;
 	use frame_system::pallet_prelude::*;
+	use sp_runtime::Saturating;
 	use sp_std::vec;
 	use sp_std::vec::Vec;
 
@@ -58,12 +59,16 @@ pub mod pallet {
 
 	pub trait WeightInfo {
 		fn set_shard_config() -> Weight;
+		fn try_elect_shard(b: u32) -> Weight;
 	}
 
 	impl WeightInfo for () {
 		fn set_shard_config() -> Weight {
 			Weight::default()
 		}
+		fn try_elect_shard(_: u32) -> Weight {
+			Weight::default()
+		}
 	}
 
 	#[pallet::pallet]
@@ -159,12 +164,9 @@ pub mod pallet {
 		fn on_initialize(_: BlockNumberFor<T>) -> Weight {
 			log::info!("on_initialize begin");
 			let mut weight = Weight::default();
-			for (network, _, _) in Unassigned::<T>::iter() {
-				weight = weight
-					// 1 Read of Unassigned per Loop
-					.saturating_add(T::DbWeight::get().reads(1))
-					.saturating_add(Self::try_elect_shard(network));
-			}
+			Unassigned::<T>::iter().map(|(n, _, _)| n).for_each(|network| {
+				weight = weight.saturating_add(Self::try_elect_shard(network));
+			});
 			log::info!("on_initialize end");
 			weight
 		}
@@ -254,11 +256,9 @@ pub mod pallet {
 		///    1. Removes the member from the [`Unassigned`] storage for the given network.
 		///    2. Notifies the `Shards` interface about the member going offline.
 		///    3. Returns the weight of the operation.
-		fn member_offline(member: &AccountId, network: NetworkId) -> Weight {
+		fn member_offline(member: &AccountId, network: NetworkId) {
 			Unassigned::<T>::remove(network, member);
-			T::DbWeight::get()
-				.writes(1)
-				.saturating_add(T::Shards::member_offline(member, network))
+			T::Shards::member_offline(member, network);
 		}
 	}
 
@@ -273,18 +273,16 @@ pub mod pallet {
 		///    1. Calls `new_shard_members` to get a list of new shard members.
 		///    2. If a new shard can be formed, removes the selected members from [`Unassigned`] storage.
 		///    3. Creates a new shard using the `Shards` interface with the selected members and current shard threshold.
-		fn try_elect_shard(network: NetworkId) -> Weight {
-			match Self::new_shard_members(network) {
-				(Some(members), r) => {
+		pub(crate) fn try_elect_shard(network: NetworkId) -> Weight {
+			let num_unassigned = match Self::new_shard_members(network) {
+				(Some(members), n) => {
 					members.iter().for_each(|m| Unassigned::<T>::remove(network, m));
-					let weight = T::DbWeight::get()
-						.reads_writes(r, members.len().try_into().unwrap_or_default());
-					T::Shards::create_shard(network, members, ShardThreshold::<T>::get())
-						.1
-						.saturating_add(weight)
+					T::Shards::create_shard(network, members, ShardThreshold::<T>::get());
+					n
 				},
-				(None, r) => T::DbWeight::get().reads(r),
-			}
+				(None, n) => n,
+			};
+			T::WeightInfo::try_elect_shard(num_unassigned)
 		}
 
 		///  Determines the members for a new shard.
@@ -294,21 +292,21 @@ pub mod pallet {
 		///    3. Returns `None` if there are not enough members to form a shard.
 		///    4. If there are just enough members, returns the list of members.
 		///    5. If there are more members than needed, sorts them by their stake and selects the top members to form the shard.
-		fn new_shard_members(network: NetworkId) -> (Option<Vec<AccountId>>, u64) {
-			let mut reads: u64 = 0;
-			let shard_members_len = ShardSize::<T>::get() as usize;
+		fn new_shard_members(network: NetworkId) -> (Option<Vec<AccountId>>, u32) {
+			let mut num_unassigned: u32 = 0;
 			let mut members = Vec::new();
 			for (m, _) in Unassigned::<T>::iter_prefix(network) {
 				if T::Members::is_member_online(&m) {
 					members.push(m);
 				}
-				reads = reads.saturating_add(2);
+				num_unassigned = num_unassigned.saturating_plus_one();
 			}
+			let shard_members_len = ShardSize::<T>::get() as usize;
 			if members.len() < shard_members_len {
-				return (None, reads);
+				return (None, num_unassigned);
 			}
 			if members.len() == shard_members_len {
-				return (Some(members), reads);
+				return (Some(members), num_unassigned);
 			}
 			// else members.len() > shard_members_len:
 			members.sort_unstable_by(|a, b| {
@@ -318,7 +316,7 @@ pub mod pallet {
 					.then_with(|| a.cmp(b))
 					.reverse()
 			});
-			(Some(members.into_iter().take(shard_members_len).collect()), reads)
+			(Some(members.into_iter().take(shard_members_len).collect()), num_unassigned)
 		}
 	}
 }

pallets/elections/src/mock.rs

@@ -4,41 +4,18 @@ use polkadot_sdk::*;
 
 use frame_support::derive_impl;
 use frame_support::traits::OnInitialize;
-use sp_core::{ConstU128, ConstU64};
+use sp_core::{ConstU128, ConstU32, ConstU64};
 use sp_runtime::{
 	traits::{IdentifyAccount, IdentityLookup, Verify},
-	BuildStorage, DispatchResult, MultiSignature,
+	BuildStorage, MultiSignature,
 };
-use time_primitives::{MembersInterface, NetworkId, PeerId, PublicKey, ShardId, TasksInterface};
+use time_primitives::{NetworkId, ShardId, TasksInterface};
 
 pub type UncheckedExtrinsic = frame_system::mocking::MockUncheckedExtrinsic<Test>;
 type Block = frame_system::mocking::MockBlock<Test>;
 pub type AccountId = <<Signature as Verify>::Signer as IdentifyAccount>::AccountId;
 pub type Signature = MultiSignature;
 
-pub struct MockMembers;
-
-impl MembersInterface for MockMembers {
-	fn member_stake(_: &AccountId) -> u128 {
-		0u128
-	}
-	fn member_peer_id(_: &AccountId) -> Option<PeerId> {
-		None
-	}
-	fn member_public_key(_account: &AccountId) -> Option<PublicKey> {
-		None
-	}
-	fn is_member_online(_: &AccountId) -> bool {
-		true
-	}
-	fn total_stake() -> u128 {
-		0u128
-	}
-	fn transfer_stake(_: &AccountId, _: &AccountId, _: u128) -> DispatchResult {
-		Ok(())
-	}
-}
-
 pub struct MockTasks;
 
 impl TasksInterface for MockTasks {
@@ -54,6 +31,7 @@ frame_support::construct_runtime!(
 		Balances: pallet_balances::{Pallet, Call, Storage, Event<T>},
 		Shards: pallet_shards::{Pallet, Call, Storage, Event<T>},
 		Elections: pallet_elections::{Pallet, Call, Storage, Config<T>, Event<T>},
+		Members: pallet_members,
 	}
 );
 
@@ -85,19 +63,29 @@ impl pallet_elections::Config for Test {
 	type AdminOrigin = frame_system::EnsureRoot<AccountId>;
 	type WeightInfo = ();
 	type Shards = Shards;
-	type Members = MockMembers;
+	type Members = Members;
 }
 
 impl pallet_shards::Config for Test {
 	type RuntimeEvent = RuntimeEvent;
 	type AdminOrigin = frame_system::EnsureRoot<AccountId>;
 	type WeightInfo = ();
 	type Tasks = MockTasks;
-	type Members = MockMembers;
+	type Members = Members;
 	type Elections = Elections;
+	type MaxTimeoutsPerBlock = ConstU32<100>;
 	type DkgTimeout = ConstU64<10>;
 }
 
+impl pallet_members::Config for Test {
+	type WeightInfo = ();
+	type RuntimeEvent = RuntimeEvent;
+	type Elections = Elections;
+	type MinStake = ConstU128<5>;
+	type HeartbeatTimeout = ConstU64<10>;
+	type MaxTimeoutsPerBlock = ConstU32<100>;
+}
+
 impl<LocalCall> frame_system::offchain::CreateSignedTransaction<LocalCall> for Test
 where
 	RuntimeCall: From<LocalCall>,

pallets/elections/src/tests.rs

@@ -4,7 +4,7 @@ use polkadot_sdk::*;
 
 use frame_support::{assert_noop, assert_ok};
 use frame_system::RawOrigin;
-
+use pallet_members::MemberOnline;
 use time_primitives::{ElectionsInterface, NetworkId};
 
 const ETHEREUM: NetworkId = 0;
@@ -25,12 +25,15 @@ fn shard_size_new_members_online_creates_shard() {
 	let c: AccountId = [3u8; 32].into();
 	new_test_ext().execute_with(|| {
 		Elections::member_online(&a, ETHEREUM);
+		MemberOnline::<Test>::insert(&a, ());
 		roll(1);
 		assert!(Unassigned::<Test>::get(ETHEREUM, &a).is_some());
 		Elections::member_online(&b, ETHEREUM);
+		MemberOnline::<Test>::insert(&b, ());
 		roll(1);
 		assert!(Unassigned::<Test>::get(ETHEREUM, &b).is_some());
 		Elections::member_online(&c, ETHEREUM);
+		MemberOnline::<Test>::insert(&c, ());
 		roll(1);
 		System::assert_last_event(pallet_shards::Event::<Test>::ShardCreated(0, ETHEREUM).into());
 		for member in [a, b, c] {
@@ -58,8 +61,11 @@ fn shard_offline_automatically_creates_new_shard() {
 	let c: AccountId = [3u8; 32].into();
 	new_test_ext().execute_with(|| {
 		Elections::member_online(&a, ETHEREUM);
+		MemberOnline::<Test>::insert(&a, ());
 		Elections::member_online(&b, ETHEREUM);
+		MemberOnline::<Test>::insert(&b, ());
 		Elections::member_online(&c, ETHEREUM);
+		MemberOnline::<Test>::insert(&c, ());
 		roll(1);
 		System::assert_last_event(pallet_shards::Event::<Test>::ShardCreated(0, ETHEREUM).into());
 		Elections::shard_offline(ETHEREUM, [a, b, c].to_vec());

pallets/members/src/benchmarking.rs

@@ -6,14 +6,17 @@ use polkadot_sdk::*;
 use frame_benchmarking::benchmarks;
 use frame_support::traits::{Currency, Get};
 use frame_system::RawOrigin;
-
 use time_primitives::{AccountId, NetworkId, PublicKey};
 
 pub const ALICE: [u8; 32] = [1u8; 32];
 pub const ETHEREUM: NetworkId = 1;
 
 fn public_key() -> PublicKey {
-	PublicKey::Sr25519(sp_core::sr25519::Public::from_raw(ALICE))
+	pk_from_account(ALICE)
+}
+
+fn pk_from_account(r: [u8; 32]) -> PublicKey {
+	PublicKey::Sr25519(sp_core::sr25519::Public::from_raw(r))
 }
 
 benchmarks! {
@@ -46,5 +49,27 @@ benchmarks! {
 	}: _(RawOrigin::Signed(caller))
 	verify { }
 
+	timeout_heartbeats {
+		let b in 1..T::MaxTimeoutsPerBlock::get();
+		for i in 0..b {
+			let raw = [i as u8; 32];
+			let caller: AccountId = raw.into();
+			pallet_balances::Pallet::<T>::resolve_creating(
+				&caller,
+				pallet_balances::Pallet::<T>::issue(<T as Config>::MinStake::get() * 100),
+			);
+			Pallet::<T>::register_member(RawOrigin::Signed(caller.clone()).into(), ETHEREUM, pk_from_account(raw), caller.clone().into(), <T as Config>::MinStake::get())?;
+			assert!(MemberOnline::<T>::get(&caller).is_some());
+			assert!(Heartbeat::<T>::take(&caller).is_some());
+		}
+	}: {
+		Pallet::<T>::timeout_heartbeats(0u32.into());
+	} verify {
+		for i in 0..b {
+			let caller: AccountId = [i as u8; 32].into();
+			assert!(MemberOnline::<T>::get(&caller).is_none());
+		}
+	}
+
 	impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test);
 }