extract determine_new_blocks into a separate utility (#3261)

* extract determine_new_blocks into a separate utility

* rework docs

polkadot/node/core/approval-voting/src/import.rs

@@ -34,6 +34,7 @@ use polkadot_node_subsystem::{
 	},
 	SubsystemContext, SubsystemError, SubsystemResult,
 };
+use polkadot_node_subsystem_util::determine_new_blocks;
 use polkadot_node_subsystem_util::rolling_session_window::{
 	RollingSessionWindow, SessionWindowUpdate,
 };
@@ -63,112 +64,6 @@ use crate::time::{slot_number_to_tick, Tick};
 
 use super::{LOG_TARGET, State, DBReader};
 
-// Given a new chain-head hash, this determines the hashes of all new blocks we should track
-// metadata for, given this head. The list will typically include the `head` hash provided unless
-// that block is already known, in which case the list should be empty. This is guaranteed to be
-// a subset of the ancestry of `head`, as well as `head`, starting from `head` and moving
-// backwards.
-//
-// This returns the entire ancestry up to the last finalized block's height or the last item we
-// have in the DB. This may be somewhat expensive when first recovering from major sync.
-async fn determine_new_blocks(
-	ctx: &mut impl SubsystemContext,
-	db: &impl DBReader,
-	head: Hash,
-	header: &Header,
-	finalized_number: BlockNumber,
-) -> SubsystemResult<Vec<(Hash, Header)>> {
-	const ANCESTRY_STEP: usize = 4;
-
-	// Early exit if the block is in the DB or too early.
-	{
-		let already_known = db.load_block_entry(&head)?
-			.is_some();
-
-		let before_relevant = header.number <= finalized_number;
-
-		if already_known || before_relevant {
-			return Ok(Vec::new());
-		}
-	}
-
-	let mut ancestry = vec![(head, header.clone())];
-
-	// Early exit if the parent hash is in the DB.
-	if db.load_block_entry(&header.parent_hash)?
-		.is_some()
-	{
-		return Ok(ancestry);
-	}
-
-	'outer: loop {
-		let &(ref last_hash, ref last_header) = ancestry.last()
-			.expect("ancestry has length 1 at initialization and is only added to; qed");
-
-		// If we iterated back to genesis, which can happen at the beginning of chains.
-		if last_header.number <= 1 {
-			break 'outer
-		}
-
-		let (tx, rx) = oneshot::channel();
-		ctx.send_message(ChainApiMessage::Ancestors {
-			hash: *last_hash,
-			k: ANCESTRY_STEP,
-			response_channel: tx,
-		}.into()).await;
-
-		// Continue past these errors.
-		let batch_hashes = match rx.await {
-			Err(_) | Ok(Err(_)) => break 'outer,
-			Ok(Ok(ancestors)) => ancestors,
-		};
-
-		let batch_headers = {
-			let (batch_senders, batch_receivers) = (0..batch_hashes.len())
-				.map(|_| oneshot::channel())
-				.unzip::<_, _, Vec<_>, Vec<_>>();
-
-			for (hash, sender) in batch_hashes.iter().cloned().zip(batch_senders) {
-				ctx.send_message(ChainApiMessage::BlockHeader(hash, sender).into()).await;
-			}
-
-			let mut requests = futures::stream::FuturesOrdered::new();
-			batch_receivers.into_iter().map(|rx| async move {
-				match rx.await {
-					Err(_) | Ok(Err(_)) => None,
-					Ok(Ok(h)) => h,
-				}
-			})
-				.for_each(|x| requests.push(x));
-
-			let batch_headers: Vec<_> = requests
-				.flat_map(|x: Option<Header>| stream::iter(x))
-				.collect()
-				.await;
-
-			// Any failed header fetch of the batch will yield a `None` result that will
-			// be skipped. Any failure at this stage means we'll just ignore those blocks
-			// as the chain DB has failed us.
-			if batch_headers.len() != batch_hashes.len() { break 'outer }
-			batch_headers
-		};
-
-		for (hash, header) in batch_hashes.into_iter().zip(batch_headers) {
-			let is_known = db.load_block_entry(&hash)?.is_some();
-
-			let is_relevant = header.number > finalized_number;
-
-			if is_known || !is_relevant {
-				break 'outer
-			}
-
-			ancestry.push((hash, header));
-		}
-	}
-
-	Ok(ancestry)
-}
-
 struct ImportedBlockInfo {
 	included_candidates: Vec<(CandidateHash, CandidateReceipt, CoreIndex, GroupIndex)>,
 	session_index: SessionIndex,
@@ -445,9 +340,15 @@ pub(crate) async fn handle_new_head(
 
 	// If we've just started the node and haven't yet received any finality notifications,
 	// we don't do any look-back. Approval voting is only for nodes were already online.
-	let finalized_number = finalized_number.unwrap_or(header.number.saturating_sub(1));
+	let lower_bound_number = finalized_number.unwrap_or(header.number.saturating_sub(1));
 
-	let new_blocks = determine_new_blocks(ctx, &state.db, head, &header, finalized_number)
+	let new_blocks = determine_new_blocks(
+		ctx.sender(),
+		|h| state.db.load_block_entry(h).map(|e| e.is_some()),
+		head,
+		&header,
+		lower_bound_number,
+	)
 		.map_err(|e| SubsystemError::with_origin("approval-voting", e))
 		.await?;
 
@@ -741,86 +642,6 @@ mod tests {
 		}
 	}
 
-	#[derive(Clone)]
-	struct TestChain {
-		start_number: BlockNumber,
-		headers: Vec<Header>,
-		numbers: HashMap<Hash, BlockNumber>,
-	}
-
-	impl TestChain {
-		fn new(start: BlockNumber, len: usize) -> Self {
-			assert!(len > 0, "len must be at least 1");
-
-			let base = Header {
-				digest: Default::default(),
-				extrinsics_root: Default::default(),
-				number: start,
-				state_root: Default::default(),
-				parent_hash: Default::default(),
-			};
-
-			let base_hash = base.hash();
-
-			let mut chain = TestChain {
-				start_number: start,
-				headers: vec![base],
-				numbers: vec![(base_hash, start)].into_iter().collect(),
-			};
-
-			for _ in 1..len {
-				chain.grow()
-			}
-
-			chain
-		}
-
-		fn grow(&mut self) {
-			let next = {
-				let last = self.headers.last().unwrap();
-				Header {
-					digest: Default::default(),
-					extrinsics_root: Default::default(),
-					number: last.number + 1,
-					state_root: Default::default(),
-					parent_hash: last.hash(),
-				}
-			};
-
-			self.numbers.insert(next.hash(), next.number);
-			self.headers.push(next);
-		}
-
-		fn header_by_number(&self, number: BlockNumber) -> Option<&Header> {
-			if number < self.start_number {
-				None
-			} else {
-				self.headers.get((number - self.start_number) as usize)
-			}
-		}
-
-		fn header_by_hash(&self, hash: &Hash) -> Option<&Header> {
-			self.numbers.get(hash).and_then(|n| self.header_by_number(*n))
-		}
-
-		fn hash_by_number(&self, number: BlockNumber) -> Option<Hash> {
-			self.header_by_number(number).map(|h| h.hash())
-		}
-
-		fn ancestry(&self, hash: &Hash, k: BlockNumber) -> Vec<Hash> {
-			let n = match self.numbers.get(hash) {
-				None => return Vec::new(),
-				Some(&n) => n,
-			};
-
-			(0..k)
-				.map(|i| i + 1)
-				.filter_map(|i| self.header_by_number(n - i))
-				.map(|h| h.hash())
-				.collect()
-		}
-	}
-
 	struct MockAssignmentCriteria;
 
 	impl AssignmentCriteria for MockAssignmentCriteria {
@@ -856,340 +677,6 @@ mod tests {
 		(VRFOutput(o.to_output()), VRFProof(p))
 	}
 
-	#[test]
-	fn determine_new_blocks_back_to_finalized() {
-		let pool = TaskExecutor::new();
-		let (mut ctx, mut handle) = make_subsystem_context::<(), _>(pool.clone());
-
-		let db = TestDB::default();
-
-		let chain = TestChain::new(10, 9);
-
-		let head = chain.header_by_number(18).unwrap().clone();
-		let head_hash = head.hash();
-		let finalized_number = 12;
-
-		// Finalized block should be omitted. The head provided to `determine_new_blocks`
-		// should be included.
-		let expected_ancestry = (13..=18)
-			.map(|n| chain.header_by_number(n).map(|h| (h.hash(), h.clone())).unwrap())
-			.rev()
-			.collect::<Vec<_>>();
-
-		let test_fut = Box::pin(async move {
-			let ancestry = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				finalized_number,
-			).await.unwrap();
-
-			assert_eq!(
-				ancestry,
-				expected_ancestry,
-			);
-		});
-
-		let aux_fut = Box::pin(async move {
-			assert_matches!(
-				handle.recv().await,
-				AllMessages::ChainApi(ChainApiMessage::Ancestors {
-					hash: h,
-					k,
-					response_channel: tx,
-				}) => {
-					assert_eq!(h, head_hash);
-					assert_eq!(k, 4);
-					let _ = tx.send(Ok(chain.ancestry(&h, k as _)));
-				}
-			);
-
-			for _ in 0..4 {
-				assert_matches!(
-					handle.recv().await,
-					AllMessages::ChainApi(ChainApiMessage::BlockHeader(h, tx)) => {
-						let _ = tx.send(Ok(chain.header_by_hash(&h).map(|h| h.clone())));
-					}
-				);
-			}
-
-			assert_matches!(
-				handle.recv().await,
-				AllMessages::ChainApi(ChainApiMessage::Ancestors {
-					hash: h,
-					k,
-					response_channel: tx,
-				}) => {
-					assert_eq!(h, chain.hash_by_number(14).unwrap());
-					assert_eq!(k, 4);
-					let _ = tx.send(Ok(chain.ancestry(&h, k as _)));
-				}
-			);
-
-			for _ in 0..4 {
-				assert_matches!(
-					handle.recv().await,
-					AllMessages::ChainApi(ChainApiMessage::BlockHeader(h, tx)) => {
-						let _ = tx.send(Ok(chain.header_by_hash(&h).map(|h| h.clone())));
-					}
-				);
-			}
-
-		});
-
-		futures::executor::block_on(futures::future::join(test_fut, aux_fut));
-	}
-
-	#[test]
-	fn determine_new_blocks_back_to_known() {
-		let pool = TaskExecutor::new();
-		let (mut ctx, mut handle) = make_subsystem_context::<(), _>(pool.clone());
-
-		let mut db = TestDB::default();
-
-		let chain = TestChain::new(10, 9);
-
-		let head = chain.header_by_number(18).unwrap().clone();
-		let head_hash = head.hash();
-		let finalized_number = 12;
-		let known_number = 15;
-		let known_hash = chain.hash_by_number(known_number).unwrap();
-
-		db.block_entries.insert(
-			known_hash,
-			crate::approval_db::v1::BlockEntry {
-				block_hash: known_hash,
-				parent_hash: Default::default(),
-				block_number: known_number,
-				session: 1,
-				slot: Slot::from(100),
-				relay_vrf_story: Default::default(),
-				candidates: Vec::new(),
-				approved_bitfield: Default::default(),
-				children: Vec::new(),
-			}.into(),
-		);
-
-		// Known block should be omitted. The head provided to `determine_new_blocks`
-		// should be included.
-		let expected_ancestry = (16..=18)
-			.map(|n| chain.header_by_number(n).map(|h| (h.hash(), h.clone())).unwrap())
-			.rev()
-			.collect::<Vec<_>>();
-
-		let test_fut = Box::pin(async move {
-			let ancestry = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				finalized_number,
-			).await.unwrap();
-
-			assert_eq!(
-				ancestry,
-				expected_ancestry,
-			);
-		});
-
-		let aux_fut = Box::pin(async move {
-			assert_matches!(
-				handle.recv().await,
-				AllMessages::ChainApi(ChainApiMessage::Ancestors {
-					hash: h,
-					k,
-					response_channel: tx,
-				}) => {
-					assert_eq!(h, head_hash);
-					assert_eq!(k, 4);
-					let _ = tx.send(Ok(chain.ancestry(&h, k as _)));
-				}
-			);
-
-			for _ in 0u32..4 {
-				assert_matches!(
-					handle.recv().await,
-					AllMessages::ChainApi(ChainApiMessage::BlockHeader(h, tx)) => {
-						let _ = tx.send(Ok(chain.header_by_hash(&h).map(|h| h.clone())));
-					}
-				);
-			}
-		});
-
-		futures::executor::block_on(futures::future::join(test_fut, aux_fut));
-	}
-
-	#[test]
-	fn determine_new_blocks_already_known_is_empty() {
-		let pool = TaskExecutor::new();
-		let (mut ctx, _handle) = make_subsystem_context::<(), _>(pool.clone());
-
-		let mut db = TestDB::default();
-
-		let chain = TestChain::new(10, 9);
-
-		let head = chain.header_by_number(18).unwrap().clone();
-		let head_hash = head.hash();
-		let finalized_number = 0;
-
-		db.block_entries.insert(
-			head_hash,
-			crate::approval_db::v1::BlockEntry {
-				block_hash: head_hash,
-				parent_hash: Default::default(),
-				block_number: 18,
-				session: 1,
-				slot: Slot::from(100),
-				relay_vrf_story: Default::default(),
-				candidates: Vec::new(),
-				approved_bitfield: Default::default(),
-				children: Vec::new(),
-			}.into(),
-		);
-
-		// Known block should be omitted.
-		let expected_ancestry = Vec::new();
-
-		let test_fut = Box::pin(async move {
-			let ancestry = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				finalized_number,
-			).await.unwrap();
-
-			assert_eq!(
-				ancestry,
-				expected_ancestry,
-			);
-		});
-
-		futures::executor::block_on(test_fut);
-	}
-
-	#[test]
-	fn determine_new_blocks_parent_known_is_fast() {
-		let pool = TaskExecutor::new();
-		let (mut ctx, _handle) = make_subsystem_context::<(), _>(pool.clone());
-
-		let mut db = TestDB::default();
-
-		let chain = TestChain::new(10, 9);
-
-		let head = chain.header_by_number(18).unwrap().clone();
-		let head_hash = head.hash();
-		let finalized_number = 0;
-		let parent_hash = chain.hash_by_number(17).unwrap();
-
-		db.block_entries.insert(
-			parent_hash,
-			crate::approval_db::v1::BlockEntry {
-				block_hash: parent_hash,
-				parent_hash: Default::default(),
-				block_number: 18,
-				session: 1,
-				slot: Slot::from(100),
-				relay_vrf_story: Default::default(),
-				candidates: Vec::new(),
-				approved_bitfield: Default::default(),
-				children: Vec::new(),
-			}.into(),
-		);
-
-		// New block should be the only new one.
-		let expected_ancestry = vec![(head_hash, head.clone())];
-
-		let test_fut = Box::pin(async move {
-			let ancestry = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				finalized_number,
-			).await.unwrap();
-
-			assert_eq!(
-				ancestry,
-				expected_ancestry,
-			);
-		});
-
-		futures::executor::block_on(test_fut);
-	}
-
-	#[test]
-	fn determine_new_block_before_finality_is_empty() {
-		let pool = TaskExecutor::new();
-		let (mut ctx, _handle) = make_subsystem_context::<(), _>(pool.clone());
-
-		let chain = TestChain::new(10, 9);
-
-		let head = chain.header_by_number(18).unwrap().clone();
-		let head_hash = head.hash();
-		let parent_hash = chain.hash_by_number(17).unwrap();
-		let mut db = TestDB::default();
-
-		db.block_entries.insert(
-			parent_hash,
-			crate::approval_db::v1::BlockEntry {
-				block_hash: parent_hash,
-				parent_hash: Default::default(),
-				block_number: 18,
-				session: 1,
-				slot: Slot::from(100),
-				relay_vrf_story: Default::default(),
-				candidates: Vec::new(),
-				approved_bitfield: Default::default(),
-				children: Vec::new(),
-			}.into(),
-		);
-
-		let test_fut = Box::pin(async move {
-			let after_finality = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				17,
-			).await.unwrap();
-
-			let at_finality = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				18,
-			).await.unwrap();
-
-			let before_finality = determine_new_blocks(
-				&mut ctx,
-				&db,
-				head_hash,
-				&head,
-				19,
-			).await.unwrap();
-
-			assert_eq!(
-				after_finality,
-				vec![(head_hash, head.clone())],
-			);
-
-			assert_eq!(
-				at_finality,
-				Vec::new(),
-			);
-
-			assert_eq!(
-				before_finality,
-				Vec::new(),
-			);
-		});
-
-		futures::executor::block_on(test_fut);
-	}
-
 	fn dummy_session_info(index: SessionIndex) -> SessionInfo {
 		SessionInfo {
 			validators: Vec::new(),
@@ -1678,8 +1165,7 @@ mod tests {
 
 		let slot = Slot::from(10);
 
-		let chain = TestChain::new(4, 1);
-		let parent_hash = chain.header_by_number(4).unwrap().hash();
+		let parent_hash = Hash::repeat_byte(0x01);
 
 		let header = Header {
 			digest: {