1234567891011121314151617181920212223242526272829303132333435use super::*;3637#[allow(unused)]38use crate::Pallet as CollatorSelection;39use frame_benchmarking::{account, benchmarks, impl_benchmark_test_suite, whitelisted_caller};40use frame_support::{41 assert_ok,42 codec::Decode,43 traits::{Currency, EnsureOrigin, Get},44};45use frame_system::{EventRecord, RawOrigin};46use pallet_authorship::EventHandler;47use pallet_session::{self as session, SessionManager};48use pallet_configuration::{49 self as configuration, BalanceOf,50 CollatorSelectionDesiredCollatorsOverride as DesiredCollators,51 CollatorSelectionLicenseBondOverride as LicenseBond,52};53use sp_std::prelude::*;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::make_free_balance_be(&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 + configuration::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 + configuration::Config>(count: u32) {131 let candidates = (0..count)132 .map(|c| account("candidate", c, SEED))133 .collect::<Vec<_>>();134 assert!(135 <LicenseBond<T>>::get() > 0u32.into(),136 "Bond cannot be zero!"137 );138139 for who in candidates {140 T::Currency::make_free_balance_be(&who, <LicenseBond<T>>::get() * 2u32.into());141 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();142 <CollatorSelection<T>>::onboard(RawOrigin::Signed(who).into()).unwrap();143 }144}145146fn get_licenses<T: Config + configuration::Config>(count: u32) {147 let candidates = (0..count)148 .map(|c| account("candidate", c, SEED))149 .collect::<Vec<_>>();150 assert!(151 <LicenseBond<T>>::get() > 0u32.into(),152 "Bond cannot be zero!"153 );154155 for who in candidates {156 T::Currency::make_free_balance_be(&who, <LicenseBond<T>>::get() * 2u32.into());157 <CollatorSelection<T>>::get_license(RawOrigin::Signed(who.clone()).into()).unwrap();158 }159}160161162163fn balance_unit<T: Config>() -> <T::Currency as Currency<T::AccountId>>::Balance {164 200u32.into()165}166167168const INITIAL_INVULNERABLES: u32 = 2;169170benchmarks! {171 where_clause { where T: pallet_authorship::Config + session::Config + configuration::Config }172173 174 175 176 add_invulnerable {177 let b in 1 .. T::MaxCollators::get() - INITIAL_INVULNERABLES - 1;178 register_validators::<T>(b);179 register_invulnerables::<T>(b);180181 182183 let new_invulnerable: T::AccountId = whitelisted_caller();184 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();185 T::Currency::make_free_balance_be(&new_invulnerable, bond.clone());186187 <session::Pallet<T>>::set_keys(188 RawOrigin::Signed(new_invulnerable.clone()).into(),189 keys::<T>(b + 1),190 Vec::new()191 ).unwrap();192193 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();194 }: {195 assert_ok!(196 <CollatorSelection<T>>::add_invulnerable(root_origin, new_invulnerable.clone())197 );198 }199 verify {200 assert_last_event::<T>(Event::InvulnerableAdded{invulnerable: new_invulnerable}.into());201 }202203 remove_invulnerable {204 let b in 1 .. T::MaxCollators::get() - INITIAL_INVULNERABLES - 1;205 register_validators::<T>(b);206 register_invulnerables::<T>(b);207208 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();209 let leaving = <Invulnerables<T>>::get().last().unwrap().clone();210 whitelist!(leaving);211 }: {212 assert_ok!(213 <CollatorSelection<T>>::remove_invulnerable(root_origin, leaving.clone())214 );215 }216 verify {217 assert_last_event::<T>(Event::InvulnerableRemoved{invulnerable: leaving}.into());218 }219220 get_license {221 let c in 1 .. T::MaxCollators::get() - 1;222223 <LicenseBond<T>>::put(balance_unit::<T>());224225 register_validators::<T>(c);226 get_licenses::<T>(c);227228 let caller: T::AccountId = whitelisted_caller();229 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();230 T::Currency::make_free_balance_be(&caller, bond.clone());231232 <session::Pallet<T>>::set_keys(233 RawOrigin::Signed(caller.clone()).into(),234 keys::<T>(c + 1),235 Vec::new()236 ).unwrap();237238 }: _(RawOrigin::Signed(caller.clone()))239 verify {240 assert_last_event::<T>(Event::LicenseObtained{account_id: caller, deposit: bond / 2u32.into()}.into());241 }242243 244 245 onboard {246 let c in 1 .. T::MaxCollators::get() - INITIAL_INVULNERABLES - 1;247248 <LicenseBond<T>>::put(balance_unit::<T>());249 <DesiredCollators<T>>::put(c + INITIAL_INVULNERABLES + 1);250251 register_validators::<T>(c);252 register_candidates::<T>(c);253254 let caller: T::AccountId = whitelisted_caller();255 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();256 T::Currency::make_free_balance_be(&caller, bond.clone());257258 let origin = RawOrigin::Signed(caller.clone());259260 <session::Pallet<T>>::set_keys(261 origin.clone().into(),262 keys::<T>(c + 1),263 Vec::new()264 ).unwrap();265266 assert_ok!(267 <CollatorSelection<T>>::get_license(origin.clone().into())268 );269 }: _(origin)270 verify {271 assert_last_event::<T>(Event::CandidateAdded{account_id: caller}.into());272 }273274 275 offboard {276 let c in 1 .. T::MaxCollators::get();277 <LicenseBond<T>>::put(balance_unit::<T>());278 <DesiredCollators<T>>::put(c + INITIAL_INVULNERABLES);279280 register_validators::<T>(c);281 register_candidates::<T>(c);282283 let leaving = <Candidates<T>>::get().last().unwrap().clone();284 whitelist!(leaving);285 }: _(RawOrigin::Signed(leaving.clone()))286 verify {287 assert_last_event::<T>(Event::CandidateRemoved{account_id: leaving}.into());288 }289290 291 release_license {292 let c in 1 .. T::MaxCollators::get();293 let bond = balance_unit::<T>();294 <LicenseBond<T>>::put(bond);295 <DesiredCollators<T>>::put(c + INITIAL_INVULNERABLES);296297 register_validators::<T>(c);298 register_candidates::<T>(c);299300 let leaving = <Candidates<T>>::get().last().unwrap().clone();301 whitelist!(leaving);302 }: _(RawOrigin::Signed(leaving.clone()))303 verify {304 assert_last_event::<T>(Event::LicenseReleased{account_id: leaving, deposit_returned: bond}.into());305 }306307 308 force_release_license {309 let c in 1 .. T::MaxCollators::get();310 let bond = balance_unit::<T>();311 <LicenseBond<T>>::put(bond);312 <DesiredCollators<T>>::put(c + INITIAL_INVULNERABLES);313314 register_validators::<T>(c);315 register_candidates::<T>(c);316317 let leaving = <Candidates<T>>::get().last().unwrap().clone();318 whitelist!(leaving);319 let origin = T::UpdateOrigin::try_successful_origin().unwrap();320 }: {321 assert_ok!(322 <CollatorSelection<T>>::force_release_license(origin, leaving.clone())323 );324 }325 verify {326 assert_last_event::<T>(Event::LicenseReleased{account_id: leaving, deposit_returned: bond}.into());327 }328329 330 note_author {331 <LicenseBond<T>>::put(balance_unit::<T>());332 T::Currency::make_free_balance_be(333 &<CollatorSelection<T>>::account_id(),334 balance_unit::<T>() * 4u32.into(),335 );336 let author = account("author", 0, SEED);337 let new_block: T::BlockNumber = 10u32.into();338339 frame_system::Pallet::<T>::set_block_number(new_block);340 assert!(T::Currency::free_balance(&author) == 0u32.into());341 }: {342 <CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone())343 } verify {344 assert!(T::Currency::free_balance(&author) > 0u32.into());345 assert_eq!(frame_system::Pallet::<T>::block_number(), new_block);346 }347348 349 new_session {350 let r in 1 .. T::MaxCollators::get();351 let c in 1 .. T::MaxCollators::get();352353 <LicenseBond<T>>::put(balance_unit::<T>());354 <DesiredCollators<T>>::put(c + INITIAL_INVULNERABLES);355 frame_system::Pallet::<T>::set_block_number(0u32.into());356357 register_validators::<T>(c);358 register_candidates::<T>(c);359360 let new_block: T::BlockNumber = 1800u32.into();361 let zero_block: T::BlockNumber = 0u32.into();362 let candidates = <Candidates<T>>::get();363364 let non_removals = c.saturating_sub(r);365366 for i in 0..c {367 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), zero_block);368 }369370 if non_removals > 0 {371 for i in 0..non_removals {372 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);373 }374 } else {375 for i in 0..c {376 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);377 }378 }379380 let pre_length = <Candidates<T>>::get().len();381382 frame_system::Pallet::<T>::set_block_number(new_block);383384 assert!(<Candidates<T>>::get().len() == c as usize);385 }: {386 <CollatorSelection<T> as SessionManager<_>>::new_session(0)387 } verify {388 if c > r {389 assert!(<Candidates<T>>::get().len() < pre_length);390 } else {391 assert!(<Candidates<T>>::get().len() == pre_length);392 }393 }394}395396impl_benchmark_test_suite!(397 CollatorSelection,398 crate::mock::new_test_ext(),399 crate::mock::Test,400);