runtime/parachains: Add pruning logic

runtime/parachains/src/disputes.rs

@@ -619,10 +619,12 @@ impl<T: Config> Pallet<T> {
 			for to_prune in to_prune {
 				// This should be small, as disputes are rare, so `None` is fine.
 				<Disputes<T>>::remove_prefix(to_prune, None);
-
-				// This is larger, and will be extracted to the `shared` module for more proper pruning.
-				// TODO: https://github.com/paritytech/polkadot/issues/3469
-				shared::Pallet::<T>::prune_included_candidates(to_prune);
+				// Mark the session as pruneable so that its candidates can be incrementally
+				// removed over the course of many block inclusions.
+				shared::Pallet::<T>::mark_session_pruneable(to_prune);
+				// TODO(ladi): remove this call, currently allows unit tests to pass. Need to
+				// figure out how to invoke paras_inherent::enter in run_to_block.
+				shared::Pallet::<T>::prune_ancient_sessions(shared::MAX_CANDIDATES_TO_PRUNE);
 				SpamSlots::<T>::remove(to_prune);
 			}
 

runtime/parachains/src/paras_inherent.rs

@@ -222,11 +222,7 @@ pub mod pallet {
 					None => continue,
 				};
 
-				T::DisputesHandler::process_included(
-					current_session,
-					*candidate_hash,
-					revert_to,
-				);
+				T::DisputesHandler::process_included(current_session, *candidate_hash, revert_to);
 			}
 
 			// Handle timeouts for any availability core work.
@@ -281,6 +277,9 @@ pub mod pallet {
 			// And track that we've finished processing the inherent for this block.
 			Included::<T>::set(Some(()));
 
+			// Prune candidates incrementally with each block inclusion.
+			shared::Pallet::<T>::prune_ancient_sessions(shared::MAX_CANDIDATES_TO_PRUNE);
+
 			Ok(Some(
 				MINIMAL_INCLUSION_INHERENT_WEIGHT +
 					(backed_candidates_len * BACKED_CANDIDATE_WEIGHT),

runtime/parachains/src/shared.rs

@@ -36,6 +36,11 @@ pub use pallet::*;
 // which guarantees that at least one full session has passed before any changes are applied.
 pub(crate) const SESSION_DELAY: SessionIndex = 2;
 
+// `MAX_CANDIDATES_TO_PRUNE` is used to upper bound the number of ancient candidates that can be
+// pruned when a new block is included. This is used to distribute the cost of pruning over
+// multiple blocks rather than pruning all historical candidates upon starting a new session.
+pub(crate) const MAX_CANDIDATES_TO_PRUNE: usize = 200;
+
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
@@ -79,6 +84,15 @@ pub mod pallet {
 		(T::BlockNumber, CoreIndex),
 	>;
 
+	/// The set of all past sessions that have yet to fully pruned. Sessions are added to this set
+	/// upon moving to new current session, and removed after all of their included candidates
+	/// have been removed. This set is tracked independently so that the cost of removing included
+	/// candidates can be amortized over multiple blocks rather than removing all candidates upon
+	/// transitioning to a new session.
+	#[pallet::storage]
+	#[pallet::getter(fn pruneable_sessions)]
+	pub(super) type PruneableSessions<T: Config> = StorageMap<_, Twox64Concat, SessionIndex, ()>;
+
 	#[pallet::call]
 	impl<T: Config> Pallet<T> {}
 }
@@ -129,6 +143,58 @@ impl<T: Config> Pallet<T> {
 		Self::session_index().saturating_add(SESSION_DELAY)
 	}
 
+	/// Adds a session that is no longer in the dispute window as pruneable. Subsequent block
+	/// inclusions will incrementally remove old candidates to avoid taking the performance hit of
+	/// removing all candidates at once.
+	pub(crate) fn mark_session_pruneable(session: SessionIndex) {
+		PruneableSessions::<T>::insert(session, ());
+	}
+
+	/// Prunes up to `max_candidates_to_prune` candidates from `IncludedCandidates` that belong to
+	/// non-active sessions.
+	pub(crate) fn prune_ancient_sessions(max_candidates_to_prune: usize) {
+		let mut to_prune = Vec::new();
+		let mut n_candidates = 0;
+		let mut incomplete_session = None;
+		for session in PruneableSessions::<T>::iter_keys() {
+			let mut hashes = Vec::new();
+			for candidate_hash in IncludedCandidates::<T>::iter_key_prefix(session) {
+				// Exit condition when this session still has more candidates to prune; mark the
+				// session as incomplete so the remaining candidates can be pruned in a subsequent
+				// invocation.
+				if n_candidates >= max_candidates_to_prune {
+					incomplete_session = Some(session);
+					break
+				}
+				hashes.push(candidate_hash);
+				n_candidates += 1;
+			}
+
+			to_prune.push((session, hashes));
+			// Exit condition when all candidates from this session were selected for pruning.
+			if n_candidates >= max_candidates_to_prune {
+				break
+			}
+		}
+
+		for (session, candidate_hashes) in to_prune {
+			for candidate_hash in candidate_hashes {
+				IncludedCandidates::<T>::remove(session, candidate_hash);
+			}
+
+			// Prune the session only if it was not marked as incomplete.
+			match incomplete_session {
+				Some(incomplete_session) =>
+					if incomplete_session != session {
+						PruneableSessions::<T>::remove(session);
+					},
+				None => {
+					PruneableSessions::<T>::remove(session);
+				},
+			}
+		}
+	}
+
 	/// Records an included candidate, returning the block height that should be reverted to if the
 	/// block is found to be invalid. This method will return `None` if and only if `included_in`
 	/// is zero.
@@ -164,9 +230,9 @@ impl<T: Config> Pallet<T> {
 		<IncludedCandidates<T>>::get(session, candidate_hash)
 	}
 
-	/// Prunes all candidates that were included in the `to_prune` session.
-	pub(crate) fn prune_included_candidates(to_prune: SessionIndex) {
-		<IncludedCandidates<T>>::remove_prefix(to_prune, None);
+	#[cfg(test)]
+	pub(crate) fn is_pruneable_session(session: &SessionIndex) -> bool {
+		<PruneableSessions<T>>::contains_key(session)
 	}
 
 	#[cfg(test)]
@@ -327,34 +393,174 @@ mod tests {
 	}
 
 	#[test]
-	fn prune_included_candidate_removes_all_candidates_with_same_session() {
+	fn prune_ancient_sessions_no_incomplete_session() {
 		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
 			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
 			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
-			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
+			let block_number = 1;
+			let core_index = CoreIndex(0);
 
-			assert!(
-				ParasShared::note_included_candidate(1, candidate_hash1, 1, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash1,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
-			assert!(
-				ParasShared::note_included_candidate(1, candidate_hash2, 1, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash2,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
+			);
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(2);
+
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session);
+			assert!(ParasShared::is_pruneable_session(&session));
+
+			// Prune the sessions, which should remove both candidates. The session should not be
+			// marked as incomplete since there are exactly two candidates.
+			ParasShared::prune_ancient_sessions(2);
+
+			assert!(!ParasShared::is_pruneable_session(&session));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash2));
+		})
+	}
+
+	#[test]
+	fn prune_ancient_sessions_incomplete_session() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session = 1;
+			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
+			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
+			let block_number = 1;
+			let core_index = CoreIndex(0);
+
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash1,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
-			assert!(
-				ParasShared::note_included_candidate(2, candidate_hash3, 2, CoreIndex(0)).is_some()
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session,
+					candidate_hash2,
+					block_number,
+					core_index,
+				),
+				Some(block_number - 1),
 			);
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 2);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(1);
 
-			ParasShared::prune_included_candidates(1);
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 1);
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session);
+			assert!(ParasShared::is_pruneable_session(&session));
 
-			ParasShared::prune_included_candidates(2);
+			// Prune the sessions, which should remove one of the candidates. The session will be
+			// marked as incomplete so the session should remain unpruned.
+			ParasShared::prune_ancient_sessions(1);
 
-			assert_eq!(ParasShared::included_candidates_iter_prefix(1).count(), 0);
-			assert_eq!(ParasShared::included_candidates_iter_prefix(2).count(), 0);
-		});
+			assert!(ParasShared::is_pruneable_session(&session));
+			assert!(!ParasShared::is_included_candidate(&session, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session, &candidate_hash2));
+		})
+	}
+
+	#[test]
+	fn prune_ancient_sessions_complete_and_incomplete_sessions() {
+		new_test_ext(MockGenesisConfig::default()).execute_with(|| {
+			let session1 = 1;
+			let session2 = 2;
+			let candidate_hash1 = CandidateHash(sp_core::H256::repeat_byte(1));
+			let candidate_hash2 = CandidateHash(sp_core::H256::repeat_byte(2));
+			let candidate_hash3 = CandidateHash(sp_core::H256::repeat_byte(3));
+			let block_number1 = 1;
+			let block_number2 = 2;
+			let core_index = CoreIndex(0);
+
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session1,
+					candidate_hash1,
+					block_number1,
+					core_index,
+				),
+				Some(block_number1 - 1),
+			);
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session2,
+					candidate_hash2,
+					block_number2,
+					core_index,
+				),
+				Some(block_number2 - 1),
+			);
+			assert_eq!(
+				ParasShared::note_included_candidate(
+					session2,
+					candidate_hash3,
+					block_number2,
+					core_index,
+				),
+				Some(block_number2 - 1),
+			);
+			assert!(ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash3));
+
+			// Prune before any sessions are marked pruneable.
+			ParasShared::prune_ancient_sessions(2);
+
+			// Both candidates should still exist.
+			assert!(ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash3));
+
+			// Mark the candidates' session as pruneable.
+			ParasShared::mark_session_pruneable(session1);
+			ParasShared::mark_session_pruneable(session2);
+			assert!(ParasShared::is_pruneable_session(&session1));
+			assert!(ParasShared::is_pruneable_session(&session2));
+
+			// Prune the sessions, which should remove one candidate from each session. The first
+			// session should be pruned while the second session will be marked as incomplete, and
+			// so should remain in the set of pruneable sessions.
+			ParasShared::prune_ancient_sessions(2);
+
+			assert!(!ParasShared::is_pruneable_session(&session1));
+			assert!(ParasShared::is_pruneable_session(&session2));
+			assert!(!ParasShared::is_included_candidate(&session1, &candidate_hash1));
+			assert!(ParasShared::is_included_candidate(&session2, &candidate_hash2));
+			assert!(!ParasShared::is_included_candidate(&session2, &candidate_hash3));
+		})
 	}
 }