Cumulus test service cleanup (#2887)

closes #2567 

Followup for https://github.com/paritytech/polkadot-sdk/pull/2331

This PR contains multiple internal cleanups:

1. This gets rid of the functionality in `generate_genesis_block` which
was only used in one benchmark
2. Fixed `transaction_pool` and `transaction_throughput` benchmarks
failing since they require a tokio runtime now.
3. Removed `parachain_id` CLI option from the test parachain
4. Removed `expect` call from `RuntimeResolver`

cumulus/client/cli/src/lib.rs

@@ -127,6 +127,27 @@ impl sc_cli::CliConfiguration for PurgeChainCmd {
 	}
 }
 
+/// Get the SCALE encoded genesis header of the parachain.
+pub fn get_raw_genesis_header<B, C>(client: Arc<C>) -> sc_cli::Result<Vec<u8>>
+where
+	B: BlockT,
+	C: HeaderBackend<B> + 'static,
+{
+	let genesis_hash =
+		client
+			.hash(Zero::zero())?
+			.ok_or(sc_cli::Error::Client(sp_blockchain::Error::Backend(
+				"Failed to lookup genesis block hash when exporting genesis head data.".into(),
+			)))?;
+	let genesis_header = client.header(genesis_hash)?.ok_or(sc_cli::Error::Client(
+		sp_blockchain::Error::Backend(
+			"Failed to lookup genesis header by hash when exporting genesis head data.".into(),
+		),
+	))?;
+
+	Ok(genesis_header.encode())
+}
+
 /// Command for exporting the genesis head data of the parachain
 #[derive(Debug, clap::Parser)]
 pub struct ExportGenesisHeadCommand {
@@ -150,22 +171,11 @@ impl ExportGenesisHeadCommand {
 		B: BlockT,
 		C: HeaderBackend<B> + 'static,
 	{
-		let genesis_hash = client.hash(Zero::zero())?.ok_or(sc_cli::Error::Client(
-			sp_blockchain::Error::Backend(
-				"Failed to lookup genesis block hash when exporting genesis head data.".into(),
-			),
-		))?;
-		let genesis_header = client.header(genesis_hash)?.ok_or(sc_cli::Error::Client(
-			sp_blockchain::Error::Backend(
-				"Failed to lookup genesis header by hash when exporting genesis head data.".into(),
-			),
-		))?;
-
-		let raw_header = genesis_header.encode();
+		let raw_header = get_raw_genesis_header(client)?;
 		let output_buf = if self.raw {
 			raw_header
 		} else {
-			format!("0x{:?}", HexDisplay::from(&genesis_header.encode())).into_bytes()
+			format!("0x{:?}", HexDisplay::from(&raw_header)).into_bytes()
 		};
 
 		if let Some(output) = &self.output {

cumulus/polkadot-parachain/src/command.rs

@@ -60,29 +60,29 @@ enum Runtime {
 }
 
 trait RuntimeResolver {
-	fn runtime(&self) -> Runtime;
+	fn runtime(&self) -> Result<Runtime>;
 }
 
 impl RuntimeResolver for dyn ChainSpec {
-	fn runtime(&self) -> Runtime {
-		runtime(self.id())
+	fn runtime(&self) -> Result<Runtime> {
+		Ok(runtime(self.id()))
 	}
 }
 
 /// Implementation, that can resolve [`Runtime`] from any json configuration file
 impl RuntimeResolver for PathBuf {
-	fn runtime(&self) -> Runtime {
+	fn runtime(&self) -> Result<Runtime> {
 		#[derive(Debug, serde::Deserialize)]
 		struct EmptyChainSpecWithId {
 			id: String,
 		}
 
-		let file = std::fs::File::open(self).expect("Failed to open file");
+		let file = std::fs::File::open(self)?;
 		let reader = std::io::BufReader::new(file);
-		let chain_spec: EmptyChainSpecWithId = serde_json::from_reader(reader)
-			.expect("Failed to read 'json' file with ChainSpec configuration");
+		let chain_spec: EmptyChainSpecWithId =
+			serde_json::from_reader(reader).map_err(|e| sc_cli::Error::Application(Box::new(e)))?;
 
-		runtime(&chain_spec.id)
+		Ok(runtime(&chain_spec.id))
 	}
 }
 
@@ -394,7 +394,7 @@ impl SubstrateCli for RelayChainCli {
 /// Creates partial components for the runtimes that are supported by the benchmarks.
 macro_rules! construct_partials {
 	($config:expr, |$partials:ident| $code:expr) => {
-		match $config.chain_spec.runtime() {
+		match $config.chain_spec.runtime()? {
 			Runtime::AssetHubPolkadot => {
 				let $partials = new_partial::<AssetHubPolkadotRuntimeApi, _>(
 					&$config,
@@ -444,7 +444,7 @@ macro_rules! construct_partials {
 macro_rules! construct_async_run {
 	(|$components:ident, $cli:ident, $cmd:ident, $config:ident| $( $code:tt )* ) => {{
 		let runner = $cli.create_runner($cmd)?;
-		match runner.config().chain_spec.runtime() {
+		match runner.config().chain_spec.runtime()? {
 			Runtime::AssetHubPolkadot => {
 				runner.async_run(|$config| {
 					let $components = new_partial::<AssetHubPolkadotRuntimeApi, _>(
@@ -686,7 +686,7 @@ pub fn run() -> Result<()> {
 				info!("Parachain Account: {}", parachain_account);
 				info!("Is collating: {}", if config.role.is_authority() { "yes" } else { "no" });
 
-				match config.chain_spec.runtime() {
+				match config.chain_spec.runtime()? {
 					AssetHubPolkadot => crate::service::start_asset_hub_node::<
 						AssetHubPolkadotRuntimeApi,
 						AssetHubPolkadotAuraId,
@@ -1032,30 +1032,30 @@ mod tests {
 			&temp_dir,
 			Box::new(create_default_with_extensions("shell-1", Extensions1::default())),
 		);
-		assert_eq!(Runtime::Shell, path.runtime());
+		assert_eq!(Runtime::Shell, path.runtime().unwrap());
 
 		let path = store_configuration(
 			&temp_dir,
 			Box::new(create_default_with_extensions("shell-2", Extensions2::default())),
 		);
-		assert_eq!(Runtime::Shell, path.runtime());
+		assert_eq!(Runtime::Shell, path.runtime().unwrap());
 
 		let path = store_configuration(
 			&temp_dir,
 			Box::new(create_default_with_extensions("seedling", Extensions2::default())),
 		);
-		assert_eq!(Runtime::Seedling, path.runtime());
+		assert_eq!(Runtime::Seedling, path.runtime().unwrap());
 
 		let path = store_configuration(
 			&temp_dir,
 			Box::new(crate::chain_spec::rococo_parachain::rococo_parachain_local_config()),
 		);
-		assert_eq!(Runtime::Default, path.runtime());
+		assert_eq!(Runtime::Default, path.runtime().unwrap());
 
 		let path = store_configuration(
 			&temp_dir,
 			Box::new(crate::chain_spec::contracts::contracts_rococo_local_config()),
 		);
-		assert_eq!(Runtime::ContractsRococo, path.runtime());
+		assert_eq!(Runtime::ContractsRococo, path.runtime().unwrap());
 	}
 }

cumulus/test/service/benches/transaction_throughput.rs

@@ -17,6 +17,7 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion, Throughput};
+use cumulus_client_cli::get_raw_genesis_header;
 use cumulus_test_runtime::{AccountId, BalancesCall, ExistentialDeposit, SudoCall};
 use futures::{future, StreamExt};
 use sc_transaction_pool_api::{TransactionPool as _, TransactionSource, TransactionStatus};
@@ -24,9 +25,8 @@ use sp_core::{crypto::Pair, sr25519};
 use sp_runtime::OpaqueExtrinsic;
 
 use cumulus_primitives_core::ParaId;
-use cumulus_test_service::{
-	construct_extrinsic, fetch_nonce, initial_head_data, Client, Keyring::*, TransactionPool,
-};
+use cumulus_test_service::{construct_extrinsic, fetch_nonce, Client, Keyring::*, TransactionPool};
+use polkadot_primitives::HeadData;
 
 fn create_accounts(num: usize) -> Vec<sr25519::Pair> {
 	(0..num)
@@ -159,6 +159,13 @@ fn transaction_throughput_benchmarks(c: &mut Criterion) {
 		None,
 	);
 
+	// Run charlie as parachain collator
+	let charlie = runtime.block_on(
+		cumulus_test_service::TestNodeBuilder::new(para_id, tokio_handle.clone(), Charlie)
+			.enable_collator()
+			.connect_to_relay_chain_nodes(vec![&alice, &bob])
+			.build(),
+	);
 	// Register parachain
 	runtime
 		.block_on(
@@ -167,19 +174,14 @@ fn transaction_throughput_benchmarks(c: &mut Criterion) {
 				cumulus_test_service::runtime::WASM_BINARY
 					.expect("You need to build the WASM binary to run this test!")
 					.to_vec(),
-				initial_head_data(para_id),
+				HeadData(
+					get_raw_genesis_header(charlie.client.clone())
+						.expect("Unable to get genesis HeadData."),
+				),
 			),
 		)
 		.unwrap();
 
-	// Run charlie as parachain collator
-	let charlie = runtime.block_on(
-		cumulus_test_service::TestNodeBuilder::new(para_id, tokio_handle.clone(), Charlie)
-			.enable_collator()
-			.connect_to_relay_chain_nodes(vec![&alice, &bob])
-			.build(),
-	);
-
 	// Run dave as parachain collator
 	let dave = runtime.block_on(
 		cumulus_test_service::TestNodeBuilder::new(para_id, tokio_handle.clone(), Dave)

cumulus/test/service/src/cli.rs

@@ -38,9 +38,6 @@ pub struct TestCollatorCli {
 	#[command(flatten)]
 	pub run: cumulus_client_cli::RunCmd,
 
-	#[arg(default_value_t = 2000u32)]
-	pub parachain_id: u32,
-
 	/// Relay chain arguments
 	#[arg(raw = true)]
 	pub relaychain_args: Vec<String>,
@@ -256,9 +253,8 @@ impl SubstrateCli for TestCollatorCli {
 
 	fn load_spec(&self, id: &str) -> std::result::Result<Box<dyn sc_service::ChainSpec>, String> {
 		Ok(match id {
-			"" => Box::new(cumulus_test_service::get_chain_spec(Some(ParaId::from(
-				self.parachain_id,
-			)))) as Box<_>,
+			"" =>
+				Box::new(cumulus_test_service::get_chain_spec(Some(ParaId::from(2000)))) as Box<_>,
 			path => {
 				let chain_spec =
 					cumulus_test_service::chain_spec::ChainSpec::from_json_file(path.into())?;

cumulus/test/service/src/genesis.rs

@@ -1,69 +0,0 @@
-// Copyright (C) Parity Technologies (UK) Ltd.
-// This file is part of Cumulus.
-
-// Cumulus 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.
-
-// Cumulus 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 Cumulus.  If not, see <http://www.gnu.org/licenses/>.
-
-use codec::Encode;
-use cumulus_primitives_core::ParaId;
-use cumulus_test_runtime::Block;
-use polkadot_primitives::HeadData;
-use sc_chain_spec::ChainSpec;
-use sp_runtime::{
-	traits::{Block as BlockT, Hash as HashT, Header as HeaderT, Zero},
-	StateVersion,
-};
-
-/// Generate a simple test genesis block from a given ChainSpec.
-pub fn generate_genesis_block<Block: BlockT>(
-	chain_spec: &dyn ChainSpec,
-	genesis_state_version: StateVersion,
-) -> Result<Block, String> {
-	let storage = chain_spec.build_storage()?;
-
-	let child_roots = storage.children_default.iter().map(|(sk, child_content)| {
-		let state_root = <<<Block as BlockT>::Header as HeaderT>::Hashing as HashT>::trie_root(
-			child_content.data.clone().into_iter().collect(),
-			genesis_state_version,
-		);
-		(sk.clone(), state_root.encode())
-	});
-	let state_root = <<<Block as BlockT>::Header as HeaderT>::Hashing as HashT>::trie_root(
-		storage.top.clone().into_iter().chain(child_roots).collect(),
-		genesis_state_version,
-	);
-
-	let extrinsics_root = <<<Block as BlockT>::Header as HeaderT>::Hashing as HashT>::trie_root(
-		Vec::new(),
-		genesis_state_version,
-	);
-
-	Ok(Block::new(
-		<<Block as BlockT>::Header as HeaderT>::new(
-			Zero::zero(),
-			extrinsics_root,
-			state_root,
-			Default::default(),
-			Default::default(),
-		),
-		Default::default(),
-	))
-}
-
-/// Returns the initial head data for a parachain ID.
-pub fn initial_head_data(para_id: ParaId) -> HeadData {
-	let spec = crate::chain_spec::get_chain_spec(Some(para_id));
-	let block: Block = generate_genesis_block(&spec, sp_runtime::StateVersion::V1).unwrap();
-	let genesis_state = block.header().encode();
-	genesis_state.into()
-}

cumulus/test/service/src/lib.rs

@@ -23,9 +23,6 @@ pub mod bench_utils;
 
 pub mod chain_spec;
 
-/// Utilities for creating test genesis block and head data
-pub mod genesis;
-
 use runtime::AccountId;
 use sc_executor::{HeapAllocStrategy, WasmExecutor, DEFAULT_HEAP_ALLOC_STRATEGY};
 use std::{
@@ -88,7 +85,6 @@ use substrate_test_client::{
 
 pub use chain_spec::*;
 pub use cumulus_test_runtime as runtime;
-pub use genesis::*;
 pub use sp_keyring::Sr25519Keyring as Keyring;
 
 const LOG_TARGET: &str = "cumulus-test-service";
@@ -922,7 +918,7 @@ pub fn run_relay_chain_validator_node(
 ) -> polkadot_test_service::PolkadotTestNode {
 	let mut config = polkadot_test_service::node_config(
 		storage_update_func,
-		tokio_handle,
+		tokio_handle.clone(),
 		key,
 		boot_nodes,
 		true,
@@ -936,5 +932,7 @@ pub fn run_relay_chain_validator_node(
 	workers_path.pop();
 	workers_path.pop();
 
-	polkadot_test_service::run_validator_node(config, Some(workers_path))
+	tokio_handle.block_on(async move {
+		polkadot_test_service::run_validator_node(config, Some(workers_path))
+	})
 }

cumulus/test/service/src/main.rs

@@ -19,9 +19,8 @@ mod cli;
 use std::sync::Arc;
 
 use cli::{RelayChainCli, Subcommand, TestCollatorCli};
-use cumulus_primitives_core::{relay_chain::CollatorPair, ParaId};
-use cumulus_test_service::{new_partial, AnnounceBlockFn};
-use polkadot_service::runtime_traits::AccountIdConversion;
+use cumulus_primitives_core::relay_chain::CollatorPair;
+use cumulus_test_service::{chain_spec, new_partial, AnnounceBlockFn};
 use sc_cli::{CliConfiguration, SubstrateCli};
 use sp_core::Pair;
 
@@ -68,24 +67,21 @@ fn main() -> Result<(), sc_cli::Error> {
 				.create_configuration(&cli, tokio_handle.clone())
 				.expect("Should be able to generate config");
 
-			let parachain_id = ParaId::from(cli.parachain_id);
 			let polkadot_cli = RelayChainCli::new(
 				&config,
 				[RelayChainCli::executable_name()].iter().chain(cli.relaychain_args.iter()),
 			);
 
-			let parachain_account =
-				AccountIdConversion::<polkadot_primitives::AccountId>::into_account_truncating(
-					&parachain_id,
-				);
-
 			let tokio_handle = config.tokio_handle.clone();
 			let polkadot_config =
 				SubstrateCli::create_configuration(&polkadot_cli, &polkadot_cli, tokio_handle)
 					.map_err(|err| format!("Relay chain argument error: {}", err))?;
 
+			let parachain_id = chain_spec::Extensions::try_get(&*config.chain_spec)
+				.map(|e| e.para_id)
+				.ok_or("Could not find parachain extension in chain-spec.")?;
+
 			tracing::info!("Parachain id: {:?}", parachain_id);
-			tracing::info!("Parachain Account: {}", parachain_account);
 			tracing::info!(
 				"Is collating: {}",
 				if config.role.is_authority() { "yes" } else { "no" }
@@ -109,7 +105,7 @@ fn main() -> Result<(), sc_cli::Error> {
 					config,
 					collator_key,
 					polkadot_config,
-					parachain_id,
+					parachain_id.into(),
 					cli.disable_block_announcements.then(wrap_announce_block),
 					cli.fail_pov_recovery,
 					|_| Ok(jsonrpsee::RpcModule::new(())),

substrate/bin/node/cli/benches/transaction_pool.rs

@@ -55,7 +55,7 @@ fn new_node(tokio_handle: Handle) -> node_cli::service::NewFullBase {
 		impl_name: "BenchmarkImpl".into(),
 		impl_version: "1.0".into(),
 		role: Role::Authority,
-		tokio_handle,
+		tokio_handle: tokio_handle.clone(),
 		transaction_pool: TransactionPoolOptions {
 			ready: PoolLimit { count: 100_000, total_bytes: 100 * 1024 * 1024 },
 			future: PoolLimit { count: 100_000, total_bytes: 100 * 1024 * 1024 },
@@ -97,7 +97,9 @@ fn new_node(tokio_handle: Handle) -> node_cli::service::NewFullBase {
 		wasm_runtime_overrides: None,
 	};
 
-	node_cli::service::new_full_base(config, None, false, |_, _| ()).expect("Creates node")
+	tokio_handle.block_on(async move {
+		node_cli::service::new_full_base(config, None, false, |_, _| ()).expect("Creates node")
+	})
 }
 
 fn create_accounts(num: usize) -> Vec<sr25519::Pair> {