Improve paras runtime BenchBuilder api

runtime/parachains/src/builder.rs

@@ -63,13 +63,31 @@ fn byte32_slice_from(n: u32) -> [u8; 32] {
 
 /// Paras inherent `enter` benchmark scenario builder.
 pub(crate) struct BenchBuilder<T: paras_inherent::Config> {
+	/// Active validators. Validators should be declared prior to all other setup.
 	validators: Option<Vec<ValidatorId>>,
+	/// Starting block number; we expect it to get incremented on session setup.
 	block_number: T::BlockNumber,
+	/// Starting session; we expect it to get incremented on session setup.
 	session: SessionIndex,
+	/// Session we want the scenario to take place in. We will roll to this session.
 	target_session: u32,
+	/// Optionally set the max validators per core; otherwise uses the configuration value.
 	max_validators_per_core: Option<u32>,
+	/// Optionally set the max validators; otherwise uses the configuration value.
 	max_validators: Option<u32>,
+	/// Optionally set the number of dispute statements for each candidate.
 	dispute_statements: BTreeMap<u32, u32>,
+	/// Session index of for each dispute. Index of slice corresponds to a core,
+	/// which is offset by the number of entries for `backed_and_concluding_cores`. I.E. if
+	/// `backed_and_concluding_cores` has 3 entries, the first index of `dispute_sessions`
+	/// will correspond to core index 3. There must be one entry for each core with a dispute
+	/// statement set.
+	dispute_sessions: Vec<u32>,
+	/// Map from core seed to number of validity votes.
+	backed_and_concluding_cores: BTreeMap<u32, u32>,
+	/// Make every candidate include a code upgrade by setting this to `Some` where the interior
+	/// value is the byte length of the new code.
+	code_upgrade: Option<u32>,
 	_phantom: sp_std::marker::PhantomData<T>,
 }
 
@@ -86,24 +104,48 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 	/// of the functions in this implementation.
 	pub(crate) fn new() -> Self {
 		BenchBuilder {
-			// Validators should be declared prior to all other setup.
 			validators: None,
-			// Starting block number; we expect it to get incremented on session setup.
 			block_number: Zero::zero(),
-			// Starting session; we expect it to get incremented on session setup.
 			session: SessionIndex::from(0u32),
-			// Session we want the scenario to take place in. We roll to to this session.
 			target_session: 2u32,
-			// Optionally set the max validators per core; otherwise uses the configuration value.
 			max_validators_per_core: None,
-			// Optionally set the max validators; otherwise uses the configuration value.
 			max_validators: None,
-			// Optionally set the number of dispute statements for each candidate,
 			dispute_statements: BTreeMap::new(),
+			dispute_sessions: Default::default(),
+			backed_and_concluding_cores: Default::default(),
+			code_upgrade: None,
 			_phantom: sp_std::marker::PhantomData::<T>,
 		}
 	}
 
+	/// Set the session index for each dispute statement set (in other words, set the session the
+	/// the dispute statement set's relay chain block is from). Indexes of `dispute_sessions`
+	/// correspond to a core, which is offset by the number of entries for
+	/// `backed_and_concluding_cores`. I.E. if `backed_and_concluding_cores` cores has 3 entries,
+	/// the first index of `dispute_sessions` will correspond to core index 3.
+	///
+	/// Note that there must be an entry for each core with a dispute statement set.
+	pub(crate) fn set_dispute_sessions(mut self, dispute_sessions: impl AsRef<[u32]>) -> Self {
+		self.dispute_sessions = dispute_sessions.as_ref().to_vec();
+		self
+	}
+
+	/// Set a map from core/para id seed to number of validity votes.
+	pub(crate) fn set_backed_and_concluding_cores(
+		mut self,
+		backed_and_concluding_cores: BTreeMap<u32, u32>,
+	) -> Self {
+		self.backed_and_concluding_cores = backed_and_concluding_cores;
+		self
+	}
+
+	/// Set to include a code upgrade for all backed candidates. The value will be the byte length
+	/// of the code.
+	pub(crate) fn set_code_upgrade(mut self, code_upgrade: impl Into<Option<u32>>) -> Self {
+		self.code_upgrade = code_upgrade.into();
+		self
+	}
+
 	/// Mock header.
 	pub(crate) fn header(block_number: T::BlockNumber) -> T::Header {
 		T::Header::new(
@@ -133,6 +175,7 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		self.max_validators.unwrap_or(Self::fallback_max_validators())
 	}
 
+	/// Set the maximum number of active validators.
 	#[cfg(not(feature = "runtime-benchmarks"))]
 	pub(crate) fn set_max_validators(mut self, n: u32) -> Self {
 		self.max_validators = Some(n);
@@ -145,16 +188,18 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		configuration::Pallet::<T>::config().max_validators_per_core.unwrap_or(5)
 	}
 
-	/// Specify a mapping of core index, para id, group index seed to the number of dispute statements for the
-	/// corresponding dispute statement set. Note that if the number of disputes is not specified it fallbacks
-	/// to having a dispute per every validator. Additionally, an entry is not guaranteed to have a dispute - it
-	/// must line up with the cores marked as disputed as defined in `Self::Build`.
+	/// Specify a mapping of core index/ para id to the number of dispute statements for the
+	/// corresponding dispute statement set. Note that if the number of disputes is not specified
+	/// it fallbacks to having a dispute per every validator. Additionally, an entry is not
+	/// guaranteed to have a dispute - it must line up with the cores marked as disputed as defined
+	/// in `Self::Build`.
 	#[cfg(not(feature = "runtime-benchmarks"))]
 	pub(crate) fn set_dispute_statements(mut self, m: BTreeMap<u32, u32>) -> Self {
 		self.dispute_statements = m;
 		self
 	}
 
+	/// Get the maximum number of validators per core.
 	fn max_validators_per_core(&self) -> u32 {
 		self.max_validators_per_core.unwrap_or(Self::fallback_max_validators_per_core())
 	}
@@ -166,17 +211,18 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		self
 	}
 
-	/// Maximum number of cores we expect from this configuration.
+	/// Get the maximum number of cores we expect from this configuration.
 	pub(crate) fn max_cores(&self) -> u32 {
 		self.max_validators() / self.max_validators_per_core()
 	}
 
-	/// Minimum number of validity votes in order for a backed candidate to be included.
+	/// Get the minimum number of validity votes in order for a backed candidate to be included.
 	#[cfg(feature = "runtime-benchmarks")]
 	pub(crate) fn fallback_min_validity_votes() -> u32 {
 		(Self::fallback_max_validators() / 2) + 1
 	}
 
+	/// Create para id, core index, and grab the associated group index from the scheduler pallet.
 	fn create_indexes(&self, seed: u32) -> (ParaId, CoreIndex, GroupIndex) {
 		let para_id = ParaId::from(seed);
 		let core_idx = CoreIndex(seed);
@@ -186,6 +232,7 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		(para_id, core_idx, group_idx)
 	}
 
+	/// Create a mock of `CandidatePendingAvailability`.
 	fn candidate_availability_mock(
 		group_idx: GroupIndex,
 		core_idx: CoreIndex,
@@ -204,6 +251,11 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		)
 	}
 
+	/// Add `CandidatePendingAvailability` and `CandidateCommitments` to the relevant storage items.
+	///
+	/// NOTE: the default `CandidateCommitments` used does not include any data that would lead to
+	/// heavy code paths in `enact_candidate`. But enact_candidates does return a weight which will
+	/// get taken into account.
 	fn add_availability(
 		para_id: ParaId,
 		core_idx: CoreIndex,
@@ -217,14 +269,13 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 			candidate_hash,
 			availability_votes,
 		);
-		// NOTE: commitments does not include any data that would lead to heavy code
-		// paths in `enact_candidate`. But enact_candidates does return a weight which will get
-		// taken into account.
 		let commitments = CandidateCommitments::<u32>::default();
 		inclusion::PendingAvailability::<T>::insert(para_id, candidate_availability);
 		inclusion::PendingAvailabilityCommitments::<T>::insert(&para_id, commitments);
 	}
 
+	/// Create an `AvailabilityBitfield` where `concluding` is a map where each key is a core index
+	/// that is concluding and `cores` is the total number of cores in the system.
 	fn availability_bitvec(concluding: &BTreeMap<u32, u32>, cores: u32) -> AvailabilityBitfield {
 		let mut bitfields = bitvec::bitvec![bitvec::order::Lsb0, u8; 0; 0];
 		for i in 0..cores {
@@ -238,6 +289,8 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		bitfields.into()
 	}
 
+	/// Run to block number `to`, calling `initializer` `on_initialize` and `on_finalize` along the
+	/// way.
 	fn run_to_block(to: u32) {
 		let to = to.into();
 		while frame_system::Pallet::<T>::block_number() < to {
@@ -250,6 +303,10 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		}
 	}
 
+	/// Register `cores` count of parachains.
+	///
+	/// Note that this must be called at least 2 sessions before the target session as there is a
+	/// n+2 session delay for the scheduled actions to take effect.
 	fn setup_para_ids(cores: u32) {
 		// make sure parachains exist prior to session change.
 		for i in 0..cores {
@@ -288,6 +345,7 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		}
 	}
 
+	/// Create a bitvec of `validators` length with all yes votes.
 	fn validator_availability_votes_yes(validators: usize) -> BitVec<bitvec::order::Lsb0, u8> {
 		// every validator confirms availability.
 		bitvec::bitvec![bitvec::order::Lsb0, u8; 1; validators as usize]
@@ -338,6 +396,10 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		self
 	}
 
+	/// Create a `UncheckedSigned<AvailabilityBitfield> for each validator where each core in
+	/// `concluding_cores` is fully available. Additionally set up storage such that each
+	/// `concluding_cores`is pending becoming fully available so the generated bitfields will be
+	///  to the cores successfully being freed from the candidates being marked as available.
 	fn create_availability_bitfields(
 		&self,
 		concluding_cores: &BTreeMap<u32, u32>,
@@ -364,8 +426,7 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 			.collect();
 
 		for (seed, _) in concluding_cores.iter() {
-			// make sure the candidates that are concluding by becoming available are marked as
-			// pending availability.
+			// make sure the candidates that will be concluding are marked as pending availability.
 			let (para_id, core_idx, group_idx) = self.create_indexes(seed.clone());
 			Self::add_availability(
 				para_id,
@@ -483,19 +544,25 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 			.collect()
 	}
 
+	/// Fill cores `start..last` with dispute statement sets. The statement sets will have 3/4th of
+	/// votes be valid, and 1/4th of votes be invalid.
 	fn create_disputes_with_no_spam(
 		&self,
 		start: u32,
 		last: u32,
-		dispute_sessions: &[u32],
+		dispute_sessions: impl AsRef<[u32]>,
 	) -> Vec<DisputeStatementSet> {
 		let validators =
 			self.validators.as_ref().expect("must have some validators prior to calling");
 
+		let dispute_sessions = dispute_sessions.as_ref();
 		(start..last)
 			.map(|seed| {
-				let session =
-					dispute_sessions.get(seed as usize).cloned().unwrap_or(self.target_session);
+				let dispute_session_idx = (seed - start) as usize;
+				let session = dispute_sessions
+					.get(dispute_session_idx)
+					.cloned()
+					.unwrap_or(self.target_session);
 
 				let (para_id, core_idx, group_idx) = self.create_indexes(seed);
 				let candidate_hash = CandidateHash(H256::from(byte32_slice_from(seed)));
@@ -537,46 +604,39 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 
 	/// Build a scenario for testing or benchmarks.
 	///
-	/// - `backed_and_concluding_cores`: Map from core/para id/group index seed to number of
-	/// validity votes.
-	/// - `dispute_sessions`: Session index of for each dispute. Index of slice corresponds to core.
-	/// The length of this must equal total cores used. Seed index for disputes starts at
-	/// `backed_and_concluding_cores.len()`, so `dispute_sessions` needs to be left padded by
-	/// `backed_and_concluding_cores.len()` values which effectively get ignored.
-	/// TODO we should fix this.
-	pub(crate) fn build(
-		self,
-		backed_and_concluding_cores: BTreeMap<u32, u32>,
-		dispute_sessions: &[u32],
-		includes_code_upgrade: Option<u32>,
-	) -> Bench<T> {
+	/// Note that this API only allows building scenarios where the `backed_and_concluding_cores`
+	/// are mutually exclusive with the cores for disputes. So
+	/// `backed_and_concluding_cores.len() + dispute_sessions.len()` must be less than the max
+	/// number of cores.
+	pub(crate) fn build(self) -> Bench<T> {
 		// Make sure relevant storage is cleared. This is just to get the asserts to work when
 		// running tests because it seems the storage is not cleared in between.
 		inclusion::PendingAvailabilityCommitments::<T>::remove_all(None);
 		inclusion::PendingAvailability::<T>::remove_all(None);
 
 		// We don't allow a core to have both disputes and be marked fully available at this block.
 		let cores = self.max_cores();
-		let used_cores = dispute_sessions.len() as u32;
+		let used_cores =
+			(self.dispute_sessions.len() + self.backed_and_concluding_cores.len()) as u32;
 		assert!(used_cores <= cores);
 
-		// NOTE: there is an n+2 session delay for these actions to take effect
-		// We are currently in Session 0, so these changes will take effect in Session 2
+		// NOTE: there is an n+2 session delay for these actions to take effect.
+		// We are currently in Session 0, so these changes will take effect in Session 2.
 		Self::setup_para_ids(used_cores);
 
 		let validator_ids = Self::generate_validator_pairs(self.max_validators());
 		let target_session = SessionIndex::from(self.target_session);
 		let builder = self.setup_session(target_session, validator_ids, used_cores);
 
 		let bitfields =
-			builder.create_availability_bitfields(&backed_and_concluding_cores, used_cores);
-		let backed_candidates =
-			builder.create_backed_candidates(&backed_and_concluding_cores, includes_code_upgrade);
+			builder.create_availability_bitfields(&builder.backed_and_concluding_cores, used_cores);
+		let backed_candidates = builder
+			.create_backed_candidates(&builder.backed_and_concluding_cores, builder.code_upgrade);
 
 		let disputes = builder.create_disputes_with_no_spam(
-			backed_and_concluding_cores.len() as u32,
+			builder.backed_and_concluding_cores.len() as u32,
 			used_cores,
-			dispute_sessions,
+			builder.dispute_sessions.as_slice(),
 		);
 
 		assert_eq!(
@@ -585,7 +645,7 @@ impl<T: paras_inherent::Config> BenchBuilder<T> {
 		);
 		assert_eq!(inclusion::PendingAvailability::<T>::iter().count(), used_cores as usize,);
 
-		// Mark all the use cores as occupied. We expect that their are `backed_and_concluding_cores`
+		// Mark all the used cores as occupied. We expect that their are `backed_and_concluding_cores`
 		// that are pending availability and that there are `used_cores - backed_and_concluding_cores `
 		// which are about to be disputed.
 		scheduler::AvailabilityCores::<T>::set(vec![