follow-chain testing mode for try-runtime (and revamp CLI configs). (#9788)

* deadlock, need to ask someone to help now

* Finally it seems to be working.. at least for a few blocks

* self-review

* major mega revamp

* some small fixes

* another mega refactor

* add license

* Apply suggestions from code review

* hack around signature verification

* Some fixes

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Zeke Mostov <32168567+emostov@users.noreply.github.com>

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Zeke Mostov <32168567+emostov@users.noreply.github.com>

* Update utils/frame/try-runtime/cli/src/lib.rs

Co-authored-by: Zeke Mostov <32168567+emostov@users.noreply.github.com>

* final tweaks, hopefully.

* a little self-review

* Add the ext root check

Co-authored-by: Zeke Mostov <32168567+emostov@users.noreply.github.com>

substrate/utils/frame/try-runtime/cli/src/commands/execute_block.rs

@@ -0,0 +1,182 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{
+	build_executor, ensure_matching_spec, extract_code, full_extensions, hash_of, local_spec,
+	state_machine_call, SharedParams, State, LOG_TARGET,
+};
+use remote_externalities::rpc_api;
+use sc_service::{Configuration, NativeExecutionDispatch};
+use sp_core::storage::well_known_keys;
+use sp_runtime::traits::{Block as BlockT, Header as HeaderT, NumberFor};
+use std::{fmt::Debug, str::FromStr};
+
+/// Configurations of the [`Command::ExecuteBlock`].
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct ExecuteBlockCmd {
+	/// Overwrite the wasm code in state or not.
+	#[structopt(long)]
+	overwrite_wasm_code: bool,
+
+	/// If set, then the state root check is disabled by the virtue of calling into
+	/// `TryRuntime_execute_block_no_check` instead of
+	/// `Core_execute_block`.
+	#[structopt(long)]
+	no_check: bool,
+
+	/// The block hash at which to fetch the block.
+	///
+	/// If the `live` state type is being used, then this can be omitted, and is equal to whatever
+	/// the `state::at` is. Only use this (with care) when combined with a snapshot.
+	#[structopt(
+		long,
+		multiple = false,
+		parse(try_from_str = crate::parse::hash)
+	)]
+	block_at: Option<String>,
+
+	/// The ws uri from which to fetch the block.
+	///
+	/// If the `live` state type is being used, then this can be omitted, and is equal to whatever
+	/// the `state::uri` is. Only use this (with care) when combined with a snapshot.
+	#[structopt(
+		long,
+		multiple = false,
+		parse(try_from_str = crate::parse::url)
+	)]
+	block_ws_uri: Option<String>,
+
+	/// The state type to use.
+	///
+	/// For this command only, if the `live` is used, then state of the parent block is fetched.
+	///
+	/// If `block_at` is provided, then the [`State::Live::at`] is being ignored.
+	#[structopt(subcommand)]
+	state: State,
+}
+
+impl ExecuteBlockCmd {
+	fn block_at<Block: BlockT>(&self) -> sc_cli::Result<Block::Hash>
+	where
+		Block::Hash: FromStr,
+		<Block::Hash as FromStr>::Err: Debug,
+	{
+		match (&self.block_at, &self.state) {
+			(Some(block_at), State::Snap { .. }) => hash_of::<Block>(&block_at),
+			(Some(block_at), State::Live { .. }) => {
+				log::warn!(target: LOG_TARGET, "--block-at is provided while state type is live. the `Live::at` will be ignored");
+				hash_of::<Block>(&block_at)
+			},
+			(None, State::Live { at: Some(at), .. }) => hash_of::<Block>(&at),
+			_ => {
+				panic!("either `--block-at` must be provided, or state must be `live with a proper `--at``");
+			},
+		}
+	}
+
+	fn block_ws_uri<Block: BlockT>(&self) -> String
+	where
+		Block::Hash: FromStr,
+		<Block::Hash as FromStr>::Err: Debug,
+	{
+		match (&self.block_ws_uri, &self.state) {
+			(Some(block_ws_uri), State::Snap { .. }) => block_ws_uri.to_owned(),
+			(Some(block_ws_uri), State::Live { .. }) => {
+				log::error!(target: LOG_TARGET, "--block-uri is provided while state type is live, Are you sure you know what you are doing?");
+				block_ws_uri.to_owned()
+			},
+			(None, State::Live { uri, .. }) => uri.clone(),
+			(None, State::Snap { .. }) => {
+				panic!("either `--block-uri` must be provided, or state must be `live`");
+			},
+		}
+	}
+}
+
+pub(crate) async fn execute_block<Block, ExecDispatch>(
+	shared: SharedParams,
+	command: ExecuteBlockCmd,
+	config: Configuration,
+) -> sc_cli::Result<()>
+where
+	Block: BlockT + serde::de::DeserializeOwned,
+	Block::Hash: FromStr,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let executor = build_executor::<ExecDispatch>(&shared, &config);
+	let execution = shared.execution;
+
+	let block_at = command.block_at::<Block>()?;
+	let block_ws_uri = command.block_ws_uri::<Block>();
+	let block: Block = rpc_api::get_block::<Block, _>(block_ws_uri.clone(), block_at).await?;
+	let parent_hash = block.header().parent_hash();
+	log::info!(
+		target: LOG_TARGET,
+		"fetched block from {:?}, parent_hash to fetch the state {:?}",
+		block_ws_uri,
+		parent_hash
+	);
+
+	let ext = {
+		let builder = command
+			.state
+			.builder::<Block>()?
+			// make sure the state is being build with the parent hash, if it is online.
+			.overwrite_online_at(parent_hash.to_owned());
+
+		let builder = if command.overwrite_wasm_code {
+			let (code_key, code) = extract_code(&config.chain_spec)?;
+			builder.inject_key_value(&[(code_key, code)])
+		} else {
+			builder.inject_hashed_key(well_known_keys::CODE)
+		};
+
+		builder.build().await?
+	};
+
+	// A digest item gets added when the runtime is processing the block, so we need to pop
+	// the last one to be consistent with what a gossiped block would contain.
+	let (mut header, extrinsics) = block.deconstruct();
+	header.digest_mut().pop();
+	let block = Block::new(header, extrinsics);
+
+	let (expected_spec_name, expected_spec_version) =
+		local_spec::<Block, ExecDispatch>(&ext, &executor);
+	ensure_matching_spec::<Block>(
+		block_ws_uri.clone(),
+		expected_spec_name,
+		expected_spec_version,
+		shared.no_spec_name_check,
+	)
+	.await;
+
+	let _ = state_machine_call::<Block, ExecDispatch>(
+		&ext,
+		&executor,
+		execution,
+		if command.no_check { "TryRuntime_execute_block_no_check" } else { "Core_execute_block" },
+		block.encode().as_ref(),
+		full_extensions(),
+	)?;
+
+	log::info!(target: LOG_TARGET, "Core_execute_block executed without errors.");
+
+	Ok(())
+}

substrate/utils/frame/try-runtime/cli/src/commands/follow_chain.rs

@@ -0,0 +1,176 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{
+	build_executor, ensure_matching_spec, extract_code, full_extensions, local_spec, parse,
+	state_machine_call, SharedParams, LOG_TARGET,
+};
+use jsonrpsee_ws_client::{
+	types::{traits::SubscriptionClient, v2::params::JsonRpcParams, Subscription},
+	WsClientBuilder,
+};
+use parity_scale_codec::Decode;
+use remote_externalities::{rpc_api, Builder, Mode, OnlineConfig};
+use sc_executor::NativeExecutionDispatch;
+use sc_service::Configuration;
+use sp_core::H256;
+use sp_runtime::traits::{Block as BlockT, Header, NumberFor};
+use std::{fmt::Debug, str::FromStr};
+
+const SUB: &'static str = "chain_subscribeFinalizedHeads";
+const UN_SUB: &'static str = "chain_unsubscribeFinalizedHeads";
+
+/// Configurations of the [`Command::FollowChain`].
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct FollowChainCmd {
+	/// The url to connect to.
+	#[structopt(
+			short,
+			long,
+			parse(try_from_str = parse::url),
+		)]
+	uri: String,
+}
+
+pub(crate) async fn follow_chain<Block, ExecDispatch>(
+	shared: SharedParams,
+	command: FollowChainCmd,
+	config: Configuration,
+) -> sc_cli::Result<()>
+where
+	Block: BlockT<Hash = H256> + serde::de::DeserializeOwned,
+	Block::Hash: FromStr,
+	Block::Header: serde::de::DeserializeOwned,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let mut maybe_state_ext = None;
+
+	let client = WsClientBuilder::default()
+		.connection_timeout(std::time::Duration::new(20, 0))
+		.max_notifs_per_subscription(1024)
+		.max_request_body_size(u32::MAX)
+		.build(&command.uri)
+		.await
+		.unwrap();
+
+	log::info!(target: LOG_TARGET, "subscribing to {:?} / {:?}", SUB, UN_SUB);
+	let mut subscription: Subscription<Block::Header> =
+		client.subscribe(&SUB, JsonRpcParams::NoParams, &UN_SUB).await.unwrap();
+
+	let (code_key, code) = extract_code(&config.chain_spec)?;
+	let executor = build_executor::<ExecDispatch>(&shared, &config);
+	let execution = shared.execution;
+
+	loop {
+		let header = match subscription.next().await {
+			Ok(Some(header)) => header,
+			Ok(None) => {
+				log::warn!("subscription returned `None`. Probably decoding has failed.");
+				break
+			},
+			Err(why) => {
+				log::warn!("subscription returned error: {:?}.", why);
+				continue
+			},
+		};
+
+		let hash = header.hash();
+		let number = header.number();
+
+		let block = rpc_api::get_block::<Block, _>(&command.uri, hash).await.unwrap();
+
+		log::debug!(
+			target: LOG_TARGET,
+			"new block event: {:?} => {:?}, extrinsics: {}",
+			hash,
+			number,
+			block.extrinsics().len()
+		);
+
+		// create an ext at the state of this block, whatever is the first subscription event.
+		if maybe_state_ext.is_none() {
+			let builder = Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
+				transport: command.uri.clone().into(),
+				at: Some(header.parent_hash().clone()),
+				..Default::default()
+			}));
+
+			let new_ext =
+				builder.inject_key_value(&[(code_key.clone(), code.clone())]).build().await?;
+			log::info!(
+				target: LOG_TARGET,
+				"initialized state externalities at {:?}, storage root {:?}",
+				number,
+				new_ext.as_backend().root()
+			);
+
+			let (expected_spec_name, expected_spec_version) =
+				local_spec::<Block, ExecDispatch>(&new_ext, &executor);
+			ensure_matching_spec::<Block>(
+				command.uri.clone(),
+				expected_spec_name,
+				expected_spec_version,
+				shared.no_spec_name_check,
+			)
+			.await;
+
+			maybe_state_ext = Some(new_ext);
+		}
+
+		let state_ext =
+			maybe_state_ext.as_mut().expect("state_ext either existed or was just created");
+
+		let (mut changes, encoded_result) = state_machine_call::<Block, ExecDispatch>(
+			&state_ext,
+			&executor,
+			execution,
+			"TryRuntime_execute_block_no_check",
+			block.encode().as_ref(),
+			full_extensions(),
+		)?;
+
+		let consumed_weight = <u64 as Decode>::decode(&mut &*encoded_result)
+			.map_err(|e| format!("failed to decode output: {:?}", e))?;
+
+		let storage_changes = changes
+			.drain_storage_changes::<_, _, NumberFor<Block>>(
+				&state_ext.backend,
+				None,
+				Default::default(),
+				&mut Default::default(),
+			)
+			.unwrap();
+		state_ext.backend.apply_transaction(
+			storage_changes.transaction_storage_root,
+			storage_changes.transaction,
+		);
+
+		log::info!(
+			target: LOG_TARGET,
+			"executed block {}, consumed weight {}, new storage root {:?}",
+			number,
+			consumed_weight,
+			state_ext.as_backend().root(),
+		);
+	}
+
+	log::error!(target: LOG_TARGET, "ws subscription must have terminated.");
+	Ok(())
+}

substrate/utils/frame/try-runtime/cli/src/commands/mod.rs

@@ -0,0 +1,21 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub(crate) mod execute_block;
+pub(crate) mod follow_chain;
+pub(crate) mod offchain_worker;
+pub(crate) mod on_runtime_upgrade;

substrate/utils/frame/try-runtime/cli/src/commands/offchain_worker.rs

@@ -0,0 +1,165 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::{
+	build_executor, ensure_matching_spec, extract_code, full_extensions, hash_of, local_spec,
+	parse, state_machine_call, SharedParams, State, LOG_TARGET,
+};
+use parity_scale_codec::Encode;
+use remote_externalities::rpc_api;
+use sc_executor::NativeExecutionDispatch;
+use sc_service::Configuration;
+use sp_core::storage::well_known_keys;
+use sp_runtime::traits::{Block as BlockT, Header, NumberFor};
+use std::{fmt::Debug, str::FromStr};
+
+/// Configurations of the [`Command::OffchainWorker`].
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct OffchainWorkerCmd {
+	/// Overwrite the wasm code in state or not.
+	#[structopt(long)]
+	overwrite_wasm_code: bool,
+
+	/// The block hash at which to fetch the header.
+	///
+	/// If the `live` state type is being used, then this can be omitted, and is equal to whatever
+	/// the `state::at` is. Only use this (with care) when combined with a snapshot.
+	#[structopt(
+		long,
+		multiple = false,
+		parse(try_from_str = parse::hash)
+	)]
+	header_at: Option<String>,
+
+	/// The ws uri from which to fetch the header.
+	///
+	/// If the `live` state type is being used, then this can be omitted, and is equal to whatever
+	/// the `state::uri` is. Only use this (with care) when combined with a snapshot.
+	#[structopt(
+		long,
+		multiple = false,
+		parse(try_from_str = parse::url)
+	)]
+	header_ws_uri: Option<String>,
+
+	/// The state type to use.
+	#[structopt(subcommand)]
+	pub state: State,
+}
+
+impl OffchainWorkerCmd {
+	fn header_at<Block: BlockT>(&self) -> sc_cli::Result<Block::Hash>
+	where
+		Block::Hash: FromStr,
+		<Block::Hash as FromStr>::Err: Debug,
+	{
+		match (&self.header_at, &self.state) {
+			(Some(header_at), State::Snap { .. }) => hash_of::<Block>(&header_at),
+			(Some(header_at), State::Live { .. }) => {
+				log::error!(target: LOG_TARGET, "--header-at is provided while state type is live, this will most likely lead to a nonsensical result.");
+				hash_of::<Block>(&header_at)
+			},
+			(None, State::Live { at: Some(at), .. }) => hash_of::<Block>(&at),
+			_ => {
+				panic!("either `--header-at` must be provided, or state must be `live` with a proper `--at`");
+			},
+		}
+	}
+
+	fn header_ws_uri<Block: BlockT>(&self) -> String
+	where
+		Block::Hash: FromStr,
+		<Block::Hash as FromStr>::Err: Debug,
+	{
+		match (&self.header_ws_uri, &self.state) {
+			(Some(header_ws_uri), State::Snap { .. }) => header_ws_uri.to_owned(),
+			(Some(header_ws_uri), State::Live { .. }) => {
+				log::error!(target: LOG_TARGET, "--header-uri is provided while state type is live, this will most likely lead to a nonsensical result.");
+				header_ws_uri.to_owned()
+			},
+			(None, State::Live { uri, .. }) => uri.clone(),
+			(None, State::Snap { .. }) => {
+				panic!("either `--header-uri` must be provided, or state must be `live`");
+			},
+		}
+	}
+}
+
+pub(crate) async fn offchain_worker<Block, ExecDispatch>(
+	shared: SharedParams,
+	command: OffchainWorkerCmd,
+	config: Configuration,
+) -> sc_cli::Result<()>
+where
+	Block: BlockT + serde::de::DeserializeOwned,
+	Block::Hash: FromStr,
+	Block::Header: serde::de::DeserializeOwned,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let executor = build_executor(&shared, &config);
+	let execution = shared.execution;
+
+	let header_at = command.header_at::<Block>()?;
+	let header_ws_uri = command.header_ws_uri::<Block>();
+
+	let header = rpc_api::get_header::<Block, _>(header_ws_uri.clone(), header_at).await?;
+	log::info!(
+		target: LOG_TARGET,
+		"fetched header from {:?}, block number: {:?}",
+		header_ws_uri,
+		header.number()
+	);
+
+	let ext = {
+		let builder = command.state.builder::<Block>()?;
+
+		let builder = if command.overwrite_wasm_code {
+			let (code_key, code) = extract_code(&config.chain_spec)?;
+			builder.inject_key_value(&[(code_key, code)])
+		} else {
+			builder.inject_hashed_key(well_known_keys::CODE)
+		};
+
+		builder.build().await?
+	};
+
+	let (expected_spec_name, expected_spec_version) =
+		local_spec::<Block, ExecDispatch>(&ext, &executor);
+	ensure_matching_spec::<Block>(
+		header_ws_uri,
+		expected_spec_name,
+		expected_spec_version,
+		shared.no_spec_name_check,
+	)
+	.await;
+
+	let _ = state_machine_call::<Block, ExecDispatch>(
+		&ext,
+		&executor,
+		execution,
+		"OffchainWorkerApi_offchain_worker",
+		header.encode().as_ref(),
+		full_extensions(),
+	)?;
+
+	log::info!(target: LOG_TARGET, "OffchainWorkerApi_offchain_worker executed without errors.");
+
+	Ok(())
+}

substrate/utils/frame/try-runtime/cli/src/commands/on_runtime_upgrade.rs

@@ -0,0 +1,92 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: Apache-2.0
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::{fmt::Debug, str::FromStr};
+
+use parity_scale_codec::Decode;
+use sc_executor::NativeExecutionDispatch;
+use sc_service::Configuration;
+use sp_runtime::traits::{Block as BlockT, NumberFor};
+
+use crate::{
+	build_executor, ensure_matching_spec, extract_code, local_spec, state_machine_call,
+	SharedParams, State, LOG_TARGET,
+};
+
+/// Configurations of the [`Command::OnRuntimeUpgrade`].
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct OnRuntimeUpgradeCmd {
+	/// The state type to use.
+	#[structopt(subcommand)]
+	pub state: State,
+}
+
+pub(crate) async fn on_runtime_upgrade<Block, ExecDispatch>(
+	shared: SharedParams,
+	command: OnRuntimeUpgradeCmd,
+	config: Configuration,
+) -> sc_cli::Result<()>
+where
+	Block: BlockT + serde::de::DeserializeOwned,
+	Block::Hash: FromStr,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let executor = build_executor(&shared, &config);
+	let execution = shared.execution;
+
+	let ext = {
+		let builder = command.state.builder::<Block>()?;
+		let (code_key, code) = extract_code(&config.chain_spec)?;
+		builder.inject_key_value(&[(code_key, code)]).build().await?
+	};
+
+	if let Some(uri) = command.state.live_uri() {
+		let (expected_spec_name, expected_spec_version) =
+			local_spec::<Block, ExecDispatch>(&ext, &executor);
+		ensure_matching_spec::<Block>(
+			uri,
+			expected_spec_name,
+			expected_spec_version,
+			shared.no_spec_name_check,
+		)
+		.await;
+	}
+
+	let (_, encoded_result) = state_machine_call::<Block, ExecDispatch>(
+		&ext,
+		&executor,
+		execution,
+		"TryRuntime_on_runtime_upgrade",
+		&[],
+		Default::default(), // we don't really need any extensions here.
+	)?;
+
+	let (weight, total_weight) = <(u64, u64) as Decode>::decode(&mut &*encoded_result)
+		.map_err(|e| format!("failed to decode output: {:?}", e))?;
+	log::info!(
+		target: LOG_TARGET,
+		"TryRuntime_on_runtime_upgrade executed without errors. Consumed weight = {}, total weight = {} ({})",
+		weight,
+		total_weight,
+		weight as f64 / total_weight.max(1) as f64
+	);
+
+	Ok(())
+}