ParaInherent create: update apply_weight_limit wrt elastic scaling

polkadot/runtime/parachains/src/paras_inherent/mod.rs

@@ -788,7 +788,7 @@ fn random_sel<X, F: Fn(&X) -> Weight>(
 /// Assumes disputes are already filtered by the time this is called.
 ///
 /// Returns the total weight consumed by `bitfields` and `candidates`.
-fn apply_weight_limit<T: Config + inclusion::Config>(
+pub(crate) fn apply_weight_limit<T: Config + inclusion::Config>(
 	candidates: &mut Vec<BackedCandidate<<T>::Hash>>,
 	bitfields: &mut UncheckedSignedAvailabilityBitfields,
 	max_consumable_weight: Weight,
@@ -805,35 +805,71 @@ fn apply_weight_limit<T: Config + inclusion::Config>(
 		return total
 	}
 
-	// Prefer code upgrades, they tend to be large and hence stand no chance to be picked
-	// late while maintaining the weight bounds.
-	let preferred_indices = candidates
+	// Invariant: block author provides candidate in the order in which they form a chain
+	// wrt elastic scaling. If the invariant is broken, we'd fail later when filtering candidates
+	// which are unchained.
+
+	let mut chained_candidates: Vec<Vec<_>> = Vec::new();
+	let mut current_para_id = None;
+
+	for candidate in sp_std::mem::take(candidates).into_iter() {
+		let candidate_para_id = candidate.descriptor().para_id;
+		if Some(candidate_para_id) == current_para_id {
+			let chain = chained_candidates
+				.last_mut()
+				.expect("if the current_para_id is Some, then vec is not empty; qed");
+			chain.push(candidate);
+		} else {
+			current_para_id = Some(candidate_para_id);
+			chained_candidates.push(vec![candidate]);
+		}
+	}
+
+	// Elastic scaling: we prefer chains that have a code upgrade among the candidates,
+	// as the candidates containing the upgrade tend to be large and hence stand no chance to
+	// be picked late while maintaining the weight bounds.
+	//
+	// Limitations: For simplicity if total weight of a chain of candidates is larger than
+	// the remaining weight, the chain will still not be included while it could still be possible
+	// to include part of that chain.
+	let preferred_chain_indices = chained_candidates
 		.iter()
 		.enumerate()
-		.filter_map(|(idx, candidate)| {
-			candidate.candidate().commitments.new_validation_code.as_ref().map(|_code| idx)
+		.filter_map(|(idx, candidates)| {
+			// Check if any of the candidate in chain contains a code upgrade.
+			if candidates
+				.iter()
+				.any(|candidate| candidate.candidate().commitments.new_validation_code.is_some())
+			{
+				Some(idx)
+			} else {
+				None
+			}
 		})
 		.collect::<Vec<usize>>();
 
-	// There is weight remaining to be consumed by a subset of candidates
+	// There is weight remaining to be consumed by a subset of chained candidates
 	// which are going to be picked now.
 	if let Some(max_consumable_by_candidates) =
 		max_consumable_weight.checked_sub(&total_bitfields_weight)
 	{
-		let (acc_candidate_weight, indices) =
-			random_sel::<BackedCandidate<<T as frame_system::Config>::Hash>, _>(
+		let (acc_candidate_weight, chained_indices) =
+			random_sel::<Vec<BackedCandidate<<T as frame_system::Config>::Hash>>, _>(
 				rng,
-				&candidates,
-				preferred_indices,
-				|c| backed_candidate_weight::<T>(c),
+				&chained_candidates,
+				preferred_chain_indices,
+				|candidates| backed_candidates_weight::<T>(&candidates),
 				max_consumable_by_candidates,
 			);
-		log::debug!(target: LOG_TARGET, "Indices Candidates: {:?}, size: {}", indices, candidates.len());
-		candidates.indexed_retain(|idx, _backed_candidate| indices.binary_search(&idx).is_ok());
+		log::debug!(target: LOG_TARGET, "Indices Candidates: {:?}, size: {}", chained_indices, candidates.len());
+		chained_candidates
+			.indexed_retain(|idx, _backed_candidates| chained_indices.binary_search(&idx).is_ok());
 		// pick all bitfields, and
 		// fill the remaining space with candidates
 		let total_consumed = acc_candidate_weight.saturating_add(total_bitfields_weight);
 
+		*candidates = chained_candidates.into_iter().flatten().collect::<Vec<_>>();
+
 		return total_consumed
 	}
 

polkadot/runtime/parachains/src/paras_inherent/tests.rs

@@ -22,7 +22,7 @@ use super::*;
 #[cfg(not(feature = "runtime-benchmarks"))]
 mod enter {
 
-	use super::*;
+	use super::{inclusion::tests::TestCandidateBuilder, *};
 	use crate::{
 		builder::{Bench, BenchBuilder},
 		mock::{mock_assigner, new_test_ext, BlockLength, BlockWeights, MockGenesisConfig, Test},
@@ -925,6 +925,129 @@ mod enter {
 		});
 	}
 
+	// Helper fn that builds chained dummy candidates for elastic scaling tests
+	fn build_backed_candidate_chain(
+		para_id: ParaId,
+		len: usize,
+		start_core_index: usize,
+		code_upgrade_index: Option<usize>,
+	) -> Vec<BackedCandidate> {
+		if let Some(code_upgrade_index) = code_upgrade_index {
+			assert!(code_upgrade_index < len, "Code upgrade index out of bounds");
+		}
+
+		(0..len)
+			.into_iter()
+			.map(|idx| {
+				let mut builder = TestCandidateBuilder::default();
+				builder.para_id = para_id;
+				let mut ccr = builder.build();
+
+				if Some(idx) == code_upgrade_index {
+					ccr.commitments.new_validation_code = Some(vec![1, 2, 3, 4].into());
+				}
+
+				ccr.commitments.processed_downward_messages = idx as u32;
+				let core_index = start_core_index + idx;
+
+				BackedCandidate::new(
+					ccr.into(),
+					Default::default(),
+					Default::default(),
+					Some(CoreIndex(core_index as u32)),
+				)
+			})
+			.collect::<Vec<_>>()
+	}
+
+	// Ensure that overweight parachain inherents are always rejected by the runtime.
+	// Runtime should panic and return `InherentOverweight` error.
+	#[test]
+	fn test_backed_candidates_apply_weight_works_for_elastic_scaling() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let seed = [
+				1, 0, 52, 0, 0, 0, 0, 0, 1, 0, 10, 0, 22, 32, 0, 0, 2, 0, 55, 49, 0, 11, 0, 0, 3,
+				0, 0, 0, 0, 0, 2, 92,
+			];
+			let mut rng = rand_chacha::ChaChaRng::from_seed(seed);
+
+			// Create an overweight inherent and oversized block
+			let mut backed_and_concluding = BTreeMap::new();
+
+			for i in 0..30 {
+				backed_and_concluding.insert(i, i);
+			}
+
+			let scenario = make_inherent_data(TestConfig {
+				dispute_statements: Default::default(),
+				dispute_sessions: vec![], // 3 cores with disputes
+				backed_and_concluding,
+				num_validators_per_core: 5,
+				code_upgrade: None,
+				fill_claimqueue: false,
+			});
+
+			let mut para_inherent_data = scenario.data.clone();
+
+			// Check the para inherent data is as expected:
+			// * 1 bitfield per validator (5 validators per core, 30 backed candidates, 0 disputes
+			//   => 5*30 = 150)
+			assert_eq!(para_inherent_data.bitfields.len(), 150);
+			// * 30 backed candidates
+			assert_eq!(para_inherent_data.backed_candidates.len(), 30);
+
+			let mut input_candidates =
+				build_backed_candidate_chain(ParaId::from(1000), 3, 0, Some(1));
+			let chained_candidates_weight = backed_candidates_weight::<Test>(&input_candidates);
+
+			input_candidates.append(&mut para_inherent_data.backed_candidates);
+			let input_bitfields = para_inherent_data.bitfields;
+
+			// Test if weight insufficient even for 1 candidate (which doesn't contain a code
+			// upgrade).
+			let max_weight = backed_candidate_weight::<Test>(&input_candidates[0]) +
+				signed_bitfields_weight::<Test>(&input_bitfields);
+			let mut backed_candidates = input_candidates.clone();
+			let mut bitfields = input_bitfields.clone();
+			apply_weight_limit::<Test>(
+				&mut backed_candidates,
+				&mut bitfields,
+				max_weight,
+				&mut rng,
+			);
+
+			// The chained candidates are not picked, instead a single other candidate is picked
+			assert_eq!(backed_candidates.len(), 1);
+			assert_ne!(backed_candidates[0].descriptor().para_id, ParaId::from(1000));
+
+			// All bitfields are kept.
+			assert_eq!(bitfields.len(), 150);
+
+			// Test if para_id 1000 chained candidates make it if there is enough room for its 3
+			// candidates.
+			let max_weight =
+				chained_candidates_weight + signed_bitfields_weight::<Test>(&input_bitfields);
+			let mut backed_candidates = input_candidates.clone();
+			let mut bitfields = input_bitfields.clone();
+			apply_weight_limit::<Test>(
+				&mut backed_candidates,
+				&mut bitfields,
+				max_weight,
+				&mut rng,
+			);
+
+			// Only the chained candidates should pass filter.
+			assert_eq!(backed_candidates.len(), 3);
+			// Check the actual candidates
+			assert_eq!(backed_candidates[0].descriptor().para_id, ParaId::from(1000));
+			assert_eq!(backed_candidates[1].descriptor().para_id, ParaId::from(1000));
+			assert_eq!(backed_candidates[2].descriptor().para_id, ParaId::from(1000));
+
+			// All bitfields are kept.
+			assert_eq!(bitfields.len(), 150);
+		});
+	}
+
 	// Ensure that overweight parachain inherents are always rejected by the runtime.
 	// Runtime should panic and return `InherentOverweight` error.
 	#[test]