Fix statement distribution benchmark

polkadot/node/subsystem-bench/src/lib/availability/mod.rs

@@ -22,9 +22,7 @@ use crate::{
 		av_store::{MockAvailabilityStore, NetworkAvailabilityState},
 		chain_api::{ChainApiState, MockChainApi},
 		network_bridge::{self, MockNetworkBridgeRx, MockNetworkBridgeTx},
-		runtime_api::{
-			node_features_with_chunk_mapping_enabled, MockRuntimeApi, MockRuntimeApiCoreState,
-		},
+		runtime_api::{default_node_features, MockRuntimeApi, MockRuntimeApiCoreState},
 		AlwaysSupportsParachains,
 	},
 	network::new_network,
@@ -394,7 +392,7 @@ pub async fn benchmark_availability_write(
 				expected_erasure_root: backed_candidate.descriptor().erasure_root(),
 				tx,
 				core_index: CoreIndex(core_index as u32),
-				node_features: node_features_with_chunk_mapping_enabled(),
+				node_features: default_node_features(),
 			},
 		))
 		.await;

polkadot/node/subsystem-bench/src/lib/availability/test_state.rs

@@ -17,7 +17,7 @@
 use crate::{
 	configuration::{TestAuthorities, TestConfiguration},
 	environment::GENESIS_HASH,
-	mock::runtime_api::node_features_with_chunk_mapping_enabled,
+	mock::runtime_api::default_node_features,
 };
 use bitvec::bitvec;
 use codec::Encode;
@@ -118,7 +118,7 @@ impl TestState {
 		test_state.chunk_indices = (0..config.n_cores)
 			.map(|core_index| {
 				availability_chunk_indices(
-					Some(&node_features_with_chunk_mapping_enabled()),
+					Some(&default_node_features()),
 					config.n_validators,
 					CoreIndex(core_index as u32),
 				)

polkadot/node/subsystem-bench/src/lib/mock/runtime_api.rs

@@ -28,12 +28,13 @@ use polkadot_node_subsystem_types::OverseerSignal;
 use polkadot_primitives::{
 	node_features,
 	vstaging::{CandidateEvent, CandidateReceiptV2 as CandidateReceipt, CoreState, OccupiedCore},
-	ApprovalVotingParams, AsyncBackingParams, GroupIndex, GroupRotationInfo, IndexedVec,
-	NodeFeatures, ScheduledCore, SessionIndex, SessionInfo, ValidationCode, ValidatorIndex,
+	ApprovalVotingParams, AsyncBackingParams, CoreIndex, GroupIndex, GroupRotationInfo,
+	Id as ParaId, IndexedVec, NodeFeatures, ScheduledCore, SessionIndex, SessionInfo,
+	ValidationCode, ValidatorIndex,
 };
 use sp_consensus_babe::Epoch as BabeEpoch;
 use sp_core::H256;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap, VecDeque};
 
 const LOG_TARGET: &str = "subsystem-bench::runtime-api-mock";
 
@@ -51,6 +52,8 @@ pub struct RuntimeApiState {
 	babe_epoch: Option<BabeEpoch>,
 	// The session child index,
 	session_index: SessionIndex,
+	// The claim queue
+	claim_queue: BTreeMap<CoreIndex, VecDeque<ParaId>>,
 }
 
 #[derive(Clone)]
@@ -80,7 +83,25 @@ impl MockRuntimeApi {
 		core_state: MockRuntimeApiCoreState,
 	) -> MockRuntimeApi {
 		// Enable chunk mapping feature to make systematic av-recovery possible.
-		let node_features = node_features_with_chunk_mapping_enabled();
+		let node_features = default_node_features();
+		let validator_group_count =
+			session_info_for_peers(&config, &authorities).validator_groups.len();
+
+		// Each para gets one core assigned and there is only one candidate per
+		// parachain per relay chain block (no elastic scaling).
+		let claim_queue = candidate_hashes
+			.iter()
+			.next()
+			.expect("Candidates are generated at test start")
+			.1
+			.iter()
+			.enumerate()
+			.map(|(index, candidate_receipt)| {
+				// Ensure test breaks if badly configured.
+				assert!(index < validator_group_count);
+				(CoreIndex(index as u32), vec![candidate_receipt.descriptor.para_id()].into())
+			})
+			.collect();
 
 		Self {
 			state: RuntimeApiState {
@@ -90,6 +111,7 @@ impl MockRuntimeApi {
 				babe_epoch,
 				session_index,
 				node_features,
+				claim_queue,
 			},
 			config,
 			core_state,
@@ -305,6 +327,9 @@ impl MockRuntimeApi {
 							if let Err(err) = tx.send(Ok(ApprovalVotingParams::default())) {
 								gum::error!(target: LOG_TARGET, ?err, "Voting params weren't received");
 							},
+						RuntimeApiMessage::Request(_parent, RuntimeApiRequest::ClaimQueue(tx)) => {
+							tx.send(Ok(self.state.claim_queue.clone())).unwrap();
+						},
 						// Long term TODO: implement more as needed.
 						message => {
 							unimplemented!("Unexpected runtime-api message: {:?}", message)
@@ -316,9 +341,12 @@ impl MockRuntimeApi {
 	}
 }
 
-pub fn node_features_with_chunk_mapping_enabled() -> NodeFeatures {
+pub fn default_node_features() -> NodeFeatures {
 	let mut node_features = NodeFeatures::new();
-	node_features.resize(node_features::FeatureIndex::AvailabilityChunkMapping as usize + 1, false);
+	node_features.resize(node_features::FeatureIndex::FirstUnassigned as usize, false);
 	node_features.set(node_features::FeatureIndex::AvailabilityChunkMapping as u8 as usize, true);
+	node_features.set(node_features::FeatureIndex::ElasticScalingMVP as u8 as usize, true);
+	node_features.set(node_features::FeatureIndex::CandidateReceiptV2 as u8 as usize, true);
+
 	node_features
 }

polkadot/node/subsystem-bench/src/lib/statement/test_state.rs

@@ -45,8 +45,9 @@ use polkadot_primitives::{
 		CandidateReceiptV2 as CandidateReceipt,
 		CommittedCandidateReceiptV2 as CommittedCandidateReceipt, MutateDescriptorV2,
 	},
-	BlockNumber, CandidateHash, CompactStatement, Hash, Header, Id, PersistedValidationData,
-	SessionInfo, SignedStatement, SigningContext, UncheckedSigned, ValidatorIndex, ValidatorPair,
+	BlockNumber, CandidateHash, CompactStatement, CoreIndex, Hash, Header, Id,
+	PersistedValidationData, SessionInfo, SignedStatement, SigningContext, UncheckedSigned,
+	ValidatorIndex, ValidatorPair,
 };
 use polkadot_primitives_test_helpers::{
 	dummy_committed_candidate_receipt_v2, dummy_hash, dummy_head_data, dummy_pvd,
@@ -61,6 +62,8 @@ use std::{
 	},
 };
 
+const SESSION_INDEX: u32 = 0;
+
 #[derive(Clone)]
 pub struct TestState {
 	// Full test config
@@ -130,6 +133,8 @@ impl TestState {
 				let mut receipt = receipt_templates[candidate_index].clone();
 				receipt.descriptor.set_para_id(Id::new(core_idx as u32 + 1));
 				receipt.descriptor.set_relay_parent(block_info.hash);
+				receipt.descriptor.set_core_index(CoreIndex(core_idx as u32));
+				receipt.descriptor.set_session_index(SESSION_INDEX);
 
 				state.candidate_receipts.entry(block_info.hash).or_default().push(
 					CandidateReceipt {
@@ -193,7 +198,7 @@ fn sign_statement(
 	validator_index: ValidatorIndex,
 	pair: &ValidatorPair,
 ) -> UncheckedSigned<CompactStatement> {
-	let context = SigningContext { parent_hash: relay_parent, session_index: 0 };
+	let context = SigningContext { parent_hash: relay_parent, session_index: SESSION_INDEX };
 	let payload = statement.signing_payload(&context);
 
 	SignedStatement::new(
@@ -320,7 +325,8 @@ impl HandleNetworkMessage for TestState {
 				}
 
 				let statement = CompactStatement::Valid(candidate_hash);
-				let context = SigningContext { parent_hash: relay_parent, session_index: 0 };
+				let context =
+					SigningContext { parent_hash: relay_parent, session_index: SESSION_INDEX };
 				let payload = statement.signing_payload(&context);
 				let pair = self.test_authorities.validator_pairs.get(index).unwrap();
 				let signature = pair.sign(&payload[..]);