Validator side improvements

Author: slumber

node/network/collator-protocol/src/collator_side/mod.rs

@@ -1076,8 +1076,8 @@ where
 
 			let collations = state
 				.per_relay_parent
-				.get_mut(removed)
-				.map(|per_relay_parent| std::mem::take(&mut per_relay_parent.collations))
+				.remove(removed)
+				.map(|per_relay_parent| per_relay_parent.collations)
 				.unwrap_or_default();
 			for collation in collations.into_values() {
 				state.collation_result_senders.remove(&collation.receipt.hash());
@@ -1103,7 +1103,6 @@ where
 					),
 				}
 			}
-			state.per_relay_parent.remove(removed);
 			state.span_per_relay_parent.remove(removed);
 			state.waiting_collation_fetches.remove(removed);
 		}

node/network/collator-protocol/src/validator_side/collation.rs

@@ -134,15 +134,14 @@ pub struct Collations {
 	pub fetching_from: Option<CollatorId>,
 	/// Collation that were advertised to us, but we did not yet fetch.
 	pub waiting_queue: VecDeque<(PendingCollation, CollatorId)>,
-	/// How many collations have been seconded per parachain.
-	/// Only used when async backing is enabled.
-	pub seconded_count: HashMap<ParaId, usize>,
+	/// How many collations have been seconded.
+	pub seconded_count: usize,
 }
 
 impl Collations {
 	/// Note a seconded collation for a given para.
-	pub(super) fn note_seconded(&mut self, para_id: ParaId) {
-		*self.seconded_count.entry(para_id).or_insert(0) += 1
+	pub(super) fn note_seconded(&mut self) {
+		self.seconded_count += 1
 	}
 
 	/// Returns the next collation to fetch from the `unfetched_collations`.
@@ -172,31 +171,24 @@ impl Collations {
 		match self.status {
 			// We don't need to fetch any other collation when we already have seconded one.
 			CollationStatus::Seconded => None,
-			CollationStatus::Waiting => {
-				while let Some(next) = self.waiting_queue.pop_front() {
-					let para_id = next.0.para_id;
-					if !self.is_fetch_allowed(relay_parent_mode, para_id) {
-						continue
-					}
-
-					return Some(next)
-				}
-
-				None
-			},
+			CollationStatus::Waiting =>
+				if !self.is_seconded_limit_reached(relay_parent_mode) {
+					None
+				} else {
+					self.waiting_queue.pop_front()
+				},
 			CollationStatus::WaitingOnValidation | CollationStatus::Fetching =>
 				unreachable!("We have reset the status above!"),
 		}
 	}
 
 	/// Checks the limit of seconded candidates for a given para.
-	pub(super) fn is_fetch_allowed(
+	pub(super) fn is_seconded_limit_reached(
 		&self,
 		relay_parent_mode: ProspectiveParachainsMode,
-		para_id: ParaId,
 	) -> bool {
 		let seconded_limit =
 			if relay_parent_mode.is_enabled() { MAX_CANDIDATE_DEPTH + 1 } else { 1 };
-		self.seconded_count.get(&para_id).map_or(true, |&num| num < seconded_limit)
+		self.seconded_count < seconded_limit
 	}
 }

node/network/collator-protocol/src/validator_side/mod.rs

@@ -1072,7 +1072,7 @@ async fn handle_advertisement<Sender>(
 				PendingCollation::new(relay_parent, para_id, peer_id, prospective_candidate);
 
 			let collations = &mut per_relay_parent.collations;
-			if !collations.is_fetch_allowed(relay_parent_mode, para_id) {
+			if !collations.is_seconded_limit_reached(relay_parent_mode) {
 				gum::debug!(
 					target: LOG_TARGET,
 					peer_id = ?peer_id,
@@ -1342,7 +1342,7 @@ async fn process_msg<Context>(
 
 				if let Some(state) = state.per_relay_parent.get_mut(&parent) {
 					state.collations.status = CollationStatus::Seconded;
-					state.collations.note_seconded(pending_collation.para_id);
+					state.collations.note_seconded();
 				}
 				// If async backing is enabled, make an attempt to fetch next collation.
 				dequeue_next_collation_and_fetch(ctx, state, parent, collator_id).await;