1234567891011121314151617181920212223242526272829303132333435use frame_benchmarking::v2::{36 account, benchmarks, impl_benchmark_test_suite, impl_test_function, whitelisted_caller,37 BenchmarkError,38};39use frame_support::{40 assert_ok,41 traits::{42 fungible::{Inspect, Mutate},43 EnsureOrigin, Get,44 },45};46use frame_system::{pallet_prelude::*, EventRecord, RawOrigin};47use pallet_authorship::EventHandler;48use pallet_session::{self as session, SessionManager};49use parity_scale_codec::Decode;50use sp_std::prelude::*;5152use super::*;53#[allow(unused)]54use crate::{BalanceOf, Pallet as CollatorSelection};5556const SEED: u32 = 0;575859macro_rules! whitelist {60 ($acc:ident) => {61 frame_benchmarking::benchmarking::add_to_whitelist(62 frame_system::Account::<T>::hashed_key_for(&$acc).into(),63 );64 };65}6667fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {68 let events = frame_system::Pallet::<T>::events();69 let system_event: <T as frame_system::Config>::RuntimeEvent = generic_event.into();70 71 let EventRecord { event, .. } = &events[events.len() - 1];72 assert_eq!(event, &system_event);73}7475fn create_funded_user<T: Config>(76 string: &'static str,77 n: u32,78 balance_factor: u32,79) -> T::AccountId {80 let user = account(string, n, SEED);81 let balance = balance_unit::<T>() * balance_factor.into();82 let _ = T::Currency::set_balance(&user, balance);83 user84}8586fn keys<T: Config + session::Config>(c: u32) -> <T as session::Config>::Keys {87 use rand::{RngCore, SeedableRng};8889 let keys = {90 let mut keys = [0u8; 128];9192 if c > 0 {93 let mut rng = rand::rngs::StdRng::seed_from_u64(c as u64);94 rng.fill_bytes(&mut keys);95 }9697 keys98 };99100 Decode::decode(&mut &keys[..]).unwrap()101}102103fn validator<T: Config + session::Config>(c: u32) -> (T::AccountId, <T as session::Config>::Keys) {104 (create_funded_user::<T>("candidate", c, 1000), keys::<T>(c))105}106107fn register_validators<T: Config + session::Config>(count: u32) -> Vec<T::AccountId> {108 let validators = (0..count).map(|c| validator::<T>(c)).collect::<Vec<_>>();109110 for (who, keys) in validators.clone() {111 <session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();112 }113114 validators.into_iter().map(|(who, _)| who).collect()115}116117fn register_invulnerables<T: Config>(count: u32) {118 let candidates = (0..count)119 .map(|c| account("candidate", c, SEED))120 .collect::<Vec<_>>();121122 for who in candidates {123 <CollatorSelection<T>>::add_invulnerable(124 T::UpdateOrigin::try_successful_origin().unwrap(),125 who,126 )127 .unwrap();128 }129}130131fn register_candidates<T: Config>(count: u32) {132 let candidates = (0..count)133 .map(|c| account("candidate", c, SEED))134 .collect::<Vec<_>>();135 assert!(T::LicenseBond::get() > 0u32.into(), "Bond cannot be zero!");136137 for who in candidates {138 T::Currency::set_balance(&who, T::LicenseBond::get() * 2u32.into());139 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();140 <CollatorSelection<T>>::onboard(RawOrigin::Signed(who).into()).unwrap();141 }142}143144fn get_licenses<T: Config>(count: u32) {145 let candidates = (0..count)146 .map(|c| account("candidate", c, SEED))147 .collect::<Vec<_>>();148 assert!(T::LicenseBond::get() > 0u32.into(), "Bond cannot be zero!");149150 for who in candidates {151 T::Currency::set_balance(&who, T::LicenseBond::get() * 2u32.into());152 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();153 }154}155156157158fn balance_unit<T: Config>() -> BalanceOf<T> {159 T::LicenseBond::get()160}161162163const INITIAL_INVULNERABLES: u32 = 2;164165#[benchmarks(where T: Config + pallet_authorship::Config + session::Config)]166mod benchmarks {167 use super::*;168 const MAX_COLLATORS: u32 = 10;169 const MAX_INVULNERABLES: u32 = MAX_COLLATORS - INITIAL_INVULNERABLES;170171 172 173 174 #[benchmark]175 fn add_invulnerable<T>(b: Linear<2, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {176 let b = b - 1;177 register_validators::<T>(b);178 register_invulnerables::<T>(b);179180 181182 let new_invulnerable: T::AccountId = whitelisted_caller();183 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();184 <T as Config>::Currency::set_balance(&new_invulnerable, bond);185186 <session::Pallet<T>>::set_keys(187 RawOrigin::Signed(new_invulnerable.clone()).into(),188 keys::<T>(b + 1),189 Vec::new(),190 )191 .unwrap();192193 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();194195 #[block]196 {197 assert_ok!(<CollatorSelection<T>>::add_invulnerable(198 root_origin,199 new_invulnerable.clone()200 ));201 }202203 assert_last_event::<T>(204 Event::InvulnerableAdded {205 invulnerable: new_invulnerable,206 }207 .into(),208 );209210 Ok(())211 }212213 #[benchmark]214 fn remove_invulnerable(b: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {215 register_validators::<T>(b);216 register_invulnerables::<T>(b);217218 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();219 let leaving = <Invulnerables<T>>::get().last().unwrap().clone();220 whitelist!(leaving);221222 #[block]223 {224 assert_ok!(<CollatorSelection<T>>::remove_invulnerable(225 root_origin,226 leaving.clone()227 ));228 }229230 assert_last_event::<T>(231 Event::InvulnerableRemoved {232 invulnerable: leaving,233 }234 .into(),235 );236237 Ok(())238 }239240 #[benchmark]241 fn get_license(c: Linear<1, MAX_COLLATORS>) -> Result<(), BenchmarkError> {242 register_validators::<T>(c);243 get_licenses::<T>(c);244245 let caller: T::AccountId = whitelisted_caller();246 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();247 T::Currency::set_balance(&caller, bond);248249 <session::Pallet<T>>::set_keys(250 RawOrigin::Signed(caller.clone()).into(),251 keys::<T>(c + 1),252 Vec::new(),253 )254 .unwrap();255256 #[extrinsic_call]257 _(RawOrigin::Signed(caller.clone()));258259 assert_last_event::<T>(260 Event::LicenseObtained {261 account_id: caller,262 deposit: bond / 2u32.into(),263 }264 .into(),265 );266267 Ok(())268 }269270 271 272 #[benchmark]273 fn onboard(c: Linear<2, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {274 let c = c - 1;275 register_validators::<T>(c);276 register_candidates::<T>(c);277278 let caller: T::AccountId = whitelisted_caller();279 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();280 T::Currency::set_balance(&caller, bond);281282 let origin = RawOrigin::Signed(caller.clone());283284 <session::Pallet<T>>::set_keys(origin.clone().into(), keys::<T>(c + 1), Vec::new())285 .unwrap();286287 assert_ok!(<CollatorSelection<T>>::get_license(origin.clone().into()));288289 #[extrinsic_call]290 _(origin);291292 assert_last_event::<T>(Event::CandidateAdded { account_id: caller }.into());293294 Ok(())295 }296297 298 #[benchmark]299 fn offboard(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {300 register_validators::<T>(c);301 register_candidates::<T>(c);302303 let leaving = <Candidates<T>>::get().last().unwrap().clone();304 whitelist!(leaving);305306 #[extrinsic_call]307 _(RawOrigin::Signed(leaving.clone()));308309 assert_last_event::<T>(310 Event::CandidateRemoved {311 account_id: leaving,312 }313 .into(),314 );315316 Ok(())317 }318319 320 #[benchmark]321 fn release_license(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {322 let bond = balance_unit::<T>();323324 register_validators::<T>(c);325 register_candidates::<T>(c);326327 let leaving = <Candidates<T>>::get().last().unwrap().clone();328 whitelist!(leaving);329330 #[extrinsic_call]331 _(RawOrigin::Signed(leaving.clone()));332333 assert_last_event::<T>(334 Event::LicenseReleased {335 account_id: leaving,336 deposit_returned: bond,337 }338 .into(),339 );340341 Ok(())342 }343344 345 #[benchmark]346 fn force_release_license(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {347 let bond = balance_unit::<T>();348349 register_validators::<T>(c);350 register_candidates::<T>(c);351352 let leaving = <Candidates<T>>::get().last().unwrap().clone();353 whitelist!(leaving);354 let origin = T::UpdateOrigin::try_successful_origin().unwrap();355356 #[block]357 {358 assert_ok!(<CollatorSelection<T>>::force_release_license(359 origin,360 leaving.clone()361 ));362 }363364 assert_last_event::<T>(365 Event::LicenseReleased {366 account_id: leaving,367 deposit_returned: bond,368 }369 .into(),370 );371372 Ok(())373 }374375 376 #[benchmark]377 fn note_author() -> Result<(), BenchmarkError> {378 T::Currency::set_balance(379 &<CollatorSelection<T>>::account_id(),380 balance_unit::<T>() * 4u32.into(),381 );382 let author = account("author", 0, SEED);383 let new_block: BlockNumberFor<T> = 10u32.into();384385 frame_system::Pallet::<T>::set_block_number(new_block);386 assert!(T::Currency::balance(&author) == 0u32.into());387388 #[block]389 {390 <CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone());391 }392393 assert!(T::Currency::balance(&author) > 0u32.into());394 assert_eq!(frame_system::Pallet::<T>::block_number(), new_block);395396 Ok(())397 }398399 400 #[benchmark]401 fn new_session(402 r: Linear<1, MAX_INVULNERABLES>,403 c: Linear<1, MAX_INVULNERABLES>,404 ) -> Result<(), BenchmarkError> {405 frame_system::Pallet::<T>::set_block_number(0u32.into());406407 register_validators::<T>(c);408 register_candidates::<T>(c);409410 let new_block: BlockNumberFor<T> = 1800u32.into();411 let zero_block: BlockNumberFor<T> = 0u32.into();412 let candidates = <Candidates<T>>::get();413414 let non_removals = c.saturating_sub(r);415416 for i in 0..c {417 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), zero_block);418 }419420 if non_removals > 0 {421 for i in 0..non_removals {422 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);423 }424 } else {425 for i in 0..c {426 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);427 }428 }429430 let pre_length = <Candidates<T>>::get().len();431432 frame_system::Pallet::<T>::set_block_number(new_block);433434 assert!(<Candidates<T>>::get().len() == c as usize);435436 #[block]437 {438 <CollatorSelection<T> as SessionManager<_>>::new_session(0);439 }440441 if c > r {442 assert!(<Candidates<T>>::get().len() < pre_length);443 } else {444 assert!(<Candidates<T>>::get().len() == pre_length);445 }446447 Ok(())448 }449450 impl_benchmark_test_suite!(451 CollatorSelection,452 crate::mock::new_test_ext(),453 crate::mock::Test,454 );455}