prospective-parachains: respond with multiple backable candidates

polkadot/node/core/prospective-parachains/src/fragment_tree.rs

@@ -763,43 +763,106 @@ impl FragmentTree {
 	/// The intention of the `required_path` is to allow queries on the basis of
 	/// one or more candidates which were previously pending availability becoming
 	/// available and opening up more room on the core.
-	pub(crate) fn select_child(
+	pub(crate) fn select_children(
 		&self,
 		required_path: &[CandidateHash],
+		count: u32,
 		pred: impl Fn(&CandidateHash) -> bool,
-	) -> Option<CandidateHash> {
+	) -> Vec<CandidateHash> {
 		let base_node = {
 			// traverse the required path.
 			let mut node = NodePointer::Root;
 			for required_step in required_path {
-				node = self.node_candidate_child(node, &required_step)?;
+				if let Some(next_node) = self.node_candidate_child(node, &required_step) {
+					node = next_node;
+				} else {
+					return vec![]
+				};
 			}
 
 			node
 		};
 
-		// TODO [now]: taking the first selection might introduce bias
+		// TODO: taking the first best selection might introduce bias
 		// or become gameable.
 		//
 		// For plausibly unique parachains, this shouldn't matter much.
 		// figure out alternative selection criteria?
-		match base_node {
+		self.select_children_inner(base_node, count, count, &pred, &mut vec![])
+	}
+
+	// Try finding a candidate chain starting from `base_node` of length `expected_count`.
+	// If not possible, return the longest we could find.
+	// Does a depth-first search, since we're optimistic that there won't be more than one such
+	// chains. Assumes that the chain must be acyclic.
+	fn select_children_inner(
+		&self,
+		base_node: NodePointer,
+		expected_count: u32,
+		remaining_count: u32,
+		pred: &dyn Fn(&CandidateHash) -> bool,
+		accumulator: &mut Vec<CandidateHash>,
+	) -> Vec<CandidateHash> {
+		if remaining_count == 0 {
+			// The best option is the chain we've accumulated so far.
+			return accumulator.to_vec();
+		}
+
+		let children: Vec<_> = match base_node {
 			NodePointer::Root => self
 				.nodes
 				.iter()
-				.take_while(|n| n.parent == NodePointer::Root)
-				.filter(|n| self.scope.get_pending_availability(&n.candidate_hash).is_none())
-				.filter(|n| pred(&n.candidate_hash))
-				.map(|n| n.candidate_hash)
-				.next(),
-			NodePointer::Storage(ptr) => self.nodes[ptr]
-				.children
-				.iter()
-				.filter(|n| self.scope.get_pending_availability(&n.1).is_none())
-				.filter(|n| pred(&n.1))
-				.map(|n| n.1)
-				.next(),
+				.enumerate()
+				.take_while(|(_, n)| n.parent == NodePointer::Root)
+				.filter(|(_, n)| self.scope.get_pending_availability(&n.candidate_hash).is_none())
+				.filter(|(_, n)| pred(&n.candidate_hash))
+				.map(|(ptr, n)| (NodePointer::Storage(ptr), n.candidate_hash))
+				.collect(),
+			NodePointer::Storage(base_node_ptr) => {
+				let base_node = &self.nodes[base_node_ptr];
+
+				base_node
+					.children
+					.iter()
+					.filter(|(_, hash)| self.scope.get_pending_availability(&hash).is_none())
+					.filter(|(_, hash)| pred(&hash))
+					.map(|(ptr, hash)| (*ptr, *hash))
+					.collect()
+			},
+		};
+
+		let mut best_result = accumulator.clone();
+		for (child_ptr, child_hash) in children {
+			// We could use a HashSet/BTreeSet for tracking the visited nodes but realistically,
+			// accumulator will not hold more than 2-3 elements (the max number of cores for a
+			// parachain).
+			if accumulator.contains(&child_hash) {
+				// We've already visited this node. There's a cycle in the tree, ignore it.
+				continue
+			}
+
+			let mut accumulator_copy = accumulator.clone();
+
+			accumulator_copy.push(child_hash);
+
+			let result = self.select_children_inner(
+				child_ptr,
+				expected_count,
+				remaining_count - 1,
+				&pred,
+				&mut accumulator_copy,
+			);
+
+			// Short-circuit the search if we've found the right length. Otherwise, we'll
+			// search for a max.
+			if result.len() == expected_count as usize {
+				return result
+			} else if best_result.len() < result.len() {
+				best_result = result;
+			}
 		}
+
+		best_result
 	}
 
 	fn populate_from_bases(&mut self, storage: &CandidateStorage, initial_bases: Vec<NodePointer>) {

polkadot/node/core/prospective-parachains/src/lib.rs

@@ -146,12 +146,20 @@ async fn run_iteration<Context>(
 					handle_candidate_seconded(view, para, candidate_hash),
 				ProspectiveParachainsMessage::CandidateBacked(para, candidate_hash) =>
 					handle_candidate_backed(&mut *ctx, view, para, candidate_hash).await?,
-				ProspectiveParachainsMessage::GetBackableCandidate(
+				ProspectiveParachainsMessage::GetBackableCandidates(
 					relay_parent,
 					para,
+					n_cores,
 					required_path,
 					tx,
-				) => answer_get_backable_candidate(&view, relay_parent, para, required_path, tx),
+				) => answer_get_backable_candidate(
+					&view,
+					relay_parent,
+					para,
+					n_cores,
+					required_path,
+					tx,
+				),
 				ProspectiveParachainsMessage::GetHypotheticalFrontier(request, tx) =>
 					answer_hypothetical_frontier_request(&view, request, tx),
 				ProspectiveParachainsMessage::GetTreeMembership(para, candidate, tx) =>
@@ -556,8 +564,9 @@ fn answer_get_backable_candidate(
 	view: &View,
 	relay_parent: Hash,
 	para: ParaId,
+	n_cores: u32,
 	required_path: Vec<CandidateHash>,
-	tx: oneshot::Sender<Option<(CandidateHash, Hash)>>,
+	tx: oneshot::Sender<Vec<(CandidateHash, Hash)>>,
 ) {
 	let data = match view.active_leaves.get(&relay_parent) {
 		None => {
@@ -568,7 +577,7 @@ fn answer_get_backable_candidate(
 				"Requested backable candidate for inactive relay-parent."
 			);
 
-			let _ = tx.send(None);
+			let _ = tx.send(vec![]);
 			return
 		},
 		Some(d) => d,
@@ -583,7 +592,7 @@ fn answer_get_backable_candidate(
 				"Requested backable candidate for inactive para."
 			);
 
-			let _ = tx.send(None);
+			let _ = tx.send(vec![]);
 			return
 		},
 		Some(tree) => tree,
@@ -598,30 +607,32 @@ fn answer_get_backable_candidate(
 				"No candidate storage for active para",
 			);
 
-			let _ = tx.send(None);
+			let _ = tx.send(vec![]);
 			return
 		},
 		Some(s) => s,
 	};
 
-	let Some(child_hash) =
-		tree.select_child(&required_path, |candidate| storage.is_backed(candidate))
-	else {
-		let _ = tx.send(None);
-		return
-	};
-	let Some(candidate_relay_parent) = storage.relay_parent_by_candidate_hash(&child_hash) else {
-		gum::error!(
-			target: LOG_TARGET,
-			?child_hash,
-			para_id = ?para,
-			"Candidate is present in fragment tree but not in candidate's storage!",
-		);
-		let _ = tx.send(None);
-		return
-	};
+	let backable_children = tree
+		.select_children(&required_path, n_cores, |candidate| storage.is_backed(candidate))
+		.into_iter()
+		.filter_map(|child_hash| {
+			storage.relay_parent_by_candidate_hash(&child_hash).map_or_else(
+				|| {
+					gum::error!(
+						target: LOG_TARGET,
+						?child_hash,
+						para_id = ?para,
+						"Candidate is present in fragment tree but not in candidate's storage!",
+					);
+					None
+				},
+				|parent_hash| Some((child_hash, parent_hash)),
+			)
+		})
+		.collect();
 
-	let _ = tx.send(Some((child_hash, candidate_relay_parent)));
+	let _ = tx.send(backable_children);
 }
 
 fn answer_hypothetical_frontier_request(