feat: initialize Kurdistan SDK - independent fork of Polkadot SDK

pezkuwi/node/subsystem-bench/src/lib/availability/av_store_helpers.rs

@@ -0,0 +1,42 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+use crate::{environment::TestEnvironmentDependencies, mock::TestSyncOracle};
+use pezkuwi_node_core_av_store::{AvailabilityStoreSubsystem, Config};
+use pezkuwi_node_metrics::metrics::Metrics;
+use pezkuwi_node_subsystem_util::database::Database;
+use std::sync::Arc;
+
+mod columns {
+	pub const DATA: u32 = 0;
+	pub const META: u32 = 1;
+	pub const NUM_COLUMNS: u32 = 2;
+}
+
+const TEST_CONFIG: Config =
+	Config { col_data: columns::DATA, col_meta: columns::META, keep_finalized_for: 1 };
+
+pub fn new_av_store(dependencies: &TestEnvironmentDependencies) -> AvailabilityStoreSubsystem {
+	let metrics = Metrics::try_register(&dependencies.registry).unwrap();
+
+	AvailabilityStoreSubsystem::new(test_store(), TEST_CONFIG, Box::new(TestSyncOracle {}), metrics)
+}
+
+fn test_store() -> Arc<dyn Database> {
+	let db = kvdb_memorydb::create(columns::NUM_COLUMNS);
+	let db = pezkuwi_node_subsystem_util::database::kvdb_impl::DbAdapter::new(db, &[columns::META]);
+	Arc::new(db)
+}

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

@@ -0,0 +1,508 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+use crate::{
+	availability::av_store_helpers::new_av_store,
+	dummy_builder,
+	environment::{TestEnvironment, TestEnvironmentDependencies},
+	mock::{
+		av_store::{MockAvailabilityStore, NetworkAvailabilityState},
+		chain_api::{ChainApiState, MockChainApi},
+		network_bridge::{self, MockNetworkBridgeRx, MockNetworkBridgeTx},
+		runtime_api::{default_node_features, MockRuntimeApi, MockRuntimeApiCoreState},
+		AlwaysSupportsTeyrchains,
+	},
+	network::new_network,
+	usage::BenchmarkUsage,
+};
+use colored::Colorize;
+use futures::{channel::oneshot, stream::FuturesUnordered, StreamExt};
+
+use codec::Encode;
+use pezkuwi_availability_bitfield_distribution::BitfieldDistribution;
+use pezkuwi_availability_distribution::{
+	AvailabilityDistributionSubsystem, IncomingRequestReceivers,
+};
+use pezkuwi_availability_recovery::{AvailabilityRecoverySubsystem, RecoveryStrategyKind};
+use pezkuwi_node_core_av_store::AvailabilityStoreSubsystem;
+use pezkuwi_node_metrics::metrics::Metrics;
+use pezkuwi_node_network_protocol::{
+	request_response::{v1, v2, IncomingRequest},
+	OurView,
+};
+use pezkuwi_node_subsystem::{
+	messages::{AllMessages, AvailabilityRecoveryMessage},
+	Overseer, OverseerConnector, SpawnGlue,
+};
+use pezkuwi_node_subsystem_types::messages::{AvailabilityStoreMessage, NetworkBridgeEvent};
+use pezkuwi_overseer::{metrics::Metrics as OverseerMetrics, Handle as OverseerHandle};
+use pezkuwi_primitives::{Block, CoreIndex, GroupIndex, Hash};
+use sc_network::request_responses::{IncomingRequest as RawIncomingRequest, ProtocolConfig};
+use std::{ops::Sub, sync::Arc, time::Instant};
+use strum::Display;
+
+use sc_service::SpawnTaskHandle;
+use serde::{Deserialize, Serialize};
+pub use test_state::TestState;
+
+mod av_store_helpers;
+mod test_state;
+
+const LOG_TARGET: &str = "subsystem-bench::availability";
+
+#[derive(clap::ValueEnum, Clone, Copy, Debug, PartialEq, Serialize, Deserialize, Display)]
+#[value(rename_all = "kebab-case")]
+#[strum(serialize_all = "kebab-case")]
+pub enum Strategy {
+	/// Regular random chunk recovery. This is also the fallback for the next strategies.
+	Chunks,
+	/// Recovery from systematic chunks. Much faster than regular chunk recovery becasue it avoid
+	/// doing the reed-solomon reconstruction.
+	Systematic,
+	/// Fetch the full availability datafrom backers first. Saves CPU as we don't need to
+	/// re-construct from chunks. Typically this is only faster if nodes have enough bandwidth.
+	FullFromBackers,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, clap::Parser)]
+#[clap(rename_all = "kebab-case")]
+#[allow(missing_docs)]
+pub struct DataAvailabilityReadOptions {
+	#[clap(short, long, default_value_t = Strategy::Systematic)]
+	pub strategy: Strategy,
+}
+
+pub enum TestDataAvailability {
+	Read(DataAvailabilityReadOptions),
+	Write,
+}
+
+fn build_overseer_for_availability_read(
+	spawn_task_handle: SpawnTaskHandle,
+	runtime_api: MockRuntimeApi,
+	av_store: MockAvailabilityStore,
+	(network_bridge_tx, network_bridge_rx): (MockNetworkBridgeTx, MockNetworkBridgeRx),
+	availability_recovery: AvailabilityRecoverySubsystem,
+	dependencies: &TestEnvironmentDependencies,
+) -> (Overseer<SpawnGlue<SpawnTaskHandle>, AlwaysSupportsTeyrchains>, OverseerHandle) {
+	let overseer_connector = OverseerConnector::with_event_capacity(64000);
+	let overseer_metrics = OverseerMetrics::try_register(&dependencies.registry).unwrap();
+
+	let dummy = dummy_builder!(spawn_task_handle, overseer_metrics);
+	let builder = dummy
+		.replace_runtime_api(|_| runtime_api)
+		.replace_availability_store(|_| av_store)
+		.replace_network_bridge_tx(|_| network_bridge_tx)
+		.replace_network_bridge_rx(|_| network_bridge_rx)
+		.replace_availability_recovery(|_| availability_recovery);
+
+	let (overseer, raw_handle) =
+		builder.build_with_connector(overseer_connector).expect("Should not fail");
+
+	(overseer, OverseerHandle::new(raw_handle))
+}
+
+#[allow(clippy::too_many_arguments)]
+fn build_overseer_for_availability_write(
+	spawn_task_handle: SpawnTaskHandle,
+	runtime_api: MockRuntimeApi,
+	(network_bridge_tx, network_bridge_rx): (MockNetworkBridgeTx, MockNetworkBridgeRx),
+	availability_distribution: AvailabilityDistributionSubsystem,
+	chain_api: MockChainApi,
+	availability_store: AvailabilityStoreSubsystem,
+	bitfield_distribution: BitfieldDistribution,
+	dependencies: &TestEnvironmentDependencies,
+) -> (Overseer<SpawnGlue<SpawnTaskHandle>, AlwaysSupportsTeyrchains>, OverseerHandle) {
+	let overseer_connector = OverseerConnector::with_event_capacity(64000);
+	let overseer_metrics = OverseerMetrics::try_register(&dependencies.registry).unwrap();
+
+	let dummy = dummy_builder!(spawn_task_handle, overseer_metrics);
+	let builder = dummy
+		.replace_runtime_api(|_| runtime_api)
+		.replace_availability_store(|_| availability_store)
+		.replace_network_bridge_tx(|_| network_bridge_tx)
+		.replace_network_bridge_rx(|_| network_bridge_rx)
+		.replace_chain_api(|_| chain_api)
+		.replace_bitfield_distribution(|_| bitfield_distribution)
+		// This is needed to test own chunk recovery for `n_cores`.
+		.replace_availability_distribution(|_| availability_distribution);
+
+	let (overseer, raw_handle) =
+		builder.build_with_connector(overseer_connector).expect("Should not fail");
+
+	(overseer, OverseerHandle::new(raw_handle))
+}
+
+pub fn prepare_test(
+	state: &TestState,
+	mode: TestDataAvailability,
+	with_prometheus_endpoint: bool,
+) -> (TestEnvironment, Vec<ProtocolConfig>) {
+	let dependencies = TestEnvironmentDependencies::default();
+
+	let availability_state = NetworkAvailabilityState {
+		candidate_hashes: state.candidate_hashes.clone(),
+		candidate_hash_to_core_index: state.candidate_hash_to_core_index.clone(),
+		available_data: state.available_data.clone(),
+		chunks: state.chunks.clone(),
+		chunk_indices: state.chunk_indices.clone(),
+		req_protocol_names: state.req_protocol_names.clone(),
+	};
+
+	let mut req_cfgs = Vec::new();
+
+	let (collation_req_receiver, collation_req_cfg) = IncomingRequest::get_config_receiver::<
+		Block,
+		sc_network::NetworkWorker<Block, Hash>,
+	>(&state.req_protocol_names);
+	req_cfgs.push(collation_req_cfg);
+
+	let (pov_req_receiver, pov_req_cfg) = IncomingRequest::get_config_receiver::<
+		Block,
+		sc_network::NetworkWorker<Block, Hash>,
+	>(&state.req_protocol_names);
+	req_cfgs.push(pov_req_cfg);
+
+	let (chunk_req_v1_receiver, chunk_req_v1_cfg) =
+		IncomingRequest::<v1::ChunkFetchingRequest>::get_config_receiver::<
+			Block,
+			sc_network::NetworkWorker<Block, Hash>,
+		>(&state.req_protocol_names);
+
+	// We won't use v1 chunk fetching requests, but we need to keep the inbound queue alive.
+	// Otherwise, av-distribution subsystem will terminate.
+	std::mem::forget(chunk_req_v1_cfg);
+
+	let (chunk_req_v2_receiver, chunk_req_v2_cfg) =
+		IncomingRequest::<v2::ChunkFetchingRequest>::get_config_receiver::<
+			Block,
+			sc_network::NetworkWorker<Block, Hash>,
+		>(&state.req_protocol_names);
+
+	let (network, network_interface, network_receiver) = new_network(
+		&state.config,
+		&dependencies,
+		&state.test_authorities,
+		vec![Arc::new(availability_state.clone())],
+	);
+
+	let network_bridge_tx = network_bridge::MockNetworkBridgeTx::new(
+		network.clone(),
+		network_interface.subsystem_sender(),
+		state.test_authorities.clone(),
+	);
+	let network_bridge_rx =
+		network_bridge::MockNetworkBridgeRx::new(network_receiver, Some(chunk_req_v2_cfg));
+
+	let runtime_api = MockRuntimeApi::new(
+		state.config.clone(),
+		state.test_authorities.clone(),
+		state.candidate_receipts.clone(),
+		Default::default(),
+		Default::default(),
+		0,
+		MockRuntimeApiCoreState::Occupied,
+	);
+
+	let (overseer, overseer_handle) = match &mode {
+		TestDataAvailability::Read(options) => {
+			let subsystem = match options.strategy {
+				Strategy::FullFromBackers =>
+					AvailabilityRecoverySubsystem::with_recovery_strategy_kind(
+						collation_req_receiver,
+						&state.req_protocol_names,
+						Metrics::try_register(&dependencies.registry).unwrap(),
+						RecoveryStrategyKind::BackersFirstAlways,
+					),
+				Strategy::Chunks => AvailabilityRecoverySubsystem::with_recovery_strategy_kind(
+					collation_req_receiver,
+					&state.req_protocol_names,
+					Metrics::try_register(&dependencies.registry).unwrap(),
+					RecoveryStrategyKind::ChunksAlways,
+				),
+				Strategy::Systematic => AvailabilityRecoverySubsystem::with_recovery_strategy_kind(
+					collation_req_receiver,
+					&state.req_protocol_names,
+					Metrics::try_register(&dependencies.registry).unwrap(),
+					RecoveryStrategyKind::SystematicChunks,
+				),
+			};
+
+			// Use a mocked av-store.
+			let av_store = MockAvailabilityStore::new(
+				state.chunks.clone(),
+				state.chunk_indices.clone(),
+				state.candidate_hashes.clone(),
+				state.candidate_hash_to_core_index.clone(),
+			);
+
+			build_overseer_for_availability_read(
+				dependencies.task_manager.spawn_handle(),
+				runtime_api,
+				av_store,
+				(network_bridge_tx, network_bridge_rx),
+				subsystem,
+				&dependencies,
+			)
+		},
+		TestDataAvailability::Write => {
+			let availability_distribution = AvailabilityDistributionSubsystem::new(
+				state.test_authorities.keyring.keystore(),
+				IncomingRequestReceivers {
+					pov_req_receiver,
+					chunk_req_v1_receiver,
+					chunk_req_v2_receiver,
+				},
+				state.req_protocol_names.clone(),
+				Metrics::try_register(&dependencies.registry).unwrap(),
+			);
+
+			let chain_api_state = ChainApiState { block_headers: state.block_headers.clone() };
+			let chain_api = MockChainApi::new(chain_api_state);
+			let bitfield_distribution =
+				BitfieldDistribution::new(Metrics::try_register(&dependencies.registry).unwrap());
+			build_overseer_for_availability_write(
+				dependencies.task_manager.spawn_handle(),
+				runtime_api,
+				(network_bridge_tx, network_bridge_rx),
+				availability_distribution,
+				chain_api,
+				new_av_store(&dependencies),
+				bitfield_distribution,
+				&dependencies,
+			)
+		},
+	};
+
+	(
+		TestEnvironment::new(
+			dependencies,
+			state.config.clone(),
+			network,
+			overseer,
+			overseer_handle,
+			state.test_authorities.clone(),
+			with_prometheus_endpoint,
+		),
+		req_cfgs,
+	)
+}
+
+pub async fn benchmark_availability_read(
+	env: &mut TestEnvironment,
+	state: &TestState,
+) -> BenchmarkUsage {
+	let config = env.config().clone();
+
+	env.metrics().set_n_validators(config.n_validators);
+	env.metrics().set_n_cores(config.n_cores);
+
+	let mut batch = FuturesUnordered::new();
+	let mut availability_bytes = 0u128;
+	let mut candidates = state.candidates.clone();
+	let test_start = Instant::now();
+	for block_info in state.block_infos.iter() {
+		let block_num = block_info.number as usize;
+		gum::info!(target: LOG_TARGET, "Current block {}/{}", block_num, env.config().num_blocks);
+		env.metrics().set_current_block(block_num);
+
+		let block_start_ts = Instant::now();
+		env.import_block(block_info.clone()).await;
+
+		for candidate_num in 0..config.n_cores as u64 {
+			let candidate =
+				candidates.next().expect("We always send up to n_cores*num_blocks; qed");
+			let (tx, rx) = oneshot::channel();
+			batch.push(rx);
+
+			let message = AllMessages::AvailabilityRecovery(
+				AvailabilityRecoveryMessage::RecoverAvailableData(
+					candidate.clone(),
+					1,
+					Some(GroupIndex(
+						candidate_num as u32 % (std::cmp::max(5, config.n_cores) / 5) as u32,
+					)),
+					Some(*state.candidate_hash_to_core_index.get(&candidate.hash()).unwrap()),
+					tx,
+				),
+			);
+			env.send_message(message).await;
+		}
+
+		gum::info!(target: LOG_TARGET, "{}", format!("{} recoveries pending", batch.len()).bright_black());
+		while let Some(completed) = batch.next().await {
+			let available_data = completed.unwrap().unwrap();
+			env.metrics().on_pov_size(available_data.encoded_size());
+			availability_bytes += available_data.encoded_size() as u128;
+		}
+
+		let block_time = Instant::now().sub(block_start_ts).as_millis() as u64;
+		env.metrics().set_block_time(block_time);
+		gum::info!(target: LOG_TARGET, "All work for block completed in {}", format!("{block_time:?}ms").cyan());
+	}
+
+	let duration: u128 = test_start.elapsed().as_millis();
+	let availability_bytes = availability_bytes / 1024;
+	gum::info!(target: LOG_TARGET, "All blocks processed in {}", format!("{duration:?}ms").cyan());
+	gum::info!(target: LOG_TARGET,
+		"Throughput: {}",
+		format!("{} KiB/block", availability_bytes / env.config().num_blocks as u128).bright_red()
+	);
+	gum::info!(target: LOG_TARGET,
+		"Avg block time: {}",
+		format!("{} ms", test_start.elapsed().as_millis() / env.config().num_blocks as u128).red()
+	);
+
+	env.stop().await;
+	env.collect_resource_usage(&["availability-recovery"], false)
+}
+
+pub async fn benchmark_availability_write(
+	env: &mut TestEnvironment,
+	state: &TestState,
+) -> BenchmarkUsage {
+	let config = env.config().clone();
+
+	env.metrics().set_n_validators(config.n_validators);
+	env.metrics().set_n_cores(config.n_cores);
+
+	gum::info!(target: LOG_TARGET, "Seeding availability store with candidates ...");
+	for (core_index, backed_candidate) in state.backed_candidates.clone().into_iter().enumerate() {
+		let candidate_index = *state.candidate_hashes.get(&backed_candidate.hash()).unwrap();
+		let available_data = state.available_data[candidate_index].clone();
+		let (tx, rx) = oneshot::channel();
+		env.send_message(AllMessages::AvailabilityStore(
+			AvailabilityStoreMessage::StoreAvailableData {
+				candidate_hash: backed_candidate.hash(),
+				n_validators: config.n_validators as u32,
+				available_data,
+				expected_erasure_root: backed_candidate.descriptor().erasure_root(),
+				tx,
+				core_index: CoreIndex(core_index as u32),
+				node_features: default_node_features(),
+			},
+		))
+		.await;
+
+		rx.await
+			.unwrap()
+			.expect("Test candidates are stored nicely in availability store");
+	}
+
+	gum::info!(target: LOG_TARGET, "Done");
+
+	let test_start = Instant::now();
+	for block_info in state.block_infos.iter() {
+		let block_num = block_info.number as usize;
+		gum::info!(target: LOG_TARGET, "Current block #{}", block_num);
+		env.metrics().set_current_block(block_num);
+
+		let block_start_ts = Instant::now();
+		let relay_block_hash = block_info.hash;
+		env.import_block(block_info.clone()).await;
+
+		// Inform bitfield distribution about our view of current test block
+		let message = pezkuwi_node_subsystem_types::messages::BitfieldDistributionMessage::NetworkBridgeUpdate(
+			NetworkBridgeEvent::OurViewChange(OurView::new(vec![relay_block_hash], 0))
+		);
+		env.send_message(AllMessages::BitfieldDistribution(message)).await;
+
+		let chunk_fetch_start_ts = Instant::now();
+
+		// Request chunks of our own backed candidate from all other validators.
+		let payloads = state.chunk_fetching_requests.get(block_num - 1).expect("pregenerated");
+		let receivers = (1..config.n_validators).filter_map(|index| {
+			let (pending_response, pending_response_receiver) = oneshot::channel();
+
+			let peer_id = *env.authorities().peer_ids.get(index).expect("all validators have ids");
+			let payload = payloads.get(index).expect("pregenerated").clone();
+			let request = RawIncomingRequest { peer: peer_id, payload, pending_response };
+			let peer = env
+				.authorities()
+				.validator_authority_id
+				.get(index)
+				.expect("all validators have keys");
+
+			if env.network().is_peer_connected(peer) &&
+				env.network().send_request_from_peer(peer, request).is_ok()
+			{
+				Some(pending_response_receiver)
+			} else {
+				None
+			}
+		});
+
+		gum::info!(target: LOG_TARGET, "Waiting for all emulated peers to receive their chunk from us ...");
+
+		let responses = futures::future::try_join_all(receivers)
+			.await
+			.expect("Chunk is always served successfully");
+		// TODO: check if chunk is the one the peer expects to receive.
+		assert!(responses.iter().all(|v| v.result.is_ok()));
+
+		let chunk_fetch_duration = Instant::now().sub(chunk_fetch_start_ts).as_millis();
+		gum::info!(target: LOG_TARGET, "All chunks received in {}ms", chunk_fetch_duration);
+
+		let network = env.network().clone();
+		let authorities = env.authorities().clone();
+
+		// Spawn a task that will generate `n_validator` - 1 signed bitfields and
+		// send them from the emulated peers to the subsystem.
+		// TODO: Implement topology.
+		let messages = state.signed_bitfields.get(&relay_block_hash).expect("pregenerated").clone();
+		for index in 1..config.n_validators {
+			let from_peer = &authorities.validator_authority_id[index];
+			let message = messages.get(index).expect("pregenerated").clone();
+
+			// Send the action from peer only if it is connected to our node.
+			if network.is_peer_connected(from_peer) {
+				let _ = network.send_message_from_peer(from_peer, message);
+			}
+		}
+
+		gum::info!(
+			"Waiting for {} bitfields to be received and processed",
+			config.connected_count()
+		);
+
+		// Wait for all bitfields to be processed.
+		env.wait_until_metric(
+			"pezkuwi_teyrchain_received_availability_bitfields_total",
+			None,
+			|value| value == (config.connected_count() * block_num) as f64,
+		)
+		.await;
+
+		gum::info!(target: LOG_TARGET, "All bitfields processed");
+
+		let block_time = Instant::now().sub(block_start_ts).as_millis() as u64;
+		env.metrics().set_block_time(block_time);
+		gum::info!(target: LOG_TARGET, "All work for block completed in {}", format!("{block_time:?}ms").cyan());
+	}
+
+	let duration: u128 = test_start.elapsed().as_millis();
+	gum::info!(target: LOG_TARGET, "All blocks processed in {}", format!("{duration:?}ms").cyan());
+	gum::info!(target: LOG_TARGET,
+		"Avg block time: {}",
+		format!("{} ms", test_start.elapsed().as_millis() / env.config().num_blocks as u128).red()
+	);
+
+	env.stop().await;
+	env.collect_resource_usage(
+		&["availability-distribution", "bitfield-distribution", "availability-store"],
+		false,
+	)
+}

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

@@ -0,0 +1,306 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezkuwi.
+
+// Pezkuwi is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Pezkuwi is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Pezkuwi.  If not, see <http://www.gnu.org/licenses/>.
+
+use crate::{
+	configuration::{TestAuthorities, TestConfiguration},
+	environment::GENESIS_HASH,
+	mock::runtime_api::default_node_features,
+};
+use bitvec::bitvec;
+use codec::Encode;
+use colored::Colorize;
+use itertools::Itertools;
+use pezkuwi_node_network_protocol::{
+	request_response::{v2::ChunkFetchingRequest, ReqProtocolNames},
+	ValidationProtocols, VersionedValidationProtocol,
+};
+use pezkuwi_node_primitives::{AvailableData, BlockData, ErasureChunk, PoV};
+use pezkuwi_node_subsystem_test_helpers::{
+	derive_erasure_chunks_with_proofs_and_root, mock::new_block_import_info,
+};
+use pezkuwi_node_subsystem_util::availability_chunks::availability_chunk_indices;
+use pezkuwi_overseer::BlockInfo;
+use pezkuwi_primitives::{
+	AvailabilityBitfield, BlockNumber, CandidateHash, CandidateReceiptV2 as CandidateReceipt,
+	ChunkIndex, CoreIndex, Hash, HeadData, Header, PersistedValidationData, Signed, SigningContext,
+	ValidatorIndex,
+};
+use pezkuwi_primitives_test_helpers::{dummy_candidate_receipt_v2, dummy_hash};
+use sp_core::H256;
+use std::{collections::HashMap, iter::Cycle, sync::Arc};
+
+const LOG_TARGET: &str = "subsystem-bench::availability::test_state";
+
+#[derive(Clone)]
+pub struct TestState {
+	// Full test configuration
+	pub config: TestConfiguration,
+	// A cycle iterator on all PoV sizes used in the test.
+	pub pov_sizes: Cycle<std::vec::IntoIter<usize>>,
+	// Generated candidate receipts to be used in the test
+	pub candidates: Cycle<std::vec::IntoIter<CandidateReceipt>>,
+	// Map from pov size to candidate index
+	pub pov_size_to_candidate: HashMap<usize, usize>,
+	// Map from generated candidate hashes to candidate index in `available_data` and `chunks`.
+	pub candidate_hashes: HashMap<CandidateHash, usize>,
+	// Map from candidate hash to occupied core index.
+	pub candidate_hash_to_core_index: HashMap<CandidateHash, CoreIndex>,
+	// Per candidate index receipts.
+	pub candidate_receipt_templates: Vec<CandidateReceipt>,
+	// Per candidate index `AvailableData`
+	pub available_data: Vec<AvailableData>,
+	// Per candidate index chunks
+	pub chunks: Vec<Vec<ErasureChunk>>,
+	// Per-core ValidatorIndex -> ChunkIndex mapping
+	pub chunk_indices: Vec<Vec<ChunkIndex>>,
+	// Per relay chain block - candidate backed by our backing group
+	pub backed_candidates: Vec<CandidateReceipt>,
+	// Request protcol names
+	pub req_protocol_names: ReqProtocolNames,
+	// Relay chain block infos
+	pub block_infos: Vec<BlockInfo>,
+	// Chung fetching requests for backed candidates
+	pub chunk_fetching_requests: Vec<Vec<Vec<u8>>>,
+	// Pregenerated signed availability bitfields
+	pub signed_bitfields: HashMap<H256, Vec<VersionedValidationProtocol>>,
+	// Relay chain block headers
+	pub block_headers: HashMap<H256, Header>,
+	// Authority keys for the network emulation.
+	pub test_authorities: TestAuthorities,
+	// Map from generated candidate receipts
+	pub candidate_receipts: HashMap<H256, Vec<CandidateReceipt>>,
+}
+
+impl TestState {
+	pub fn new(config: &TestConfiguration) -> Self {
+		use pezkuwi_primitives::MutateDescriptorV2;
+		let mut test_state = Self {
+			available_data: Default::default(),
+			candidate_receipt_templates: Default::default(),
+			chunks: Default::default(),
+			pov_size_to_candidate: Default::default(),
+			pov_sizes: Vec::from(config.pov_sizes()).into_iter().cycle(),
+			candidate_hashes: HashMap::new(),
+			candidates: Vec::new().into_iter().cycle(),
+			backed_candidates: Vec::new(),
+			config: config.clone(),
+			block_infos: Default::default(),
+			chunk_fetching_requests: Default::default(),
+			signed_bitfields: Default::default(),
+			candidate_receipts: Default::default(),
+			block_headers: Default::default(),
+			test_authorities: config.generate_authorities(),
+			req_protocol_names: ReqProtocolNames::new(GENESIS_HASH, None),
+			chunk_indices: Default::default(),
+			candidate_hash_to_core_index: Default::default(),
+		};
+
+		// we use it for all candidates.
+		let persisted_validation_data = PersistedValidationData {
+			parent_head: HeadData(vec![7, 8, 9]),
+			relay_parent_number: Default::default(),
+			max_pov_size: 1024,
+			relay_parent_storage_root: Default::default(),
+		};
+
+		test_state.chunk_indices = (0..config.n_cores)
+			.map(|core_index| {
+				availability_chunk_indices(
+					&default_node_features(),
+					config.n_validators,
+					CoreIndex(core_index as u32),
+				)
+				.unwrap()
+			})
+			.collect();
+
+		// For each unique pov we create a candidate receipt.
+		for (index, pov_size) in config.pov_sizes().iter().cloned().unique().enumerate() {
+			gum::info!(target: LOG_TARGET, index, pov_size, "{}", "Generating template candidate".bright_blue());
+
+			let mut candidate_receipt = dummy_candidate_receipt_v2(dummy_hash());
+			let pov = PoV { block_data: BlockData(vec![index as u8; pov_size]) };
+
+			let new_available_data = AvailableData {
+				validation_data: persisted_validation_data.clone(),
+				pov: Arc::new(pov),
+			};
+
+			let (new_chunks, erasure_root) = derive_erasure_chunks_with_proofs_and_root(
+				config.n_validators,
+				&new_available_data,
+				|_, _| {},
+			);
+
+			candidate_receipt.descriptor.set_erasure_root(erasure_root);
+
+			test_state.chunks.push(new_chunks);
+			test_state.available_data.push(new_available_data);
+			test_state.pov_size_to_candidate.insert(pov_size, index);
+			test_state.candidate_receipt_templates.push(CandidateReceipt {
+				descriptor: candidate_receipt.descriptor,
+				commitments_hash: candidate_receipt.commitments_hash,
+			});
+		}
+
+		test_state.block_infos = (1..=config.num_blocks)
+			.map(|block_num| {
+				let relay_block_hash = Hash::repeat_byte(block_num as u8);
+				new_block_import_info(relay_block_hash, block_num as BlockNumber)
+			})
+			.collect();
+
+		test_state.block_headers = test_state
+			.block_infos
+			.iter()
+			.map(|info| {
+				(
+					info.hash,
+					Header {
+						digest: Default::default(),
+						number: info.number,
+						parent_hash: info.parent_hash,
+						extrinsics_root: Default::default(),
+						state_root: Default::default(),
+					},
+				)
+			})
+			.collect::<HashMap<_, _>>();
+
+		// Generate all candidates
+		let candidates_count = config.n_cores * config.num_blocks;
+		gum::info!(target: LOG_TARGET,"{}", format!("Pre-generating {candidates_count} candidates.").bright_blue());
+		test_state.candidates = (0..candidates_count)
+			.map(|index| {
+				let pov_size = test_state.pov_sizes.next().expect("This is a cycle; qed");
+				let candidate_index = *test_state
+					.pov_size_to_candidate
+					.get(&pov_size)
+					.expect("pov_size always exists; qed");
+				let mut candidate_receipt =
+					test_state.candidate_receipt_templates[candidate_index].clone();
+
+				// Make it unique.
+				candidate_receipt
+					.descriptor
+					.set_relay_parent(Hash::from_low_u64_be(index as u64));
+				// Store the new candidate in the state
+				test_state.candidate_hashes.insert(candidate_receipt.hash(), candidate_index);
+
+				let core_index = (index % config.n_cores) as u32;
+				test_state
+					.candidate_hash_to_core_index
+					.insert(candidate_receipt.hash(), core_index.into());
+
+				gum::debug!(target: LOG_TARGET, candidate_hash = ?candidate_receipt.hash(), "new candidate");
+
+				candidate_receipt
+			})
+			.collect::<Vec<_>>()
+			.into_iter()
+			.cycle();
+
+		// Prepare per block candidates.
+		// Genesis block is always finalized, so we start at 1.
+		for info in test_state.block_infos.iter() {
+			for _ in 0..config.n_cores {
+				let receipt = test_state.candidates.next().expect("Cycle iterator");
+				test_state.candidate_receipts.entry(info.hash).or_default().push(receipt);
+			}
+
+			// First candidate is our backed candidate.
+			test_state.backed_candidates.push(
+				test_state
+					.candidate_receipts
+					.get(&info.hash)
+					.expect("just inserted above")
+					.first()
+					.expect("just inserted above")
+					.clone(),
+			);
+		}
+
+		test_state.chunk_fetching_requests = test_state
+			.backed_candidates
+			.iter()
+			.map(|candidate| {
+				(0..config.n_validators)
+					.map(|index| {
+						ChunkFetchingRequest {
+							candidate_hash: candidate.hash(),
+							index: ValidatorIndex(index as u32),
+						}
+						.encode()
+					})
+					.collect::<Vec<_>>()
+			})
+			.collect::<Vec<_>>();
+
+		test_state.signed_bitfields = test_state
+			.block_infos
+			.iter()
+			.map(|block_info| {
+				let signing_context =
+					SigningContext { session_index: 0, parent_hash: block_info.hash };
+				let messages = (0..config.n_validators)
+					.map(|index| {
+						let validator_public = test_state
+							.test_authorities
+							.validator_public
+							.get(index)
+							.expect("All validator keys are known");
+
+						// Node has all the chunks in the world.
+						let payload: AvailabilityBitfield =
+							AvailabilityBitfield(bitvec![u8, bitvec::order::Lsb0; 1u8; 32]);
+						let signed_bitfield = Signed::<AvailabilityBitfield>::sign(
+							&test_state.test_authorities.keyring.keystore(),
+							payload,
+							&signing_context,
+							ValidatorIndex(index as u32),
+							validator_public,
+						)
+						.ok()
+						.flatten()
+						.expect("should be signed");
+
+						peer_bitfield_message_v3(block_info.hash, signed_bitfield)
+					})
+					.collect::<Vec<_>>();
+
+				(block_info.hash, messages)
+			})
+			.collect();
+
+		gum::info!(target: LOG_TARGET, "{}","Created test environment.".bright_blue());
+
+		test_state
+	}
+}
+
+fn peer_bitfield_message_v3(
+	relay_hash: H256,
+	signed_bitfield: Signed<AvailabilityBitfield>,
+) -> VersionedValidationProtocol {
+	let bitfield = pezkuwi_node_network_protocol::v3::BitfieldDistributionMessage::Bitfield(
+		relay_hash,
+		signed_bitfield.into(),
+	);
+
+	ValidationProtocols::V3(
+		pezkuwi_node_network_protocol::v3::ValidationProtocol::BitfieldDistribution(bitfield),
+	)
+}