Extract candidate backing logic in separate struct

paritytech · eskimor · Nov 11, 2022 · Oct 5, 2022 · Oct 5, 2022 · Oct 7, 2022
commit a97e29a2ba2f32882dbc48a055f7fc36013d9cc5
diff --git a/node/core/dispute-coordinator/src/scraping/candidates.rs b/node/core/dispute-coordinator/src/scraping/candidates.rs
@@ -1,5 +1,5 @@
-use polkadot_primitives::v2::CandidateHash;
-use std::collections::HashMap;
+use polkadot_primitives::v2::{BlockNumber, CandidateHash};
+use std::collections::{BTreeMap, HashMap, HashSet};
 
 /// Keeps `CandidateHash` in reference counted way.
 /// Each `insert` saves a value with `reference count == 1` or increases the reference
@@ -42,7 +42,7 @@ impl RefCountedCandidates {
 }
 
 #[cfg(test)]
-mod tests {
+mod ref_counted_candidates_tests {
 	use super::*;
 	use polkadot_primitives::v2::{BlakeTwo256, HashT};
 
@@ -78,3 +78,116 @@ mod tests {
 		assert!(!container.contains(&one));
 	}
 }
+
+/// Keeps track of backed candidates. Supports `insert`, `remove`, `remove_up_to_height`
+/// and `contains` operations.
+/// This structure has got an undesired side effect. If a candidate is removed explicitly
+/// with `remove` its entries will remain in `backed_candidates_by_block_number` until
+/// `remove_up_to_height` is called. This means we will keep data in
+/// `backed_candidates_by_block_number` a bit longer than necessary.
+/// Unfortunately fixing this is not trivial. `backed_candidates` uses reference counting
+/// because a candidate can be backed on multiple block heights. So `remove` doesn't necessary
+/// removes the candidate - it just decreases its reference count. For this reason it is not
+/// a good idea to remove any entries in `backed_candidates_by_block_number` when `remove` is
+/// called.
+/// One approach for a fix is to introduce `BTreeMap<CandidateHash, HashSet<BlockNumber>>` to
+/// keep track on which candidates exist at a specific key in `backed_candidates_by_block_number`.
+/// However when removing a candidate with `remove` we don't know (and don't care) at which block
+/// height it was backed. So no precise removal is possible. This means that we should either
+/// a) don't do any cleanup in `backed_candidates_by_block_number`
+/// b) add the 2nd `BTreeMap` mentioned above and delete a random block from
+///    `backed_candidates_by_block_number` on each removal.
+/// Option a) sounds more reasonable to me - if we can't do the right thing we'd better do nothing.
+pub struct BackedCandidates {
+	/// Main data structure which keeps the backed candidates we know about. `contains` does
+	/// lookups only here.
+	backed_candidates: RefCountedCandidates,
+	/// Keeps track at which block number a candidate was backed. Used in `remove_up_to_height`.
+	/// Without this tracking we won't be able to 'remove all candidates included in block X and before'.
+	backed_candidates_by_block_number: BTreeMap<BlockNumber, HashSet<CandidateHash>>,
+}
+
+impl BackedCandidates {
+	pub fn new() -> Self {
+		Self {
+			backed_candidates: RefCountedCandidates::new(),
+			backed_candidates_by_block_number: BTreeMap::new(),
+		}
+	}
+
+	pub fn contains(&self, candidate_hash: &CandidateHash) -> bool {
+		self.backed_candidates.contains(candidate_hash)
+	}
+
+	// Removes all candidates up to a given height. The candidates at the block height are NOT removed.
+	pub fn remove_up_to_height(&mut self, height: &BlockNumber) {
+		let not_stale = self.backed_candidates_by_block_number.split_off(&height);
+		let stale = std::mem::take(&mut self.backed_candidates_by_block_number);
+		self.backed_candidates_by_block_number = not_stale;
+		for candidates in stale.values() {
+			for c in candidates {
+				self.backed_candidates.remove(c);
+			}
+		}
+	}
+
+	// Removes a single candidate
+	pub fn remove(&mut self, candidate_hash: &CandidateHash) {
+		self.backed_candidates.remove(candidate_hash);
+	}
+
+	pub fn insert(&mut self, block_number: BlockNumber, candidate_hash: CandidateHash) {
+		self.backed_candidates.insert(candidate_hash);
+		self.backed_candidates_by_block_number
+			.entry(block_number)
+			.or_default()
+			.insert(candidate_hash);
+	}
+
+	// Used only for tests to verify the pruning doesn't leak data.
+	#[cfg(test)]
+	pub fn backed_candidates_by_block_number_is_empty(&self) -> bool {
+		self.backed_candidates_by_block_number.is_empty()
+	}
+
+	#[cfg(test)]
+	pub fn backed_candidates_by_block_number_has_key(&self, key: &BlockNumber) -> bool {
+		self.backed_candidates_by_block_number.contains_key(key)
+	}
+}
+
+#[cfg(test)]
+mod backed_candidates_tests {
+	use super::*;
+	use polkadot_primitives::v2::{BlakeTwo256, HashT};
+
+	#[test]
+	fn explicitly_removed_are_cleaned_from_block_mapping() {
+		let mut candidates = BackedCandidates::new();
+		let target = CandidateHash(BlakeTwo256::hash(&vec![1, 2, 3]));
+		candidates.insert(1, target);
+
+		assert!(candidates.contains(&target));
+
+		candidates.remove(&target);
+		assert!(!candidates.contains(&target));
+		assert!(candidates.backed_candidates_by_block_number_has_key(&1)); // undesired side effect
+
+		candidates.remove_up_to_height(&2);
+		assert!(!candidates.backed_candidates_by_block_number_has_key(&1));
+		assert!(candidates.backed_candidates_by_block_number_is_empty());
+	}
+
+	#[test]
+	fn stale_candidates_are_removed() {
+		let mut candidates = BackedCandidates::new();
+		let target = CandidateHash(BlakeTwo256::hash(&vec![1, 2, 3]));
+		candidates.insert(1, target);
+
+		assert!(candidates.contains(&target));
+
+		candidates.remove_up_to_height(&2);
+		assert!(!candidates.contains(&target));
+		assert!(candidates.backed_candidates_by_block_number_is_empty());
+	}
+}
diff --git a/node/core/dispute-coordinator/src/scraping/mod.rs b/node/core/dispute-coordinator/src/scraping/mod.rs
@@ -66,37 +66,27 @@ const LRU_OBSERVED_BLOCKS_CAPACITY: NonZeroUsize = match NonZeroUsize::new(20) {
 ///
 /// With this information it provides a `CandidateComparator` and as a return value of
 /// `process_active_leaves_update` any scraped votes.
+///
+/// If a candidate was backed and included - it will be removed once the block at its height
+/// is finalized. However a candidate can be backed and never included. To avoid
+/// leaking memory in such cases we keep track at which block height the candidate was backed.
+/// We remove backed (and
+/// not included) candidates when `BACKED_CANDIDATE_MAX_LIFETIME` more blocks are
+/// finalized after the block they were initially backed in.
+/// E.g. if `BACKED_CANDIDATE_MAX_LIFETIME = 2` when block 4 is finalized we will remove
+/// all candidates backed in block 2 which were not already cleaned up on finalization.
+///
+/// Please note that if a candidate is backed AND included it will be cleaned up on
+/// finalization and won't be kept for `BACKED_CANDIDATE_MAX_LIFETIME` blocks.
 pub struct ChainScraper {
 	/// All candidates we have seen included, which not yet have been finalized.
 	included_candidates: candidates::RefCountedCandidates,
 	/// All candidates we have seen backed
-	backed_candidates: candidates::RefCountedCandidates,
+	backed_candidates: candidates::BackedCandidates,
 	/// including block -> `CandidateHash`
 	///
 	/// We need this to clean up `included_candidates` on finalization.
 	included_candidates_by_block_number: BTreeMap<BlockNumber, HashSet<CandidateHash>>,
-	/// Keeps track at which block a `CandidateHash` was backed. If a candidate was backed
-	/// and included - it will be removed once the block at its height is finalized. However
-	/// a candidate can be backed and never included. To avoid leaking memory in such cases
-	/// we keep track at which block height the candidate was backed. We remove backed (and
-	/// not included) candidates when `BACKED_CANDIDATE_MAX_LIFETIME` more blocks are
-	/// finalized after the block they were initially backed in.
-	/// E.g. if `BACKED_CANDIDATE_MAX_LIFETIME = 2` when block 4 is finalized we will remove
-	/// all candidates backed in block 2 which were not already cleaned up on finalization.
-	///
-	/// Please note that if a candidate is backed AND included it will be cleaned up on
-	/// finalization thanks to the tracking in `included_candidates_by_block_number`.
-	///
-	/// Also please note that this `BTreeMap` keeps information when a candidate was backed
-	/// only to guarantee that it will be removed at some point instead of staying in the
-	/// `BTreeMap` forever. If a candidate was included it will be removed from
-	/// `backed_candidates` but will remain here until `BACKED_CANDIDATE_MAX_LIFETIME`
-	/// blocks are finalized from its inclusion.
-	/// An alternative to fix this is to keep track of yet another mapping -
-	/// `BTreeMap<CandidateHash, HashSet<BlockNumber>>` so that we can lookup when a candidate
-	/// was backed and clean it up. It's better to have a few stale entries in one `BTreeMap`
-	/// instead of keeping three copies of each candidate.
-	backed_candidates_by_block_number: BTreeMap<BlockNumber, HashSet<CandidateHash>>,
 	/// Latest relay blocks observed by the provider.
 	///
 	/// We assume that ancestors of cached blocks are already processed, i.e. we have saved
@@ -128,9 +118,8 @@ impl ChainScraper {
 	{
 		let mut s = Self {
 			included_candidates: candidates::RefCountedCandidates::new(),
-			backed_candidates: candidates::RefCountedCandidates::new(),
+			backed_candidates: candidates::BackedCandidates::new(),
 			included_candidates_by_block_number: BTreeMap::new(),
-			backed_candidates_by_block_number: BTreeMap::new(),
 			last_observed_blocks: LruCache::new(LRU_OBSERVED_BLOCKS_CAPACITY),
 		};
 		let update =
@@ -200,7 +189,20 @@ impl ChainScraper {
 	/// finalized, we can treat it as low priority.
 	pub fn process_finalized_block(&mut self, finalized_block_number: &BlockNumber) {
 		// First remove any stale backed and not included candidates
-		self.remove_stale_backed_candidates(&finalized_block_number);
+		// Stale means any backed but not included candidates after `BACKED_CANDIDATE_MAX_LIFETIME`
+		// has passed since the (finalized) block containing the corresponding `CandidateBacked`
+		// event.
+		// `BACKED_CANDIDATE_MAX_LIFETIME - 1` because `finalized_block_number` counts to the backed
+		// candidate lifetime. Example:
+		// finalized_block_number = 4; BACKED_CANDIDATE_MAX_LIFETIME = 2;
+		// key_to_clean = 4 - (2 - 1) = 3
+		// After `remove_all_at_block_height` at 3:
+		// 0, 1, 2 will be removed
+		// 3, 4 will be kept
+		// => We keep backed candidates in the last two finalized blocks
+		finalized_block_number
+			.checked_sub(Self::BACKED_CANDIDATE_MAX_LIFETIME - 1)
+			.map(|key_to_clean| self.backed_candidates.remove_up_to_height(&key_to_clean));
 
 		// Then handle included candidates
 		let not_finalized =
@@ -210,46 +212,10 @@ impl ChainScraper {
 		// Clean up finalized
 		for finalized_candidate in finalized.into_values().flatten() {
 			self.included_candidates.remove(&finalized_candidate);
-			// Remove the candidate from backed. Note that the candidate will remain in
-			// `backed_candidates_by_block_number`. This is not desirable but acceptable trade off.
-			// The `BTreeMap` is used to track stale candidates only. Decision if a candidate is backed
-			// or not is made by looking it up in `backed_candidates`.
-			// Entries in `backed_candidates_by_block_number` will be removed by `remove_stale_backed_candidates`
 			self.backed_candidates.remove(&finalized_candidate);
 		}
 	}
 
-	/// Remove any backed but not included candidates after `BACKED_CANDIDATE_MAX_LIFETIME` has passed
-	/// since the (finalized) block containing the corresponding `CandidateBacked` event.
-	/// Used in `process_finalized_block`.
-	fn remove_stale_backed_candidates(&mut self, finalized_block_number: &BlockNumber) {
-		// Cleanup stale backed and unincluded blocks. We want to keep the candidates backed in the last
-		// `BACKED_CANDIDATE_MAX_LIFETIME` finalzied blocks.
-		// We subtract `BACKED_CANDIDATE_MAX_LIFETIME - 1` from the finalzied block number to get the correct
-		// boundry for `split_off`.
-		//
-		// Example:
-		// finalized_block_number = 4; BACKED_CANDIDATE_MAX_LIFETIME = 2;
-		// boundry = 4 - (2 - 1) = 3
-		// After `split_off` at 3:
-		// 0, 1, 2 will be in stale
-		// 3, 4 will be in not stale
-		// => We keep backed candidates in the last two finalized blocks
-		finalized_block_number.checked_sub(Self::BACKED_CANDIDATE_MAX_LIFETIME - 1).map(
-			|key_to_clean| {
-				// let key_to_clean = key_to_clean + 1;
-				let not_stale = self.backed_candidates_by_block_number.split_off(&(key_to_clean));
-				let stale = std::mem::take(&mut self.backed_candidates_by_block_number);
-				self.backed_candidates_by_block_number = not_stale;
-				for candidates in stale.values() {
-					for c in candidates {
-						self.backed_candidates.remove(c);
-					}
-				}
-			},
-		);
-	}
-
 	/// Process candidate events of a block.
 	///
 	/// Keep track of all included and backed candidates.
@@ -287,11 +253,7 @@ impl ChainScraper {
 						?block_number,
 						"Processing backed event"
 					);
-					self.backed_candidates.insert(candidate_hash);
-					self.backed_candidates_by_block_number
-						.entry(block_number)
-						.or_default()
-						.insert(candidate_hash);
+					self.backed_candidates.insert(block_number, candidate_hash);
 				},
 				_ => {
 					// skip the rest
@@ -366,18 +328,6 @@ impl ChainScraper {
 		}
 		return Ok(ancestors)
 	}
-
-	// Used only for tests to verify the pruning doesn't leak data.
-	#[cfg(test)]
-	fn backed_candidates_by_block_number_is_empty(&self) -> bool {
-		println!("{:?}", self.backed_candidates_by_block_number);
-		self.backed_candidates_by_block_number.is_empty()
-	}
-
-	#[cfg(test)]
-	fn backed_candidates_by_block_number_has_key(&self, key: &BlockNumber) -> bool {
-		self.backed_candidates_by_block_number.contains_key(key)
-	}
 }
 
 async fn get_finalized_block_number<Sender>(sender: &mut Sender) -> FatalResult<BlockNumber>

diff --git a/node/core/dispute-coordinator/src/scraping/tests.rs b/node/core/dispute-coordinator/src/scraping/tests.rs
@@ -456,12 +456,6 @@ fn scraper_cleans_finalized_candidates() {
 
 		assert!(!scraper.is_candidate_backed(&candidate.hash()));
 		assert!(!scraper.is_candidate_included(&candidate.hash()));
-
-		// this is undesired side effect - check the comments in `struct ChainScraper` and
-		// `process_finalized_block` for details
-		finalized_block_number += 1;
-		process_finalized_block(&mut scraper, &finalized_block_number);
-		assert!(!scraper.backed_candidates_by_block_number_has_key(&TEST_TARGET_BLOCK_NUMBER));
 	});
 }
 
@@ -512,7 +506,6 @@ fn scraper_handles_backed_but_not_included_candidate() {
 
 		assert!(!scraper.is_candidate_included(&candidate.hash()));
 		assert!(!scraper.is_candidate_backed(&candidate.hash()));
-		assert!(scraper.backed_candidates_by_block_number_is_empty());
 	});
 }
 
@@ -558,10 +551,5 @@ fn scraper_handles_the_same_candidate_incuded_in_two_different_block_heights() {
 
 		assert!(!scraper.is_candidate_backed(&magic_candidate.hash()));
 		assert!(!scraper.is_candidate_included(&magic_candidate.hash()));
-
-		// On the next finalization all backed candidates should be removed
-		finalized_block_number += 1;
-		process_finalized_block(&mut scraper, &finalized_block_number);
-		assert!(scraper.backed_candidates_by_block_number_is_empty());
 	});
 }