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<2, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {175 let b = b - 1;176 register_validators::<T>(b);177 register_invulnerables::<T>(b);178179 180181 let new_invulnerable: T::AccountId = whitelisted_caller();182 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();183 <T as Config>::Currency::set_balance(&new_invulnerable, bond);184185 <session::Pallet<T>>::set_keys(186 RawOrigin::Signed(new_invulnerable.clone()).into(),187 keys::<T>(b + 1),188 Vec::new(),189 )190 .unwrap();191192 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();193194 #[block]195 {196 assert_ok!(<CollatorSelection<T>>::add_invulnerable(197 root_origin,198 new_invulnerable.clone()199 ));200 }201202 assert_last_event::<T>(203 Event::InvulnerableAdded {204 invulnerable: new_invulnerable,205 }206 .into(),207 );208209 Ok(())210 }211212 #[benchmark]213 fn remove_invulnerable(b: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {214 register_validators::<T>(b);215 register_invulnerables::<T>(b);216217 let root_origin = T::UpdateOrigin::try_successful_origin().unwrap();218 let leaving = <Invulnerables<T>>::get().last().unwrap().clone();219 whitelist!(leaving);220221 #[block]222 {223 assert_ok!(<CollatorSelection<T>>::remove_invulnerable(224 root_origin,225 leaving.clone()226 ));227 }228229 assert_last_event::<T>(230 Event::InvulnerableRemoved {231 invulnerable: leaving,232 }233 .into(),234 );235236 Ok(())237 }238239 #[benchmark]240 fn get_license(c: Linear<1, MAX_COLLATORS>) -> Result<(), BenchmarkError> {241 register_validators::<T>(c);242 get_licenses::<T>(c);243244 let caller: T::AccountId = whitelisted_caller();245 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();246 T::Currency::set_balance(&caller, bond);247248 <session::Pallet<T>>::set_keys(249 RawOrigin::Signed(caller.clone()).into(),250 keys::<T>(c + 1),251 Vec::new(),252 )253 .unwrap();254255 #[extrinsic_call]256 _(RawOrigin::Signed(caller.clone()));257258 assert_last_event::<T>(259 Event::LicenseObtained {260 account_id: caller,261 deposit: bond / 2u32.into(),262 }263 .into(),264 );265266 Ok(())267 }268269 270 271 #[benchmark]272 fn onboard(c: Linear<2, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {273 let c = c - 1;274 register_validators::<T>(c);275 register_candidates::<T>(c);276277 let caller: T::AccountId = whitelisted_caller();278 let bond: BalanceOf<T> = balance_unit::<T>() * 2u32.into();279 T::Currency::set_balance(&caller, bond);280281 let origin = RawOrigin::Signed(caller.clone());282283 <session::Pallet<T>>::set_keys(origin.clone().into(), keys::<T>(c + 1), Vec::new())284 .unwrap();285286 assert_ok!(<CollatorSelection<T>>::get_license(origin.clone().into()));287288 #[extrinsic_call]289 _(origin);290291 assert_last_event::<T>(Event::CandidateAdded { account_id: caller }.into());292293 Ok(())294 }295296 297 #[benchmark]298 fn offboard(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {299 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<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {321 let bond = balance_unit::<T>();322323 register_validators::<T>(c);324 register_candidates::<T>(c);325326 let leaving = <Candidates<T>>::get().last().unwrap().clone();327 whitelist!(leaving);328329 #[extrinsic_call]330 _(RawOrigin::Signed(leaving.clone()));331332 assert_last_event::<T>(333 Event::LicenseReleased {334 account_id: leaving,335 deposit_returned: bond,336 }337 .into(),338 );339340 Ok(())341 }342343 344 #[benchmark]345 fn force_release_license(c: Linear<1, MAX_INVULNERABLES>) -> Result<(), BenchmarkError> {346 let bond = balance_unit::<T>();347348 register_validators::<T>(c);349 register_candidates::<T>(c);350351 let leaving = <Candidates<T>>::get().last().unwrap().clone();352 whitelist!(leaving);353 let origin = T::UpdateOrigin::try_successful_origin().unwrap();354355 #[block]356 {357 assert_ok!(<CollatorSelection<T>>::force_release_license(358 origin,359 leaving.clone()360 ));361 }362363 assert_last_event::<T>(364 Event::LicenseReleased {365 account_id: leaving,366 deposit_returned: bond,367 }368 .into(),369 );370371 Ok(())372 }373374 375 #[benchmark]376 fn note_author() -> Result<(), BenchmarkError> {377 T::Currency::set_balance(378 &<CollatorSelection<T>>::account_id(),379 balance_unit::<T>() * 4u32.into(),380 );381 let author = account("author", 0, SEED);382 let new_block: BlockNumberFor<T> = 10u32.into();383384 frame_system::Pallet::<T>::set_block_number(new_block);385 assert!(T::Currency::balance(&author) == 0u32.into());386387 #[block]388 {389 <CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone());390 }391392 assert!(T::Currency::balance(&author) > 0u32.into());393 assert_eq!(frame_system::Pallet::<T>::block_number(), new_block);394395 Ok(())396 }397398 399 #[benchmark]400 fn new_session(401 r: Linear<1, MAX_INVULNERABLES>,402 c: Linear<1, MAX_INVULNERABLES>,403 ) -> Result<(), BenchmarkError> {404 frame_system::Pallet::<T>::set_block_number(0u32.into());405406 register_validators::<T>(c);407 register_candidates::<T>(c);408409 let new_block: BlockNumberFor<T> = 1800u32.into();410 let zero_block: BlockNumberFor<T> = 0u32.into();411 let candidates = <Candidates<T>>::get();412413 let non_removals = c.saturating_sub(r);414415 for i in 0..c {416 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), zero_block);417 }418419 if non_removals > 0 {420 for i in 0..non_removals {421 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);422 }423 } else {424 for i in 0..c {425 <LastAuthoredBlock<T>>::insert(candidates[i as usize].clone(), new_block);426 }427 }428429 let pre_length = <Candidates<T>>::get().len();430431 frame_system::Pallet::<T>::set_block_number(new_block);432433 assert!(<Candidates<T>>::get().len() == c as usize);434435 #[block]436 {437 <CollatorSelection<T> as SessionManager<_>>::new_session(0);438 }439440 if c > r {441 assert!(<Candidates<T>>::get().len() < pre_length);442 } else {443 assert!(<Candidates<T>>::get().len() == pre_length);444 }445446 Ok(())447 }448449 impl_benchmark_test_suite!(450 CollatorSelection,451 crate::mock::new_test_ext(),452 crate::mock::Test,453 );454}