elastic scaling: preserve candidate ordering in provisioner (#3778)

https://github.com/paritytech/polkadot-sdk/issues/3742

polkadot/node/core/provisioner/src/lib.rs

@@ -46,7 +46,7 @@ use polkadot_primitives::{
 	BackedCandidate, BlockNumber, CandidateHash, CandidateReceipt, CoreIndex, CoreState, Hash,
 	Id as ParaId, OccupiedCoreAssumption, SessionIndex, SignedAvailabilityBitfield, ValidatorIndex,
 };
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap};
 
 mod disputes;
 mod error;
@@ -598,13 +598,11 @@ async fn select_candidate_hashes_from_tracked(
 	candidates: &[CandidateReceipt],
 	relay_parent: Hash,
 	sender: &mut impl overseer::ProvisionerSenderTrait,
-) -> Result<Vec<(CandidateHash, Hash)>, Error> {
+) -> Result<HashMap<ParaId, Vec<(CandidateHash, Hash)>>, Error> {
 	let block_number = get_block_number_under_construction(relay_parent, sender).await?;
 
 	let mut selected_candidates =
-		Vec::with_capacity(candidates.len().min(availability_cores.len()));
-	let mut selected_parachains =
-		HashSet::with_capacity(candidates.len().min(availability_cores.len()));
+		HashMap::with_capacity(candidates.len().min(availability_cores.len()));
 
 	gum::debug!(
 		target: LOG_TARGET,
@@ -638,7 +636,7 @@ async fn select_candidate_hashes_from_tracked(
 			CoreState::Free => continue,
 		};
 
-		if selected_parachains.contains(&scheduled_core.para_id) {
+		if selected_candidates.contains_key(&scheduled_core.para_id) {
 			// We already picked a candidate for this parachain. Elastic scaling only works with
 			// prospective parachains mode.
 			continue
@@ -677,8 +675,10 @@ async fn select_candidate_hashes_from_tracked(
 				"Selected candidate receipt",
 			);
 
-			selected_parachains.insert(candidate.descriptor.para_id);
-			selected_candidates.push((candidate_hash, candidate.descriptor.relay_parent));
+			selected_candidates.insert(
+				candidate.descriptor.para_id,
+				vec![(candidate_hash, candidate.descriptor.relay_parent)],
+			);
 		}
 	}
 
@@ -695,12 +695,12 @@ async fn request_backable_candidates(
 	bitfields: &[SignedAvailabilityBitfield],
 	relay_parent: Hash,
 	sender: &mut impl overseer::ProvisionerSenderTrait,
-) -> Result<Vec<(CandidateHash, Hash)>, Error> {
+) -> Result<HashMap<ParaId, Vec<(CandidateHash, Hash)>>, Error> {
 	let block_number = get_block_number_under_construction(relay_parent, sender).await?;
 
 	// Record how many cores are scheduled for each paraid. Use a BTreeMap because
 	// we'll need to iterate through them.
-	let mut scheduled_cores: BTreeMap<ParaId, usize> = BTreeMap::new();
+	let mut scheduled_cores_per_para: BTreeMap<ParaId, usize> = BTreeMap::new();
 	// The on-chain ancestors of a para present in availability-cores.
 	let mut ancestors: HashMap<ParaId, Ancestors> =
 		HashMap::with_capacity(availability_cores.len());
@@ -709,7 +709,7 @@ async fn request_backable_candidates(
 		let core_idx = CoreIndex(core_idx as u32);
 		match core {
 			CoreState::Scheduled(scheduled_core) => {
-				*scheduled_cores.entry(scheduled_core.para_id).or_insert(0) += 1;
+				*scheduled_cores_per_para.entry(scheduled_core.para_id).or_insert(0) += 1;
 			},
 			CoreState::Occupied(occupied_core) => {
 				let is_available = bitfields_indicate_availability(
@@ -726,14 +726,14 @@ async fn request_backable_candidates(
 
 					if let Some(ref scheduled_core) = occupied_core.next_up_on_available {
 						// Request a new backable candidate for the newly scheduled para id.
-						*scheduled_cores.entry(scheduled_core.para_id).or_insert(0) += 1;
+						*scheduled_cores_per_para.entry(scheduled_core.para_id).or_insert(0) += 1;
 					}
 				} else if occupied_core.time_out_at <= block_number {
 					// Timed out before being available.
 
 					if let Some(ref scheduled_core) = occupied_core.next_up_on_time_out {
 						// Candidate's availability timed out, practically same as scheduled.
-						*scheduled_cores.entry(scheduled_core.para_id).or_insert(0) += 1;
+						*scheduled_cores_per_para.entry(scheduled_core.para_id).or_insert(0) += 1;
 					}
 				} else {
 					// Not timed out and not available.
@@ -747,10 +747,10 @@ async fn request_backable_candidates(
 		};
 	}
 
-	let mut selected_candidates: Vec<(CandidateHash, Hash)> =
-		Vec::with_capacity(availability_cores.len());
+	let mut selected_candidates: HashMap<ParaId, Vec<(CandidateHash, Hash)>> =
+		HashMap::with_capacity(scheduled_cores_per_para.len());
 
-	for (para_id, core_count) in scheduled_cores {
+	for (para_id, core_count) in scheduled_cores_per_para {
 		let para_ancestors = ancestors.remove(&para_id).unwrap_or_default();
 
 		// If elastic scaling MVP is disabled, only allow one candidate per parachain.
@@ -777,7 +777,7 @@ async fn request_backable_candidates(
 			continue
 		}
 
-		selected_candidates.extend(response.into_iter().take(core_count));
+		selected_candidates.insert(para_id, response);
 	}
 
 	Ok(selected_candidates)
@@ -826,33 +826,38 @@ async fn select_candidates(
 		selected_candidates.clone(),
 		tx,
 	));
-	let mut candidates = rx.await.map_err(|err| Error::CanceledBackedCandidates(err))?;
+	let candidates = rx.await.map_err(|err| Error::CanceledBackedCandidates(err))?;
 	gum::trace!(target: LOG_TARGET, leaf_hash=?relay_parent,
 				"Got {} backed candidates", candidates.len());
 
 	// keep only one candidate with validation code.
 	let mut with_validation_code = false;
-	candidates.retain(|c| {
-		if c.candidate().commitments.new_validation_code.is_some() {
-			if with_validation_code {
-				return false
+	// merge the candidates into a common collection, preserving the order
+	let mut merged_candidates = Vec::with_capacity(availability_cores.len());
+
+	for para_candidates in candidates.into_values() {
+		for candidate in para_candidates {
+			if candidate.candidate().commitments.new_validation_code.is_some() {
+				if with_validation_code {
+					break
+				} else {
+					with_validation_code = true;
+				}
 			}
 
-			with_validation_code = true;
+			merged_candidates.push(candidate);
 		}
-
-		true
-	});
+	}
 
 	gum::debug!(
 		target: LOG_TARGET,
-		n_candidates = candidates.len(),
+		n_candidates = merged_candidates.len(),
 		n_cores = availability_cores.len(),
 		?relay_parent,
 		"Selected backed candidates",
 	);
 
-	Ok(candidates)
+	Ok(merged_candidates)
 }
 
 /// Produces a block number 1 higher than that of the relay parent

polkadot/node/core/provisioner/src/tests.rs

@@ -258,6 +258,8 @@ mod select_candidates {
 		BlockNumber, CandidateCommitments, CommittedCandidateReceipt, PersistedValidationData,
 	};
 	use rstest::rstest;
+	use std::ops::Not;
+	use CoreState::{Free, Scheduled};
 
 	const BLOCK_UNDER_PRODUCTION: BlockNumber = 128;
 
@@ -323,9 +325,6 @@ mod select_candidates {
 	//     11: Occupied(next_up_on_available and available, but different successor para_id)
 	//   ]
 	fn mock_availability_cores_one_per_para() -> Vec<CoreState> {
-		use std::ops::Not;
-		use CoreState::{Free, Scheduled};
-
 		vec![
 			// 0: Free,
 			Free,
@@ -389,9 +388,6 @@ mod select_candidates {
 	// For test purposes with multiple possible cores assigned to a para, we always return this set
 	// of availability cores:
 	fn mock_availability_cores_multiple_per_para() -> Vec<CoreState> {
-		use std::ops::Not;
-		use CoreState::{Free, Scheduled};
-
 		vec![
 			// 0: Free,
 			Free,
@@ -562,7 +558,10 @@ mod select_candidates {
 		use ChainApiMessage::BlockNumber;
 		use RuntimeApiMessage::Request;
 
-		let mut backed_iter = expected.clone().into_iter();
+		let mut backed = expected.clone().into_iter().fold(HashMap::new(), |mut acc, candidate| {
+			acc.entry(candidate.descriptor().para_id).or_insert(vec![]).push(candidate);
+			acc
+		});
 
 		expected.sort_by_key(|c| c.candidate().descriptor.para_id);
 		let mut candidates_iter = expected
@@ -583,11 +582,30 @@ mod select_candidates {
 					hashes,
 					sender,
 				)) => {
-					let response: Vec<BackedCandidate> =
-						backed_iter.by_ref().take(hashes.len()).collect();
-					let expected_hashes: Vec<(CandidateHash, Hash)> = response
+					let mut response: HashMap<ParaId, Vec<BackedCandidate>> = HashMap::new();
+					for (para_id, requested_candidates) in hashes.clone() {
+						response.insert(
+							para_id,
+							backed
+								.get_mut(&para_id)
+								.unwrap()
+								.drain(0..requested_candidates.len())
+								.collect(),
+						);
+					}
+					let expected_hashes: HashMap<ParaId, Vec<(CandidateHash, Hash)>> = response
 						.iter()
-						.map(|candidate| (candidate.hash(), candidate.descriptor().relay_parent))
+						.map(|(para_id, candidates)| {
+							(
+								*para_id,
+								candidates
+									.iter()
+									.map(|candidate| {
+										(candidate.hash(), candidate.descriptor().relay_parent)
+									})
+									.collect(),
+							)
+						})
 						.collect();
 
 					assert_eq!(expected_hashes, hashes);
@@ -768,7 +786,7 @@ mod select_candidates {
 	#[rstest]
 	#[case(ProspectiveParachainsMode::Disabled)]
 	#[case(ProspectiveParachainsMode::Enabled {max_candidate_depth: 0, allowed_ancestry_len: 0})]
-	fn selects_max_one_code_upgrade(
+	fn selects_max_one_code_upgrade_one_core_per_para(
 		#[case] prospective_parachains_mode: ProspectiveParachainsMode,
 	) {
 		let mock_cores = mock_availability_cores_one_per_para();
@@ -780,7 +798,12 @@ mod select_candidates {
 		let cores = [1, 4, 7, 8, 10, 12];
 		let cores_with_code = [1, 4, 8];
 
-		let expected_cores = [1, 7, 10, 12];
+		// We can't be sure which one code upgrade the provisioner will pick. We can only assert
+		// that it only picks one. These are the possible cores for which the provisioner will
+		// supply candidates.
+		// There are multiple possibilities depending on which code upgrade it
+		// chooses.
+		let possible_expected_cores = [[1, 7, 10, 12], [4, 7, 10, 12], [7, 8, 10, 12]];
 
 		let committed_receipts: Vec<_> = (0..=mock_cores.len())
 			.map(|i| {
@@ -820,8 +843,10 @@ mod select_candidates {
 		// Then, some of them get filtered due to new validation code rule.
 		let expected_backed: Vec<_> =
 			cores.iter().map(|&idx| backed_candidates[idx].clone()).collect();
-		let expected_backed_filtered: Vec<_> =
-			expected_cores.iter().map(|&idx| candidates[idx].clone()).collect();
+		let expected_backed_filtered: Vec<Vec<_>> = possible_expected_cores
+			.iter()
+			.map(|indices| indices.iter().map(|&idx| candidates[idx].clone()).collect())
+			.collect();
 
 		let mock_cores_clone = mock_cores.clone();
 
@@ -850,13 +875,120 @@ mod select_candidates {
 
 				assert_eq!(result.len(), 4);
 
-				result.into_iter().for_each(|c| {
-					assert!(
-						expected_backed_filtered.iter().any(|c2| c.candidate().corresponds_to(c2)),
-						"Failed to find candidate: {:?}",
-						c,
-					)
-				});
+				assert!(expected_backed_filtered.iter().any(|expected_backed_filtered| {
+					result.clone().into_iter().all(|c| {
+						expected_backed_filtered.iter().any(|c2| c.candidate().corresponds_to(c2))
+					})
+				}));
+			},
+		)
+	}
+
+	#[test]
+	fn selects_max_one_code_upgrade_multiple_cores_per_para() {
+		let prospective_parachains_mode =
+			ProspectiveParachainsMode::Enabled { max_candidate_depth: 0, allowed_ancestry_len: 0 };
+		let mock_cores = vec![
+			// 0: Scheduled(default),
+			Scheduled(scheduled_core(1)),
+			// 1: Scheduled(default),
+			Scheduled(scheduled_core(2)),
+			// 2: Scheduled(default),
+			Scheduled(scheduled_core(2)),
+			// 3: Scheduled(default),
+			Scheduled(scheduled_core(2)),
+			// 4: Scheduled(default),
+			Scheduled(scheduled_core(3)),
+			// 5: Scheduled(default),
+			Scheduled(scheduled_core(3)),
+			// 6: Scheduled(default),
+			Scheduled(scheduled_core(3)),
+		];
+
+		let empty_hash = PersistedValidationData::<Hash, BlockNumber>::default().hash();
+		let cores_with_code = [0, 2, 4, 5];
+
+		// We can't be sure which one code upgrade the provisioner will pick. We can only assert
+		// that it only picks one.
+		// These are the possible cores for which the provisioner will
+		// supply candidates. There are multiple possibilities depending on which code upgrade it
+		// chooses.
+		let possible_expected_cores = [vec![0, 1], vec![1, 2, 3], vec![4, 1]];
+
+		let committed_receipts: Vec<_> = (0..mock_cores.len())
+			.map(|i| {
+				let mut descriptor = dummy_candidate_descriptor(dummy_hash());
+				descriptor.para_id = mock_cores[i].para_id().unwrap();
+				descriptor.persisted_validation_data_hash = empty_hash;
+				descriptor.pov_hash = Hash::from_low_u64_be(i as u64);
+				CommittedCandidateReceipt {
+					descriptor,
+					commitments: CandidateCommitments {
+						new_validation_code: if cores_with_code.contains(&i) {
+							Some(vec![].into())
+						} else {
+							None
+						},
+						..Default::default()
+					},
+				}
+			})
+			.collect();
+
+		// Input to select_candidates
+		let candidates: Vec<_> = committed_receipts.iter().map(|r| r.to_plain()).collect();
+		// Build possible outputs from select_candidates
+		let backed_candidates: Vec<_> = committed_receipts
+			.iter()
+			.map(|committed_receipt| {
+				BackedCandidate::new(
+					committed_receipt.clone(),
+					Vec::new(),
+					default_bitvec(MOCK_GROUP_SIZE),
+					None,
+				)
+			})
+			.collect();
+
+		// First, provisioner will request backable candidates for each scheduled core.
+		// Then, some of them get filtered due to new validation code rule.
+		let expected_backed: Vec<_> =
+			(0..mock_cores.len()).map(|idx| backed_candidates[idx].clone()).collect();
+		let expected_backed_filtered: Vec<Vec<_>> = possible_expected_cores
+			.iter()
+			.map(|indices| indices.iter().map(|&idx| candidates[idx].clone()).collect())
+			.collect();
+
+		let mock_cores_clone = mock_cores.clone();
+
+		test_harness(
+			|r| {
+				mock_overseer(
+					r,
+					mock_cores_clone,
+					expected_backed,
+					HashMap::new(),
+					prospective_parachains_mode,
+				)
+			},
+			|mut tx: TestSubsystemSender| async move {
+				let result = select_candidates(
+					&mock_cores,
+					&[],
+					&candidates,
+					prospective_parachains_mode,
+					true,
+					Default::default(),
+					&mut tx,
+				)
+				.await
+				.unwrap();
+
+				assert!(expected_backed_filtered.iter().any(|expected_backed_filtered| {
+					result.clone().into_iter().all(|c| {
+						expected_backed_filtered.iter().any(|c2| c.candidate().corresponds_to(c2))
+					}) && (expected_backed_filtered.len() == result.len())
+				}));
 			},
 		)
 	}