1234567891011121314151617181920212223242526272829303132333435use frame_benchmarking::v2::{36 account, benchmarks, impl_benchmark_test_suite, whitelisted_caller, BenchmarkError,37};38use frame_support::{39 assert_ok,40 traits::{41 fungible::{Inspect, Mutate},42 EnsureOrigin, Get,43 },44};45use frame_system::{pallet_prelude::*, EventRecord, RawOrigin};46use pallet_authorship::EventHandler;47use pallet_session::{self as session, SessionManager};48use parity_scale_codec::Decode;49use sp_std::prelude::*;5051use super::*;52#[allow(unused)]53use crate::{BalanceOf, Pallet as CollatorSelection};5455const SEED: u32 = 0;565758macro_rules! whitelist {59 ($acc:ident) => {60 frame_benchmarking::benchmarking::add_to_whitelist(61 frame_system::Account::<T>::hashed_key_for(&$acc).into(),62 );63 };64}6566fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {67 let events = frame_system::Pallet::<T>::events();68 let system_event: <T as frame_system::Config>::RuntimeEvent = generic_event.into();69 70 let EventRecord { event, .. } = &events[events.len() - 1];71 assert_eq!(event, &system_event);72}7374fn create_funded_user<T: Config>(75 string: &'static str,76 n: u32,77 balance_factor: u32,78) -> T::AccountId {79 let user = account(string, n, SEED);80 let balance = balance_unit::<T>() * balance_factor.into();81 let _ = T::Currency::set_balance(&user, balance);82 user83}8485fn keys<T: Config + session::Config>(c: u32) -> <T as session::Config>::Keys {86 use rand::{RngCore, SeedableRng};8788 let keys = {89 let mut keys = [0u8; 128];9091 if c > 0 {92 let mut rng = rand::rngs::StdRng::seed_from_u64(c as u64);93 rng.fill_bytes(&mut keys);94 }9596 keys97 };9899 Decode::decode(&mut &keys[..]).unwrap()100}101102fn validator<T: Config + session::Config>(c: u32) -> (T::AccountId, <T as session::Config>::Keys) {103 (create_funded_user::<T>("candidate", c, 1000), keys::<T>(c))104}105106fn register_validators<T: Config + session::Config>(count: u32) -> Vec<T::AccountId> {107 let validators = (0..count).map(|c| validator::<T>(c)).collect::<Vec<_>>();108109 for (who, keys) in validators.clone() {110 <session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();111 }112113 validators.into_iter().map(|(who, _)| who).collect()114}115116fn register_invulnerables<T: Config>(count: u32) {117 let candidates = (0..count)118 .map(|c| account("candidate", c, SEED))119 .collect::<Vec<_>>();120121 for who in candidates {122 <CollatorSelection<T>>::add_invulnerable(123 T::UpdateOrigin::try_successful_origin().unwrap(),124 who,125 )126 .unwrap();127 }128}129130fn register_candidates<T: Config>(count: u32) {131 let candidates = (0..count)132 .map(|c| account("candidate", c, SEED))133 .collect::<Vec<_>>();134 assert!(T::LicenseBond::get() > 0u32.into(), "Bond cannot be zero!");135136 for who in candidates {137 T::Currency::set_balance(&who, T::LicenseBond::get() * 2u32.into());138 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();139 <CollatorSelection<T>>::onboard(RawOrigin::Signed(who).into()).unwrap();140 }141}142143fn get_licenses<T: Config>(count: u32) {144 let candidates = (0..count)145 .map(|c| account("candidate", c, SEED))146 .collect::<Vec<_>>();147 assert!(T::LicenseBond::get() > 0u32.into(), "Bond cannot be zero!");148149 for who in candidates {150 T::Currency::set_balance(&who, T::LicenseBond::get() * 2u32.into());151 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();152 }153}154155156157fn balance_unit<T: Config>() -> BalanceOf<T> {158 T::LicenseBond::get()159}160161162const INITIAL_INVULNERABLES: u32 = 2;163164#[benchmarks(where T: Config + pallet_authorship::Config + session::Config)]165mod benchmarks {166 use super::*;167 const MAX_COLLATORS: u32 = 10;168 const MAX_INVULNERABLES: u32 = MAX_COLLATORS - INITIAL_INVULNERABLES;169170 171 172 173 #[benchmark]174 fn add_invulnerable<T>(b: Linear<1, MAX_COLLATORS>) -> Result<(), BenchmarkError> {175 register_validators::<T>(b);176 register_invulnerables::<T>(b);177178 179180 let new_invulnerable: T::AccountId = whitelisted_caller();181 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();182 <T as Config>::Currency::set_balance(&new_invulnerable, bond);183184 <session::Pallet<T>>::set_keys(185 RawOrigin::Signed(new_invulnerable.clone()).into(),186 keys::<T>(b + 1),187 Vec::new(),188 )189 .unwrap();190191 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();192193 #[block]194 {195 assert_ok!(<CollatorSelection<T>>::add_invulnerable(196 root_origin,197 new_invulnerable.clone()198 ));199 }200201 assert_last_event::<T>(202 Event::InvulnerableAdded {203 invulnerable: new_invulnerable,204 }205 .into(),206 );207208 Ok(())209 }210211 #[benchmark]212 fn remove_invulnerable(b: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {213 register_validators::<T>(b);214 register_invulnerables::<T>(b);215216 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();217 let leaving = <Invulnerables<T>>::get().last().unwrap().clone();218 whitelist!(leaving);219220 #[block]221 {222 assert_ok!(<CollatorSelection<T>>::remove_invulnerable(223 root_origin,224 leaving.clone()225 ));226 }227228 assert_last_event::<T>(229 Event::InvulnerableRemoved {230 invulnerable: leaving,231 }232 .into(),233 );234235 Ok(())236 }237238 #[benchmark]239 fn get_license(c: Linear<1, MAX_COLLATORS>) -> Result<(), BenchmarkError> {240 register_validators::<T>(c);241 get_licenses::<T>(c);242243 let caller: T::AccountId = whitelisted_caller();244 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();245 T::Currency::set_balance(&caller, bond);246247 <session::Pallet<T>>::set_keys(248 RawOrigin::Signed(caller.clone()).into(),249 keys::<T>(c + 1),250 Vec::new(),251 )252 .unwrap();253254 #[extrinsic_call]255 _(RawOrigin::Signed(caller.clone()));256257 assert_last_event::<T>(258 Event::LicenseObtained {259 account_id: caller,260 deposit: bond / 2u32.into(),261 }262 .into(),263 );264265 Ok(())266 }267268 269 270 #[benchmark]271 fn onboard(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {272 register_validators::<T>(c);273 register_candidates::<T>(c);274275 let caller: T::AccountId = whitelisted_caller();276 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();277 T::Currency::set_balance(&caller, bond);278279 let origin = RawOrigin::Signed(caller.clone());280281 <session::Pallet<T>>::set_keys(origin.clone().into(), keys::<T>(c + 1), Vec::new())282 .unwrap();283284 assert_ok!(<CollatorSelection<T>>::get_license(origin.clone().into()));285286 #[extrinsic_call]287 _(origin);288289 assert_last_event::<T>(Event::CandidateAdded { account_id: caller }.into());290291 Ok(())292 }293294 295 #[benchmark]296 fn offboard(c: Linear<0, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {297 let c = c + 1;298299 register_validators::<T>(c);300 register_candidates::<T>(c);301302 let leaving = <Candidates<T>>::get().last().unwrap().clone();303 whitelist!(leaving);304305 #[extrinsic_call]306 _(RawOrigin::Signed(leaving.clone()));307308 assert_last_event::<T>(309 Event::CandidateRemoved {310 account_id: leaving,311 }312 .into(),313 );314315 Ok(())316 }317318 319 #[benchmark]320 fn release_license(c: Linear<0, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {321 let c = c + 1;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<0, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {347 let c = c + 1;348 let bond = balance_unit::<T>();349350 register_validators::<T>(c);351 register_candidates::<T>(c);352353 let leaving = <Candidates<T>>::get().last().unwrap().clone();354 whitelist!(leaving);355 let origin = T::UpdateOrigin::try_successful_origin().unwrap();356357 #[block]358 {359 assert_ok!(<CollatorSelection<T>>::force_release_license(360 origin,361 leaving.clone()362 ));363 }364365 assert_last_event::<T>(366 Event::LicenseReleased {367 account_id: leaving,368 deposit_returned: bond,369 }370 .into(),371 );372373 Ok(())374 }375376 377 #[benchmark]378 fn note_author() -> Result<(), BenchmarkError> {379 T::Currency::set_balance(380 &<CollatorSelection<T>>::account_id(),381 balance_unit::<T>() * 4u32.into(),382 );383 let author = account("author", 0, SEED);384 let new_block: BlockNumberFor<T> = 10u32.into();385386 frame_system::Pallet::<T>::set_block_number(new_block);387 assert!(T::Currency::balance(&author) == 0u32.into());388389 #[block]390 {391 <CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone());392 }393394 assert!(T::Currency::balance(&author) > 0u32.into());395 assert_eq!(frame_system::Pallet::<T>::block_number(), new_block);396397 Ok(())398 }399400 401 #[benchmark]402 fn new_session(403 r: Linear<0, MAX_INVULNERABLES>,404 c: Linear<0, MAX_INVULNERABLES>,405 ) -> Result<(), BenchmarkError> {406 let r = r + 1;407 let c = c + 1;408409 frame_system::Pallet::<T>::set_block_number(0u32.into());410411 register_validators::<T>(c);412 register_candidates::<T>(c);413414 let new_block: BlockNumberFor<T> = 1800u32.into();415 let zero_block: BlockNumberFor<T> = 0u32.into();416 let candidates = <Candidates<T>>::get();417418 let non_removals = c.saturating_sub(r);419420 for i in 0..c {421 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), zero_block);422 }423424 if non_removals > 0 {425 for i in 0..non_removals {426 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);427 }428 } else {429 for i in 0..c {430 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);431 }432 }433434 let pre_length = <Candidates<T>>::get().len();435436 frame_system::Pallet::<T>::set_block_number(new_block);437438 assert!(<Candidates<T>>::get().len() == c as usize);439440 #[block]441 {442 <CollatorSelection<T> as SessionManager<_>>::new_session(0);443 }444445 if c > r {446 assert!(<Candidates<T>>::get().len() < pre_length);447 } else {448 assert!(<Candidates<T>>::get().len() == pre_length);449 }450451 Ok(())452 }453454 impl_benchmark_test_suite!(455 CollatorSelection,456 crate::mock::new_test_ext(),457 crate::mock::Test,458 );459}