Add metric to measure the time it takes to gather enough assignments

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

@@ -607,7 +607,7 @@ pub(crate) mod tests {
 	use super::*;
 	use crate::{
 		approval_db::common::{load_block_entry, DbBackend},
-		RuntimeInfo, RuntimeInfoConfig,
+		RuntimeInfo, RuntimeInfoConfig, MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
 	};
 	use ::test_helpers::{dummy_candidate_receipt, dummy_hash};
 	use assert_matches::assert_matches;
@@ -622,6 +622,7 @@ pub(crate) mod tests {
 		node_features::FeatureIndex, ExecutorParams, Id as ParaId, IndexedVec, NodeFeatures,
 		SessionInfo, ValidatorId, ValidatorIndex,
 	};
+	use schnellru::{ByLength, LruMap};
 	pub(crate) use sp_consensus_babe::{
 		digests::{CompatibleDigestItem, PreDigest, SecondaryVRFPreDigest},
 		AllowedSlots, BabeEpochConfiguration, Epoch as BabeEpoch,
@@ -658,7 +659,9 @@ pub(crate) mod tests {
 			clock: Box::new(MockClock::default()),
 			assignment_criteria: Box::new(MockAssignmentCriteria::default()),
 			spans: HashMap::new(),
-			per_block_assignments_gathering_times: Default::default(),
+			per_block_assignments_gathering_times: LruMap::new(ByLength::new(
+				MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
+			)),
 		}
 	}
 

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

@@ -65,7 +65,9 @@ use sp_consensus::SyncOracle;
 use sp_consensus_slots::Slot;
 use std::time::Instant;
 
-// The maximum block we keep track of assignments gathering times.
+// The maximum block we keep track of assignments gathering times, on normal operation
+// this would never be reached because we prune the data on finalization, but we need
+// to also ensure the data is not growing unecessarily large.
 const MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS: u32 = 100;
 
 use futures::{
@@ -819,7 +821,7 @@ struct State {
 	// assignments, this is relevant because it gives us a good idea about how many
 	// tranches we trigger and why.
 	per_block_assignments_gathering_times:
-		HashMap<BlockNumber, HashMap<(Hash, CandidateHash), AssignmentGatheringRecord>>,
+		LruMap<BlockNumber, HashMap<(Hash, CandidateHash), AssignmentGatheringRecord>>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -949,33 +951,22 @@ impl State {
 		block_hash: Hash,
 		candidate: CandidateHash,
 	) {
-		let record = self
+		if let Some(record) = self
 			.per_block_assignments_gathering_times
-			.entry(block_number)
-			.or_default()
-			.entry((block_hash, candidate))
-			.or_default();
-
-		if record.stage_start.is_none() {
-			record.stage += 1;
-			gum::debug!(
-				target: LOG_TARGET,
-				stage = ?record.stage,
-				?block_hash,
-				?candidate,
-				"Started a new assignment gathering stage",
-			);
-			record.stage_start = Some(Instant::now());
-		}
-
-		// Make sure we always cleanup if we have too many records.
-		if self.per_block_assignments_gathering_times.len() >
-			MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS as usize &&
-			block_number >= MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS
+			.get_or_insert(block_number, HashMap::new)
+			.and_then(|records| Some(records.entry((block_hash, candidate)).or_default()))
 		{
-			self.cleanup_assignments_gathering_timestamp(
-				block_number - MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
-			)
+			if record.stage_start.is_none() {
+				record.stage += 1;
+				gum::debug!(
+					target: LOG_TARGET,
+					stage = ?record.stage,
+					?block_hash,
+					?candidate,
+					"Started a new assignment gathering stage",
+				);
+				record.stage_start = Some(Instant::now());
+			}
 		}
 	}
 
@@ -987,7 +978,7 @@ impl State {
 	) -> AssignmentGatheringRecord {
 		let record = self
 			.per_block_assignments_gathering_times
-			.get_mut(&block_number)
+			.get(&block_number)
 			.and_then(|entry| entry.get_mut(&(block_hash, candidate)));
 		let stage = record.as_ref().map(|record| record.stage).unwrap_or_default();
 		AssignmentGatheringRecord {
@@ -997,8 +988,14 @@ impl State {
 	}
 
 	fn cleanup_assignments_gathering_timestamp(&mut self, keep_greater_than: BlockNumber) {
-		self.per_block_assignments_gathering_times
-			.retain(|block_number, _| *block_number > keep_greater_than)
+		while let Some((block_number, _)) = self.per_block_assignments_gathering_times.peek_oldest()
+		{
+			if *block_number < keep_greater_than {
+				self.per_block_assignments_gathering_times.pop_oldest();
+			} else {
+				break
+			}
+		}
 	}
 
 	fn observe_assignment_gathering_status(
@@ -1077,7 +1074,9 @@ where
 		clock: subsystem.clock,
 		assignment_criteria,
 		spans: HashMap::new(),
-		per_block_assignments_gathering_times: Default::default(),
+		per_block_assignments_gathering_times: LruMap::new(ByLength::new(
+			MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
+		)),
 	};
 
 	// `None` on start-up. Gets initialized/updated on leaf update

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

@@ -5063,7 +5063,9 @@ fn test_gathering_assignments_statements() {
 		clock: Box::new(MockClock::default()),
 		assignment_criteria: Box::new(MockAssignmentCriteria::check_only(|_| Ok(0))),
 		spans: HashMap::new(),
-		per_block_assignments_gathering_times: Default::default(),
+		per_block_assignments_gathering_times: LruMap::new(ByLength::new(
+			MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
+		)),
 	};
 
 	for i in 0..200i32 {
@@ -5078,7 +5080,11 @@ fn test_gathering_assignments_statements() {
 		);
 
 		assert_eq!(
-			state.per_block_assignments_gathering_times.keys().min(),
+			state
+				.per_block_assignments_gathering_times
+				.iter()
+				.map(|(block_number, _)| block_number)
+				.min(),
 			Some(max(0, i - MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS as i32 + 1) as u32).as_ref()
 		)
 	}
@@ -5087,15 +5093,28 @@ fn test_gathering_assignments_statements() {
 		MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS as usize
 	);
 
-	let nothing_changes = state.per_block_assignments_gathering_times.clone();
+	let nothing_changes = state
+		.per_block_assignments_gathering_times
+		.iter()
+		.map(|(block_number, _)| *block_number)
+		.sorted()
+		.collect::<Vec<_>>();
 
 	for i in 150..200i32 {
 		state.mark_begining_of_gathering_assignments(
 			i as u32,
 			Hash::repeat_byte(i as u8),
 			CandidateHash(Hash::repeat_byte(i as u8)),
 		);
-		assert_eq!(nothing_changes, state.per_block_assignments_gathering_times);
+		assert_eq!(
+			nothing_changes,
+			state
+				.per_block_assignments_gathering_times
+				.iter()
+				.map(|(block_number, _)| *block_number)
+				.sorted()
+				.collect::<Vec<_>>()
+		);
 	}
 
 	for i in 110..120 {
@@ -5123,6 +5142,9 @@ fn test_gathering_assignments_statements() {
 		assert!(record.stage_start.is_some());
 		assert_eq!(record.stage, 1);
 	}
+
+	state.cleanup_assignments_gathering_timestamp(200);
+	assert_eq!(state.per_block_assignments_gathering_times.len(), 0);
 }
 
 // Test we note the time we took to transition RequiredTranche  from Pending to Exact and
@@ -5135,7 +5157,9 @@ fn test_observe_assignment_gathering_status() {
 		clock: Box::new(MockClock::default()),
 		assignment_criteria: Box::new(MockAssignmentCriteria::check_only(|_| Ok(0))),
 		spans: HashMap::new(),
-		per_block_assignments_gathering_times: Default::default(),
+		per_block_assignments_gathering_times: LruMap::new(ByLength::new(
+			MAX_BLOCKS_WITH_ASSIGNMENT_TIMESTAMPS,
+		)),
 	};
 
 	let metrics_inner = MetricsInner {