Adjust parent search to allow for pending blocks with depth > 1

cumulus/client/consensus/common/src/lib.rs

@@ -250,7 +250,7 @@ pub struct ParentSearchParams {
 }
 
 /// A potential parent block returned from [`find_potential_parents`]
-#[derive(Debug, PartialEq)]
+#[derive(PartialEq)]
 pub struct PotentialParent<B: BlockT> {
 	/// The hash of the block.
 	pub hash: B::Hash,
@@ -263,6 +263,17 @@ pub struct PotentialParent<B: BlockT> {
 	pub aligned_with_pending: bool,
 }
 
+impl<B: BlockT> std::fmt::Debug for PotentialParent<B> {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("PotentialParent")
+			.field("hash", &self.hash)
+			.field("depth", &self.depth)
+			.field("aligned_with_pending", &self.aligned_with_pending)
+			.field("number", &self.header.number())
+			.finish()
+	}
+}
+
 /// Perform a recursive search through blocks to find potential
 /// parent blocks for a new block.
 ///
@@ -280,10 +291,11 @@ pub struct PotentialParent<B: BlockT> {
 ///   * the block number is within `max_depth` blocks of the included block
 pub async fn find_potential_parents<B: BlockT>(
 	params: ParentSearchParams,
-	client: &impl Backend<B>,
+	backend: &impl Backend<B>,
 	relay_client: &impl RelayChainInterface,
 ) -> Result<Vec<PotentialParent<B>>, RelayChainError> {
 	// 1. Build up the ancestry record of the relay chain to compare against.
+	tracing::trace!("Parent search parameters: {params:?}");
 	let rp_ancestry = {
 		let mut ancestry = Vec::with_capacity(params.ancestry_lookback + 1);
 		let mut current_rp = params.relay_parent;
@@ -352,24 +364,96 @@ pub async fn find_potential_parents<B: BlockT>(
 	let included_hash = included_header.hash();
 	let pending_hash = pending_header.as_ref().map(|hdr| hdr.hash());
 
+	if params.max_depth == 0 {
+		return Ok(vec![PotentialParent::<B> {
+			hash: included_hash,
+			header: included_header,
+			depth: 0,
+			aligned_with_pending: true,
+		}])
+	};
+
+	let maybe_route = pending_hash
+		.map(|pending| sp_blockchain::tree_route(backend.blockchain(), included_hash, pending))
+		.transpose()?;
+
+	// The distance between pending and included block. Is later used to check if a child
+	// is aligned with pending when it is between pending and included block.
+	let pending_distance = maybe_route.as_ref().map(|route| route.enacted().len());
+
+	// If we want to ignore alternative branches there is no reason to start
+	// the parent search at the included block. We can add the included block and
+	// the path to the pending block to the potential parents directly (limited by max_depth).
+	let (mut frontier, mut potential_parents) = if let (Some(pending), true, Some(ref route)) =
+		(pending_header, params.ignore_alternative_branches, &maybe_route)
+	{
+		let mut potential_parents = Vec::new();
+		// Included block is always a potential parent
+		potential_parents.push(PotentialParent::<B> {
+			hash: included_hash,
+			header: included_header.clone(),
+			depth: 0,
+			aligned_with_pending: true,
+		});
+
+		// Add all items on the path included -> pending - 1 to the potential parents, but not
+		// more than `max_depth`.
+		let num_parents_on_path = route.enacted().len().saturating_sub(1).min(params.max_depth);
+		for (num, block) in route.enacted().iter().take(num_parents_on_path).enumerate() {
+			let header = match backend.blockchain().header(block.hash) {
+				Ok(Some(h)) => h,
+				Ok(None) => continue,
+				Err(_) => continue,
+			};
+
+			potential_parents.push(PotentialParent::<B> {
+				hash: block.hash,
+				header,
+				depth: 1 + num,
+				aligned_with_pending: true,
+			});
+		}
+
+		// The search for additional potential parents should now start at the
+		// pending block.
+		(
+			vec![PotentialParent::<B> {
+				hash: pending.hash(),
+				header: pending.clone(),
+				depth: route.enacted().len(),
+				aligned_with_pending: true,
+			}],
+			potential_parents,
+		)
+	} else {
+		(
+			vec![PotentialParent::<B> {
+				hash: included_hash,
+				header: included_header.clone(),
+				depth: 0,
+				aligned_with_pending: true,
+			}],
+			Default::default(),
+		)
+	};
+
+	if potential_parents.len() > params.max_depth {
+		return Ok(potential_parents);
+	}
+
+	// If a block is on the path included -> pending, we consider it `aligned_with_pending`.
+	let is_child_in_path_to_pending = |hash| {
+		maybe_route
+			.as_ref()
+			.map_or(true, |route| route.enacted().iter().any(|x| x.hash == hash))
+	};
+
 	tracing::trace!(target: PARENT_SEARCH_LOG_TARGET, ?included_hash, included_num = ?included_header.number(), ?pending_hash ,?rp_ancestry, "Searching relay chain ancestry.");
-	let mut frontier = vec![PotentialParent::<B> {
-		hash: included_hash,
-		header: included_header,
-		depth: 0,
-		aligned_with_pending: true,
-	}];
-
-	// Recursive search through descendants of the included block which have acceptable
-	// relay parents.
-	let mut potential_parents = Vec::new();
 	while let Some(entry) = frontier.pop() {
-		// TODO find_potential_parents The assumption that entry.depth = 1 is the pending block is
-		// not correct if we produce sub 6s blocks.
+		// TODO Adjust once we can fetch multiple pending blocks.
 		// https://github.com/paritytech/polkadot-sdk/issues/3967
-		let is_pending =
-			entry.depth == 1 && pending_hash.as_ref().map_or(false, |h| &entry.hash == h);
-		let is_included = entry.depth == 0;
+		let is_pending = pending_hash.as_ref().map_or(false, |h| &entry.hash == h);
+		let is_included = included_hash == entry.hash;
 
 		// note: even if the pending block or included block have a relay parent
 		// outside of the expected part of the relay chain, they are always allowed
@@ -400,19 +484,19 @@ pub async fn find_potential_parents<B: BlockT>(
 		}
 
 		// push children onto search frontier.
-		for child in client.blockchain().children(hash).ok().into_iter().flatten() {
+		for child in backend.blockchain().children(hash).ok().into_iter().flatten() {
+			tracing::trace!(target: PARENT_SEARCH_LOG_TARGET, ?child, child_depth, ?pending_distance, "Looking at child.");
 			let aligned_with_pending = parent_aligned_with_pending &&
-				if child_depth == 1 {
-					pending_hash.as_ref().map_or(true, |h| &child == h)
-				} else {
-					true
-				};
+				(pending_distance.map_or(true, |dist| child_depth > dist) ||
+					pending_hash.as_ref().map_or(true, |h| &child == h) ||
+					is_child_in_path_to_pending(child));
 
 			if params.ignore_alternative_branches && !aligned_with_pending {
+				tracing::trace!(target: PARENT_SEARCH_LOG_TARGET, ?child, "Child is not aligned with pending block.");
 				continue
 			}
 
-			let header = match client.blockchain().header(child) {
+			let header = match backend.blockchain().header(child) {
 				Ok(Some(h)) => h,
 				Ok(None) => continue,
 				Err(_) => continue,

cumulus/client/consensus/common/src/parachain_consensus.rs

@@ -385,7 +385,7 @@ async fn handle_new_best_parachain_head<Block, P>(
 			unset_best_header.take();
 			tracing::debug!(
 				target: LOG_TARGET,
-				?hash,
+				included = ?hash,
 				"Importing block as new best for parachain.",
 			);
 			import_block_as_new_best(hash, parachain_head, parachain).await;

cumulus/client/consensus/common/src/tests.rs

@@ -1133,6 +1133,177 @@ fn find_potential_parents_with_max_depth() {
 	}
 }
 
+/// Test where there is an additional block between included and pending block.
+#[test]
+fn find_potential_parents_aligned_with_late_pending() {
+	sp_tracing::try_init_simple();
+
+	const NON_INCLUDED_CHAIN_LEN: usize = 5;
+
+	let backend = Arc::new(Backend::new_test(1000, 1));
+	let client = Arc::new(TestClientBuilder::with_backend(backend.clone()).build());
+	let mut para_import =
+		ParachainBlockImport::new_with_delayed_best_block(client.clone(), backend.clone());
+
+	let relay_parent = relay_hash_from_block_num(10);
+	// Choose different relay parent for alternative chain to get new hashes.
+	let search_relay_parent = relay_hash_from_block_num(11);
+	let included_block = build_and_import_block_ext(
+		&client,
+		BlockOrigin::NetworkInitialSync,
+		true,
+		&mut para_import,
+		None,
+		None,
+		Some(relay_parent),
+	);
+
+	let in_between_block = build_and_import_block_ext(
+		&client,
+		BlockOrigin::NetworkInitialSync,
+		true,
+		&mut para_import,
+		Some(included_block.header().hash()),
+		None,
+		Some(relay_parent),
+	);
+
+	let pending_block = build_and_import_block_ext(
+		&client,
+		BlockOrigin::Own,
+		true,
+		&mut para_import,
+		Some(in_between_block.header().hash()),
+		None,
+		Some(relay_parent),
+	);
+
+	let relay_chain = Relaychain::new();
+	{
+		let relay_inner = &mut relay_chain.inner.lock().unwrap();
+		relay_inner
+			.relay_chain_hash_to_header
+			.insert(search_relay_parent, included_block.header().clone());
+		relay_inner
+			.relay_chain_hash_to_header_pending
+			.insert(search_relay_parent, in_between_block.header().clone());
+		relay_inner
+			.relay_chain_hash_to_header_pending
+			.insert(search_relay_parent, pending_block.header().clone());
+	}
+
+	// Build two sibling chains from the included block.
+	let mut aligned_blocks = Vec::new();
+	let mut parent = pending_block.header().hash();
+	for _ in 2..NON_INCLUDED_CHAIN_LEN {
+		let block = build_and_import_block_ext(
+			&client,
+			BlockOrigin::Own,
+			true,
+			&mut para_import,
+			Some(parent),
+			None,
+			Some(relay_parent),
+		);
+		parent = block.header().hash();
+		aligned_blocks.push(block);
+	}
+
+	let mut alt_blocks = Vec::new();
+	let mut parent = included_block.header().hash();
+	for _ in 0..NON_INCLUDED_CHAIN_LEN {
+		let block = build_and_import_block_ext(
+			&client,
+			BlockOrigin::NetworkInitialSync,
+			true,
+			&mut para_import,
+			Some(parent),
+			None,
+			Some(search_relay_parent),
+		);
+		parent = block.header().hash();
+		alt_blocks.push(block);
+	}
+
+	// Ignore alternative branch:
+	for max_depth in 0..=NON_INCLUDED_CHAIN_LEN {
+		let potential_parents = block_on(find_potential_parents(
+			ParentSearchParams {
+				relay_parent: search_relay_parent,
+				para_id: ParaId::from(100),
+				ancestry_lookback: 1, // aligned chain is in ancestry.
+				max_depth,
+				ignore_alternative_branches: true,
+			},
+			&*backend,
+			&relay_chain,
+		))
+		.unwrap();
+
+		assert_eq!(potential_parents.len(), max_depth + 1);
+		let expected_parents: Vec<_> = [&included_block, &in_between_block, &pending_block]
+			.into_iter()
+			.chain(aligned_blocks.iter())
+			.take(max_depth + 1)
+			.collect();
+
+		for i in 0..(max_depth + 1) {
+			let parent = &potential_parents[i];
+			let expected = &expected_parents[i];
+
+			assert_eq!(parent.hash, expected.hash());
+			assert_eq!(&parent.header, expected.header());
+			assert_eq!(parent.depth, i);
+			assert!(parent.aligned_with_pending);
+		}
+	}
+
+	// Do not ignore:
+	for max_depth in 0..=NON_INCLUDED_CHAIN_LEN {
+		let potential_parents = block_on(find_potential_parents(
+			ParentSearchParams {
+				relay_parent: search_relay_parent,
+				para_id: ParaId::from(100),
+				ancestry_lookback: 1, // aligned chain is in ancestry.
+				max_depth,
+				ignore_alternative_branches: false,
+			},
+			&*backend,
+			&relay_chain,
+		))
+		.unwrap();
+
+		let expected_len = 2 * max_depth + 1;
+		assert_eq!(potential_parents.len(), expected_len);
+		let expected_aligned: Vec<_> = [&included_block, &in_between_block, &pending_block]
+			.into_iter()
+			.chain(aligned_blocks.iter())
+			.take(max_depth + 1)
+			.collect();
+		let expected_alt = alt_blocks.iter().take(max_depth);
+
+		let expected_parents: Vec<_> =
+			expected_aligned.clone().into_iter().chain(expected_alt).collect();
+		// Check correctness.
+		assert_eq!(expected_parents.len(), expected_len);
+
+		for i in 0..expected_len {
+			let parent = &potential_parents[i];
+			let expected = expected_parents
+				.iter()
+				.find(|block| block.header().hash() == parent.hash)
+				.expect("missing parent");
+
+			let is_aligned = expected_aligned.contains(&expected);
+
+			assert_eq!(parent.hash, expected.hash());
+			assert_eq!(&parent.header, expected.header());
+
+			assert_eq!(parent.aligned_with_pending, is_aligned);
+		}
+	}
+}
+
 #[test]
 fn find_potential_parents_aligned_with_pending() {
 	sp_tracing::try_init_simple();
@@ -1244,6 +1415,7 @@ fn find_potential_parents_aligned_with_pending() {
 
 	// Do not ignore:
 	for max_depth in 0..=NON_INCLUDED_CHAIN_LEN {
+		log::info!("Ran with max_depth = {max_depth}");
 		let potential_parents = block_on(find_potential_parents(
 			ParentSearchParams {
 				relay_parent: search_relay_parent,
@@ -1271,6 +1443,7 @@ fn find_potential_parents_aligned_with_pending() {
 		// Check correctness.
 		assert_eq!(expected_parents.len(), expected_len);
 
+		potential_parents.iter().for_each(|p| log::info!("result: {:?}", p));
 		for i in 0..expected_len {
 			let parent = &potential_parents[i];
 			let expected = expected_parents
@@ -1283,6 +1456,12 @@ fn find_potential_parents_aligned_with_pending() {
 			assert_eq!(parent.hash, expected.hash());
 			assert_eq!(&parent.header, expected.header());
 
+			log::info!(
+				"Check hash: {:?} expected: {} is: {}",
+				parent.hash,
+				is_aligned,
+				parent.aligned_with_pending,
+			);
 			assert_eq!(parent.aligned_with_pending, is_aligned);
 		}
 	}