Treat backed and included candidates in the same fashion

paritytech · eskimor · Nov 11, 2022 · Oct 5, 2022 · Oct 5, 2022 · Oct 7, 2022
commit 0d498dc55b258158951d308b0086bb44df597be5
diff --git a/node/core/dispute-coordinator/src/scraping/candidates.rs b/node/core/dispute-coordinator/src/scraping/candidates.rs
@@ -79,115 +79,70 @@ mod ref_counted_candidates_tests {
 	}
 }
 
-/// 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>>,
+/// Keeps track of scraoed candidates. Supports `insert`, `remove_up_to_height` and `contains`
+/// operations.
+pub struct ScrapedCandidates {
+	/// Main data structure which keeps the candidates we know about. `contains` does lookups only here.
+	candidates: RefCountedCandidates,
+	/// Keeps track at which block number a candidate was inserted. Used in `remove_up_to_height`.
+	/// Without this tracking we won't be able to 'remove all candidates inserted at block X and before'.
+	candidates_by_block_number: BTreeMap<BlockNumber, HashSet<CandidateHash>>,
 }
 
-impl BackedCandidates {
+impl ScrapedCandidates {
 	pub fn new() -> Self {
 		Self {
-			backed_candidates: RefCountedCandidates::new(),
-			backed_candidates_by_block_number: BTreeMap::new(),
+			candidates: RefCountedCandidates::new(),
+			candidates_by_block_number: BTreeMap::new(),
 		}
 	}
 
 	pub fn contains(&self, candidate_hash: &CandidateHash) -> bool {
-		self.backed_candidates.contains(candidate_hash)
+		self.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;
+		let not_stale = self.candidates_by_block_number.split_off(&height);
+		let stale = std::mem::take(&mut self.candidates_by_block_number);
+		self.candidates_by_block_number = not_stale;
 		for candidates in stale.values() {
 			for c in candidates {
-				self.backed_candidates.remove(c);
+				self.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
+		self.candidates.insert(candidate_hash);
+		self.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)
+	pub fn candidates_by_block_number_is_empty(&self) -> bool {
+		self.candidates_by_block_number.is_empty()
 	}
 }
 
 #[cfg(test)]
-mod backed_candidates_tests {
+mod scraped_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 mut candidates = ScrapedCandidates::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());
+		assert!(candidates.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
@@ -14,10 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
-use std::{
-	collections::{BTreeMap, HashSet},
-	num::NonZeroUsize,
-};
+use std::num::NonZeroUsize;
 
 use futures::channel::oneshot;
 use lru::LruCache;
@@ -71,22 +68,18 @@ const LRU_OBSERVED_BLOCKS_CAPACITY: NonZeroUsize = match NonZeroUsize::new(20) {
 /// 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
+/// not included) candidates when `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
+/// E.g. if `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.
+/// finalization and won't be kept for `CANDIDATE_MAX_LIFETIME` blocks.
 pub struct ChainScraper {
 	/// All candidates we have seen included, which not yet have been finalized.
-	included_candidates: candidates::RefCountedCandidates,
+	included_candidates: candidates::ScrapedCandidates,
 	/// All candidates we have seen backed
-	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>>,
+	backed_candidates: candidates::ScrapedCandidates,
 	/// Latest relay blocks observed by the provider.
 	///
 	/// We assume that ancestors of cached blocks are already processed, i.e. we have saved
@@ -104,7 +97,7 @@ impl ChainScraper {
 
 	/// How many blocks after finalization a backed and not included candidate should
 	/// be kept.
-	pub(crate) const BACKED_CANDIDATE_MAX_LIFETIME: BlockNumber = 2;
+	pub(crate) const CANDIDATE_MAX_LIFETIME: BlockNumber = 2;
 
 	/// Create a properly initialized `OrderingProvider`.
 	///
@@ -117,9 +110,8 @@ impl ChainScraper {
 		Sender: overseer::DisputeCoordinatorSenderTrait,
 	{
 		let mut s = Self {
-			included_candidates: candidates::RefCountedCandidates::new(),
-			backed_candidates: candidates::BackedCandidates::new(),
-			included_candidates_by_block_number: BTreeMap::new(),
+			included_candidates: candidates::ScrapedCandidates::new(),
+			backed_candidates: candidates::ScrapedCandidates::new(),
 			last_observed_blocks: LruCache::new(LRU_OBSERVED_BLOCKS_CAPACITY),
 		};
 		let update =
@@ -185,35 +177,23 @@ impl ChainScraper {
 
 	/// Prune finalized candidates.
 	///
-	/// Once a candidate lives in a relay chain block that's behind the finalized chain/got
-	/// finalized, we can treat it as low priority.
+	/// Once a candidate lives in a relay chain block that's behind `CANDIDATE_MAX_LIFETIME` blocks
+	/// after finalization 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
-		// 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;
+		// `CANDIDATE_MAX_LIFETIME - 1` because `finalized_block_number` counts to the candidate lifetime.
+		// Example:
+		// finalized_block_number = 4; 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
+		// => We keep 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 =
-			self.included_candidates_by_block_number.split_off(finalized_block_number);
-		let finalized = std::mem::take(&mut self.included_candidates_by_block_number);
-		self.included_candidates_by_block_number = not_finalized;
-		// Clean up finalized
-		for finalized_candidate in finalized.into_values().flatten() {
-			self.included_candidates.remove(&finalized_candidate);
-			self.backed_candidates.remove(&finalized_candidate);
-		}
+			.checked_sub(Self::CANDIDATE_MAX_LIFETIME - 1)
+			.map(|key_to_clean| {
+				self.backed_candidates.remove_up_to_height(&key_to_clean);
+				self.included_candidates.remove_up_to_height(&key_to_clean);
+			});
 	}
 
 	/// Process candidate events of a block.
@@ -239,11 +219,7 @@ impl ChainScraper {
 						?block_number,
 						"Processing included event"
 					);
-					self.included_candidates.insert(candidate_hash);
-					self.included_candidates_by_block_number
-						.entry(block_number)
-						.or_default()
-						.insert(candidate_hash);
+					self.included_candidates.insert(block_number, candidate_hash);
 				},
 				CandidateEvent::CandidateBacked(receipt, _, _, _) => {
 					let candidate_hash = receipt.hash();

diff --git a/node/core/dispute-coordinator/src/scraping/tests.rs b/node/core/dispute-coordinator/src/scraping/tests.rs
@@ -451,7 +451,7 @@ fn scraper_cleans_finalized_candidates() {
 		// Finalize the next block and verify that candidate is neither backed nor included.
 		// Backed AND Included candidates should be cleaned on finalization.
 		// `scraper_handles_backed_but_not_included_candidate` covers Backed and NOT Included candidates.
-		finalized_block_number += 1;
+		finalized_block_number += ChainScraper::CANDIDATE_MAX_LIFETIME + 1;
 		process_finalized_block(&mut scraper, &finalized_block_number);
 
 		assert!(!scraper.is_candidate_backed(&candidate.hash()));
@@ -500,8 +500,7 @@ fn scraper_handles_backed_but_not_included_candidate() {
 
 		// Bump the finalized block outside `BACKED_CANDIDATE_MAX_LIFETIME`.
 		// The candidate should be removed.
-		finalized_block_number +=
-			FIRST_TEST_BLOCK + ChainScraper::BACKED_CANDIDATE_MAX_LIFETIME + 1;
+		finalized_block_number += FIRST_TEST_BLOCK + ChainScraper::CANDIDATE_MAX_LIFETIME + 1;
 		process_finalized_block(&mut scraper, &finalized_block_number);
 
 		assert!(!scraper.is_candidate_included(&candidate.hash()));
@@ -538,7 +537,7 @@ fn scraper_handles_the_same_candidate_incuded_in_two_different_block_heights() {
 
 		// Finalize blocks to enforce pruning of scraped events.
 		// The magic candidate was added twice, so it shouldn't be removed if we finalize two more blocks.
-		finalized_block_number += BLOCKS_TO_SKIP as u32 - finalized_block_number;
+		finalized_block_number += ChainScraper::CANDIDATE_MAX_LIFETIME + 1; //BLOCKS_TO_SKIP as u32 - finalized_block_number;
 		process_finalized_block(&mut scraper, &finalized_block_number);
 
 		let magic_candidate = make_candidate_receipt(get_magic_candidate_hash());