[try-runtime-cli] Offchain worker support (#8966)

* make remote-ext work with ws and safe RPCs

* Update docs.

* Update utils/frame/remote-externalities/Cargo.toml

Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com>

* Fix test

* Update lock file

* Update utils/frame/remote-externalities/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Fix build again.

* checkpoint, merging the paged rpc now

* revert lifetime stuff

* WIP: remote client init not working

* Small cleanups

* use jsonrpsee alpha.7

* WIP

* Executiing without errors

* Reorg & cleanup

* Trivial cleaning

* Add txpool & keystore extension

* Small cleaning

* More :cleaning

* Flags: page-size, override-code

* WIP

* Apply suggestions from code review

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Remove heap_pages

* Dry code extraction from state

* Formatting

* More formatting

* Add issue todo

* Use jsonrpsee 0.2.0

* Try trigger gitlab

* Fix "block_import_works" test

* fix native_big_block_import_fails_on_fallback test

* fix commit should work

* Rewrite UI tests

* Revert "Rewrite UI tests"

This reverts commit ada7f670f701c21fb399946a3f6918453f537bcb.

* try again with UI

* Use const for legacy heap pages val

* Move parse module to its own file

* Move rpc_api module to its own file

* Apply suggestions from code review

Co-authored-by: Peter Goodspeed-Niklaus <coriolinus@users.noreply.github.com>

* trait names: Block, not B

* Corect HEAP_PAGES_TEST_LEGACY export

* Update utils/frame/remote-externalities/src/rpc_api.rs

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

* Revert test_ext heap_page insert; adjust storage root instead

* Doc comments for try_runtime::cli::Command

* TryRuntime stub

* trailing comma

* Remove unused dev dep in frame-executive

* Improve parse::hash variable name & error index

* Use Result for rpc_api fns

* Richer err messagges

* Remove HEAP_PAGE_TEST_LEGACY

* Update bin/node/executor/tests/basic.rs

Co-authored-by: kianenigma <kian@parity.io>
Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>
Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Peter Goodspeed-Niklaus <coriolinus@users.noreply.github.com>

substrate/Cargo.lock

@@ -1,5 +1,7 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+version = 3
+
 [[package]]
 name = "Inflector"
 version = "0.11.4"
@@ -6787,6 +6789,7 @@ dependencies = [
  "log",
  "pallet-elections-phragmen",
  "parity-scale-codec",
+ "serde",
  "serde_json",
  "sp-core",
  "sp-io",
@@ -10647,14 +10650,17 @@ dependencies = [
  "log",
  "parity-scale-codec",
  "remote-externalities",
+ "sc-chain-spec",
  "sc-cli",
  "sc-client-api",
  "sc-executor",
  "sc-service",
+ "serde",
  "sp-api",
  "sp-blockchain",
  "sp-core",
  "sp-externalities",
+ "sp-keystore",
  "sp-runtime",
  "sp-state-machine",
  "structopt",

substrate/bin/node/cli/src/cli.rs

@@ -47,11 +47,14 @@ pub enum Subcommand {
 	#[structopt(name = "benchmark", about = "Benchmark runtime pallets.")]
 	Benchmark(frame_benchmarking_cli::BenchmarkCmd),
 
-	/// Try some experimental command on the runtime. This includes migration and runtime-upgrade
+	/// Try some command against runtime state.
-	/// testing.
 	#[cfg(feature = "try-runtime")]
 	TryRuntime(try_runtime_cli::TryRuntimeCmd),
 
+	/// Try some command against runtime state. Note: `try-runtime` feature must be enabled.
+	#[cfg(not(feature = "try-runtime"))]
+	TryRuntime,
+
 	/// Verify a signature for a message, provided on STDIN, with a given (public or secret) key.
 	Verify(VerifyCmd),
 

substrate/bin/node/cli/src/command.rs

@@ -163,6 +163,11 @@ pub fn run() -> Result<()> {
 
 				Ok((cmd.run::<Block, Executor>(config), task_manager))
 			})
-		}
+		},
+		#[cfg(not(feature = "try-runtime"))]
+		Some(Subcommand::TryRuntime) => {
+			Err("TryRuntime wasn't enabled when building the node. \
+				You can enable it with `--features try-runtime`.".into())
+		},
 	}
 }

substrate/bin/node/executor/tests/basic.rs

@@ -697,6 +697,9 @@ fn native_big_block_import_succeeds() {
 fn native_big_block_import_fails_on_fallback() {
 	let mut t = new_test_ext(compact_code_unwrap(), false);
 
+	// We set the heap pages to 8 because we know that should give an OOM in WASM with the given block.
+	set_heap_pages(&mut t.ext(), 8);
+
 	assert!(
 		executor_call::<NeverNativeValue, fn() -> _>(
 			&mut t,

substrate/frame/executive/src/lib.rs

@@ -835,7 +835,7 @@ mod tests {
 				header: Header {
 					parent_hash: [69u8; 32].into(),
 					number: 1,
-					state_root: hex!("6e70de4fa07bac443dc7f8a812c8a0c941aacfa892bb373c5899f7d511d4c25b").into(),
+					state_root: hex!("ec6bb58b0e4bc7fdf0151a0f601eb825f529fbf90b5be5b2024deba30c5cbbcb").into(),
 					extrinsics_root: hex!("03170a2e7597b7b7e3d84c05391d139a62b157e78786d8c082f29dcf4c111314").into(),
 					digest: Digest { logs: vec![], },
 				},

substrate/primitives/state-machine/src/testing.rs

@@ -30,12 +30,12 @@ use crate::{
 	},
 };
 
-use codec::{Decode, Encode};
+use codec::Decode;
 use hash_db::Hasher;
 use sp_core::{
 	offchain::testing::TestPersistentOffchainDB,
 	storage::{
-		well_known_keys::{CHANGES_TRIE_CONFIG, CODE, HEAP_PAGES, is_child_storage_key},
+		well_known_keys::{CHANGES_TRIE_CONFIG, CODE, is_child_storage_key},
 		Storage,
 	},
 	traits::TaskExecutorExt,
@@ -103,7 +103,6 @@ where
 		assert!(storage.top.keys().all(|key| !is_child_storage_key(key)));
 		assert!(storage.children_default.keys().all(|key| is_child_storage_key(key)));
 
-		storage.top.insert(HEAP_PAGES.to_vec(), 8u64.encode());
 		storage.top.insert(CODE.to_vec(), code.to_vec());
 
 		let mut extensions = Extensions::default();
@@ -308,7 +307,7 @@ mod tests {
 		ext.set_storage(b"doe".to_vec(), b"reindeer".to_vec());
 		ext.set_storage(b"dog".to_vec(), b"puppy".to_vec());
 		ext.set_storage(b"dogglesworth".to_vec(), b"cat".to_vec());
-		let root = H256::from(hex!("2a340d3dfd52f5992c6b117e9e45f479e6da5afffafeb26ab619cf137a95aeb8"));
+		let root = H256::from(hex!("ed4d8c799d996add422395a6abd7545491d40bd838d738afafa1b8a4de625489"));
 		assert_eq!(H256::from_slice(ext.storage_root().as_slice()), root);
 	}
 

substrate/utils/frame/remote-externalities/Cargo.toml

@@ -20,8 +20,8 @@ hex = "0.4.0"
 env_logger = "0.8.2"
 log = "0.4.11"
 codec = { package = "parity-scale-codec", version = "2.0.0" }
-
 serde_json = "1.0"
+serde = "1.0.0"
 
 sp-io = { version = "3.0.0", path = "../../../primitives/io" }
 sp-core = { version = "3.0.0", path = "../../../primitives/core" }

substrate/utils/frame/remote-externalities/src/lib.rs

@@ -34,9 +34,11 @@ use sp_core::{
 use codec::{Encode, Decode};
 use sp_runtime::traits::Block as BlockT;
 use jsonrpsee_ws_client::{
-	WsClientBuilder, WsClient, v2::params::JsonRpcParams, traits::Client,
+	WsClientBuilder, WsClient, v2::params::JsonRpcParams,
 };
 
+pub mod rpc_api;
+
 type KeyPair = (StorageKey, StorageData);
 
 const LOG_TARGET: &str = "remote-ext";
@@ -72,7 +74,7 @@ impl<B: BlockT> Default for Mode<B> {
 	}
 }
 
-/// configuration of the online execution.
+/// Configuration of the offline execution.
 ///
 /// A state snapshot config must be present.
 #[derive(Clone)]
@@ -81,7 +83,7 @@ pub struct OfflineConfig {
 	pub state_snapshot: SnapshotConfig,
 }
 
-/// Description of the transport protocol.
+/// Description of the transport protocol (for online execution).
 #[derive(Debug)]
 pub struct Transport {
 	uri: String,
@@ -115,10 +117,17 @@ pub struct OnlineConfig<B: BlockT> {
 	pub transport: Transport,
 }
 
+impl<B: BlockT> OnlineConfig<B> {
+	/// Return rpc (ws) client.
+	fn rpc_client(&self) -> &WsClient {
+		self.transport.client.as_ref().expect("ws client must have been initialized by now; qed.")
+	}
+}
+
 impl<B: BlockT> Default for OnlineConfig<B> {
 	fn default() -> Self {
 		Self {
-			transport: Transport { uri: DEFAULT_TARGET.to_string(), client: None },
+			transport: Transport { uri: DEFAULT_TARGET.to_owned(), client: None },
 			at: None,
 			state_snapshot: None,
 			modules: vec![],
@@ -126,12 +135,6 @@ impl<B: BlockT> Default for OnlineConfig<B> {
 	}
 }
 
-impl<B: BlockT> OnlineConfig<B> {
-	/// Return rpc (ws) client.
-	fn rpc_client(&self) -> &WsClient {
-		self.transport.client.as_ref().expect("ws client must have been initialized by now; qed.")
-	}
-}
 
 /// Configuration of the state snapshot.
 #[derive(Clone)]
@@ -189,6 +192,7 @@ impl<B: BlockT> Builder<B> {
 
 // RPC methods
 impl<B: BlockT> Builder<B> {
+	/// Get the latest finalized head.
 	async fn rpc_get_head(&self) -> Result<B::Hash, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: finalized_head");
 		RpcApi::<B>::finalized_head(self.as_online().rpc_client()).await.map_err(|e| {
@@ -250,6 +254,7 @@ impl<B: BlockT> Builder<B> {
 		prefix: StorageKey,
 		at: B::Hash,
 	) -> Result<Vec<KeyPair>, &'static str> {
+		use jsonrpsee_ws_client::traits::Client;
 		use serde_json::to_value;
 		let keys = self.get_keys_paged(prefix, at).await?;
 		let keys_count = keys.len();
@@ -438,8 +443,10 @@ impl<B: BlockT> Builder<B> {
 		info!(target: LOG_TARGET, "injecting a total of {} keys", kv.len());
 		for (k, v) in kv {
 			let (k, v) = (k.0, v.0);
+			// Insert the key,value pair into the test trie backend
 			ext.insert(k, v);
 		}
+
 		Ok(ext)
 	}
 }

substrate/utils/frame/remote-externalities/src/rpc_api.rs

@@ -0,0 +1,53 @@
+// 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.
+
+//! WS RPC API for one off RPC calls to a substrate node.
+// TODO: Consolidate one off RPC calls https://github.com/paritytech/substrate/issues/8988
+
+use super::*;
+
+/// Get the header of the block identified by `at`
+pub async fn get_header<B: BlockT, S: AsRef<str>>(from: S, at: B::Hash) -> Result<B::Header, String>
+where
+	B::Header: serde::de::DeserializeOwned,
+{
+	use jsonrpsee_ws_client::traits::Client;
+	let at = serde_json::to_value(at)
+		.map_err(|e| format!("Block hash could not be converted to JSON due to {:?}", e))?;
+	let params = vec![at];
+	let client = WsClientBuilder::default()
+		.max_request_body_size(u32::MAX)
+		.build(from.as_ref())
+		.await
+		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))?;
+	client.request::<B::Header>("chain_getHeader", JsonRpcParams::Array(params))
+		.await
+		.map_err(|e| format!("chain_getHeader request failed due to {:?}", e))
+}
+
+/// Get the finalized head
+pub async fn get_finalized_head<B: BlockT, S: AsRef<str>>(from: S) -> Result<B::Hash, String> {
+	use jsonrpsee_ws_client::traits::Client;
+	let client = WsClientBuilder::default()
+		.max_request_body_size(u32::MAX)
+		.build(from.as_ref())
+		.await
+		.map_err(|e| format!("`WsClientBuilder` failed to build do to {:?}", e))?;
+	client.request::<B::Hash>("chain_getFinalizedHead", JsonRpcParams::NoParams)
+		.await
+		.map_err(|e| format!("chain_getFinalizedHead request failed due to {:?}", e))
+}

substrate/utils/frame/try-runtime/cli/Cargo.toml

@@ -15,18 +15,21 @@ targets = ["x86_64-unknown-linux-gnu"]
 [dependencies]
 log = "0.4.8"
 parity-scale-codec = { version = "2.0.0" }
+serde = "1.0.0"
+structopt = "0.3.8"
 
 sc-service = { version = "0.9.0", default-features = false, path = "../../../../client/service" }
 sc-cli = { version = "0.9.0", path = "../../../../client/cli" }
 sc-executor = { version = "0.9.0", path = "../../../../client/executor" }
 sc-client-api = { version = "3.0.0", path = "../../../../client/api" }
-structopt = "0.3.8"
+sc-chain-spec = { version = "3.0.0", path = "../../../../client/chain-spec" }
 sp-state-machine = { version = "0.9.0", path = "../../../../primitives/state-machine" }
 sp-api = { version = "3.0.0", path = "../../../../primitives/api" }
 sp-blockchain = { version = "3.0.0", path = "../../../../primitives/blockchain" }
 sp-runtime = { version = "3.0.0", path = "../../../../primitives/runtime" }
 sp-externalities = { version = "0.9.0", path = "../../../../primitives/externalities" }
 sp-core = { version = "3.0.0", path = "../../../../primitives/core" }
+sp-keystore = { version = "0.9.0", path = "../../../../primitives/keystore" }
 frame-try-runtime = { version = "0.9.0", path = "../../../../frame/try-runtime" }
 
 remote-externalities = { version = "0.9.0", path = "../../remote-externalities" }

substrate/utils/frame/try-runtime/cli/src/lib.rs

@@ -15,23 +15,61 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! `Structopt`-ready struct for `try-runtime`.
+//! `Structopt`-ready structs for `try-runtime`.
 
-use parity_scale_codec::Decode;
+use parity_scale_codec::{Decode, Encode};
-use std::{fmt::Debug, path::PathBuf, str::FromStr};
+use std::{fmt::Debug, path::PathBuf, str::FromStr, sync::Arc};
 use sc_service::Configuration;
 use sc_cli::{CliConfiguration, ExecutionStrategy, WasmExecutionMethod};
 use sc_executor::NativeExecutor;
 use sc_service::NativeExecutionDispatch;
+use sc_chain_spec::ChainSpec;
 use sp_state_machine::StateMachine;
 use sp_runtime::traits::{Block as BlockT, NumberFor};
-use sp_core::storage::{StorageData, StorageKey, well_known_keys};
+use sp_core::{
+	offchain::{
+		OffchainWorkerExt, OffchainDbExt, TransactionPoolExt,
+		testing::{TestOffchainExt, TestTransactionPoolExt}
+	},
+	storage::{StorageData, StorageKey, well_known_keys},
+};
+use sp_keystore::{KeystoreExt, testing::KeyStore};
+use remote_externalities::{Builder, Mode, SnapshotConfig, OfflineConfig, OnlineConfig, rpc_api};
 
-/// Various commands to try out the new runtime, over configurable states.
+mod parse;
-///
+
-/// For now this only assumes running the `on_runtime_upgrade` hooks.
+/// Possible subcommands of `try-runtime`.
-#[derive(Debug, structopt::StructOpt)]
+#[derive(Debug, Clone, structopt::StructOpt)]
-pub struct TryRuntimeCmd {
+pub enum Command {
+	/// Execute "TryRuntime_on_runtime_upgrade" against the given runtime state.
+	OnRuntimeUpgrade(OnRuntimeUpgradeCmd),
+	/// Execute "OffchainWorkerApi_offchain_worker" against the given runtime state.
+	OffchainWorker(OffchainWorkerCmd),
+}
+
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct OnRuntimeUpgradeCmd {
+	#[structopt(subcommand)]
+	pub state: State,
+}
+
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct OffchainWorkerCmd {
+	/// Hash of the block whose header to use to execute the offchain worker.
+	#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
+	pub header_at: String,
+
+	#[structopt(subcommand)]
+	pub state: State,
+
+	/// Whether or not to overwrite the code from state with the code from
+	/// the specified chain spec.
+	#[structopt(long)]
+	pub overwrite_code: bool,
+}
+
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct SharedParams {
 	/// The shared parameters
 	#[allow(missing_docs)]
 	#[structopt(flatten)]
@@ -43,7 +81,7 @@ pub struct TryRuntimeCmd {
 		value_name = "STRATEGY",
 		possible_values = &ExecutionStrategy::variants(),
 		case_insensitive = true,
-		default_value = "Native",
+		default_value = "Wasm",
 	)]
 	pub execution: ExecutionStrategy,
 
@@ -53,24 +91,38 @@ pub struct TryRuntimeCmd {
 		value_name = "METHOD",
 		possible_values = &WasmExecutionMethod::variants(),
 		case_insensitive = true,
-		default_value = "Interpreted"
+		default_value = "Compiled"
 	)]
 	pub wasm_method: WasmExecutionMethod,
 
-	/// The state to use to run the migration.
+	/// The number of 64KB pages to allocate for Wasm execution. Defaults to
-	#[structopt(subcommand)]
+	/// sc_service::Configuration.default_heap_pages.
-	pub state: State,
+	#[structopt(long)]
+	pub heap_pages: Option<u64>,
 }
 
-/// The state to use for a migration dry-run.
+/// Various commands to try out against runtime state at a specific block.
-#[derive(Debug, structopt::StructOpt)]
+#[derive(Debug, Clone, structopt::StructOpt)]
+pub struct TryRuntimeCmd {
+	#[structopt(flatten)]
+	pub shared: SharedParams,
+
+	#[structopt(subcommand)]
+	pub command: Command,
+}
+
+/// The source of runtime state to try operations against.
+#[derive(Debug, Clone, structopt::StructOpt)]
 pub enum State {
-	/// Use a state snapshot as state to run the migration.
+	/// Use a state snapshot as the source of runtime state. NOTE: for the offchain-worker command this
+	/// is only partially supported at the moment and you must have a relevant archive node exposed on
+	/// localhost:9944 in order to query the block header.
+	// TODO https://github.com/paritytech/substrate/issues/9027
 	Snap {
 		snapshot_path: PathBuf,
 	},
 
-	/// Use a live chain to run the migration.
+	/// Use a live chain as the source of runtime state.
 	Live {
 		/// An optional state snapshot file to WRITE to. Not written if set to `None`.
 		#[structopt(short, long)]
@@ -78,7 +130,7 @@ pub enum State {
 
 		/// The block hash at which to connect.
 		/// Will be latest finalized head if not provided.
-		#[structopt(short, long, multiple = false, parse(try_from_str = parse_hash))]
+		#[structopt(short, long, multiple = false, parse(try_from_str = parse::hash))]
 		block_at: Option<String>,
 
 		/// The modules to scrape. If empty, entire chain state will be scraped.
@@ -86,136 +138,243 @@ pub enum State {
 		modules: Option<Vec<String>>,
 
 		/// The url to connect to.
-		#[structopt(default_value = "ws://localhost:9944", parse(try_from_str = parse_url))]
+		#[structopt(default_value = "ws://localhost:9944", parse(try_from_str = parse::url))]
 		url: String,
 	},
 }
 
-fn parse_hash(block_number: &str) -> Result<String, String> {
+async fn on_runtime_upgrade<Block, ExecDispatch>(
-	let block_number = if block_number.starts_with("0x") {
+	shared: SharedParams,
-		&block_number[2..]
+	command: OnRuntimeUpgradeCmd,
+	config: Configuration
+) -> sc_cli::Result<()>
+where
+	Block: BlockT,
+	Block::Hash: FromStr,
+	<Block::Hash as FromStr>::Err: Debug,
+	NumberFor<Block>: FromStr,
+	<NumberFor<Block> as FromStr>::Err: Debug,
+	ExecDispatch: NativeExecutionDispatch + 'static,
+{
+	let wasm_method = shared.wasm_method;
+	let execution = shared.execution;
+	let heap_pages = if shared.heap_pages.is_some() {
+		shared.heap_pages
 	} else {
-		block_number
+		config.default_heap_pages
 	};
 
-	if let Some(pos) = block_number.chars().position(|c| !c.is_ascii_hexdigit()) {
+	let mut changes = Default::default();
-		Err(format!(
+	let max_runtime_instances = config.max_runtime_instances;
-			"Expected block hash, found illegal hex character at position: {}",
+	let executor = NativeExecutor::<ExecDispatch>::new(
-			2 + pos,
+		wasm_method.into(),
-		))
+		heap_pages,
-	} else {
+		max_runtime_instances,
-		Ok(block_number.into())
+	);
-	}
+
+	let ext = {
+		let builder = match command.state {
+			State::Snap { snapshot_path } => {
+				Builder::<Block>::new().mode(Mode::Offline(OfflineConfig {
+					state_snapshot: SnapshotConfig::new(snapshot_path),
+				}))
+			},
+			State::Live {
+				url,
+				snapshot_path,
+				block_at,
+				modules
+			} => Builder::<Block>::new().mode(Mode::Online(OnlineConfig {
+				transport: url.to_owned().into(),
+				state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
+				modules: modules.to_owned().unwrap_or_default(),
+				at: block_at.as_ref()
+					.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
+				..Default::default()
+			})),
+		};
+
+		let (code_key, code) = extract_code(config.chain_spec)?;
+		builder.inject(&[(code_key, code)]).build().await?
+	};
+
+	let encoded_result = StateMachine::<_, _, NumberFor<Block>, _>::new(
+		&ext.backend,
+		None,
+		&mut changes,
+		&executor,
+		"TryRuntime_on_runtime_upgrade",
+		&[],
+		ext.extensions,
+		&sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend)
+			.runtime_code()?,
+		sp_core::testing::TaskExecutor::new(),
+	)
+	.execute(execution.into())
+	.map_err(|e| format!("failed to execute 'TryRuntime_on_runtime_upgrade' due to {:?}", e))?;
+
+	let (weight, total_weight) = <(u64, u64) as Decode>::decode(&mut &*encoded_result)
+		.map_err(|e| format!("failed to decode output due to {:?}", e))?;
+	log::info!(
+		"TryRuntime_on_runtime_upgrade executed without errors. Consumed weight = {}, total weight = {} ({})",
+		weight,
+		total_weight,
+		weight as f64 / total_weight as f64
+	);
+
+	Ok(())
 }
 
-fn parse_url(s: &str) -> Result<String, &'static str> {
+async fn offchain_worker<Block, ExecDispatch>(
-	if s.starts_with("ws://") || s.starts_with("wss://") {
+	shared: SharedParams,
-		// could use Url crate as well, but lets keep it simple for now.
+	command: OffchainWorkerCmd,
-		Ok(s.to_string())
+	config: Configuration,
+)-> sc_cli::Result<()>
+where
+	Block: BlockT,
+	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 wasm_method = shared.wasm_method;
+	let execution = shared.execution;
+	let heap_pages = if shared.heap_pages.is_some() {
+		shared.heap_pages
 	} else {
-		Err("not a valid WS(S) url: must start with 'ws://' or 'wss://'")
+		config.default_heap_pages
-	}
+	};
-}
 
-impl TryRuntimeCmd {
+	let mut changes = Default::default();
-	pub async fn run<B, ExecDispatch>(&self, config: Configuration) -> sc_cli::Result<()>
+	let max_runtime_instances = config.max_runtime_instances;
-	where
+	let executor = NativeExecutor::<ExecDispatch>::new(
-		B: BlockT,
+		wasm_method.into(),
-		B::Hash: FromStr,
+		heap_pages,
-		<B::Hash as FromStr>::Err: Debug,
+		max_runtime_instances,
-		NumberFor<B>: FromStr,
+	);
-		<NumberFor<B> as FromStr>::Err: Debug,
-		ExecDispatch: NativeExecutionDispatch + 'static,
-	{
-		let spec = config.chain_spec;
-		let genesis_storage = spec.build_storage()?;
 
-		let code = StorageData(
+	let (mode, url) = match command.state {
-			genesis_storage
+			State::Live {
-				.top
+				url,
-				.get(well_known_keys::CODE)
+				snapshot_path,
-				.expect("code key must exist in genesis storage; qed")
+				block_at,
-				.to_vec(),
+				modules
-		);
+			} => {
-		let code_key = StorageKey(well_known_keys::CODE.to_vec());
+				let online_config = OnlineConfig {
-
-		let wasm_method = self.wasm_method;
-		let execution = self.execution;
-
-		let mut changes = Default::default();
-		// don't really care about these -- use the default values.
-		let max_runtime_instances = config.max_runtime_instances;
-		let heap_pages = config.default_heap_pages;
-		let executor = NativeExecutor::<ExecDispatch>::new(
-			wasm_method.into(),
-			heap_pages,
-			max_runtime_instances,
-		);
-
-		let ext = {
-			use remote_externalities::{Builder, Mode, SnapshotConfig, OfflineConfig, OnlineConfig};
-			let builder = match &self.state {
-				State::Snap { snapshot_path } => {
-					Builder::<B>::new().mode(Mode::Offline(OfflineConfig {
-						state_snapshot: SnapshotConfig::new(snapshot_path),
-					}))
-				},
-				State::Live {
-					url,
-					snapshot_path,
-					block_at,
-					modules
-				} => Builder::<B>::new().mode(Mode::Online(OnlineConfig {
 					transport: url.to_owned().into(),
 					state_snapshot: snapshot_path.as_ref().map(SnapshotConfig::new),
 					modules: modules.to_owned().unwrap_or_default(),
 					at: block_at.as_ref()
 						.map(|b| b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))).transpose()?,
 					..Default::default()
-				})),
+				};
-			};
 
-			// inject the code into this ext.
+				(Mode::Online(online_config), url)
-			builder.inject(&[(code_key, code)]).build().await?
+			},
-		};
+			State::Snap { snapshot_path } => {
+				// TODO This is a temporary hack; the url is used just to get the header. We should try
+				// and get the header out of state, OR use an arbitrary header if thats ok, OR allow
+				// the user to feed in a header via file.
+				// https://github.com/paritytech/substrate/issues/9027
+				// This assumes you have a node running on local host default
+				let url = "ws://127.0.0.1:9944".to_string();
+				let mode = Mode::Offline(OfflineConfig {
+					state_snapshot: SnapshotConfig::new(snapshot_path),
+				});
 
-		let encoded_result = StateMachine::<_, _, NumberFor<B>, _>::new(
+				(mode, url)
-			&ext.backend,
+			}
-			None,
+	};
-			&mut changes,
+	let builder = Builder::<Block>::new().mode(mode);
-			&executor,
+	let mut ext = if command.overwrite_code {
-			"TryRuntime_on_runtime_upgrade",
+		let (code_key, code) = extract_code(config.chain_spec)?;
-			&[],
+		builder.inject(&[(code_key, code)]).build().await?
-			ext.extensions,
+	} else {
-			&sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend)
+		builder.build().await?
-				.runtime_code()?,
+	};
-			sp_core::testing::TaskExecutor::new(),
-		)
-		.execute(execution.into())
-		.map_err(|e| format!("failed to execute 'TryRuntime_on_runtime_upgrade' due to {:?}", e))?;
 
-		let (weight, total_weight) = <(u64, u64) as Decode>::decode(&mut &*encoded_result)
+	// register externality extensions in order to provide host interface for OCW to the runtime.
-			.map_err(|e| format!("failed to decode output due to {:?}", e))?;
+	let (offchain, _offchain_state) = TestOffchainExt::new();
-		log::info!(
+	let (pool, _pool_state) = TestTransactionPoolExt::new();
-			"try-runtime executed without errors. Consumed weight = {}, total weight = {} ({})",
+	ext.register_extension(OffchainDbExt::new(offchain.clone()));
-			weight,
+	ext.register_extension(OffchainWorkerExt::new(offchain));
-			total_weight,
+	ext.register_extension(KeystoreExt(Arc::new(KeyStore::new())));
-			weight as f64 / total_weight as f64
+	ext.register_extension(TransactionPoolExt::new(pool));
-		);
 
-		Ok(())
+	let header_hash: Block::Hash = command.header_at
+		.parse()
+		.map_err(|e| format!("Could not parse header hash: {:?}", e))?;
+	let header = rpc_api::get_header::<Block, _>(url, header_hash).await?;
+
+	let _ = StateMachine::<_, _, NumberFor<Block>, _>::new(
+		&ext.backend,
+		None,
+		&mut changes,
+		&executor,
+		"OffchainWorkerApi_offchain_worker",
+		header.encode().as_ref(),
+		ext.extensions,
+		&sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend)
+			.runtime_code()?,
+		sp_core::testing::TaskExecutor::new(),
+	)
+	.execute(execution.into())
+	.map_err(|e| format!("failed to execute 'OffchainWorkerApi_offchain_worker' due to {:?}", e))?;
+
+	log::info!("OffchainWorkerApi_offchain_worker executed without errors.");
+
+	Ok(())
+}
+
+impl TryRuntimeCmd {
+	pub async fn run<Block, ExecDispatch>(&self, config: Configuration) -> sc_cli::Result<()>
+	where
+		Block: BlockT,
+		Block::Header: serde::de::DeserializeOwned,
+		Block::Hash: FromStr,
+		<Block::Hash as FromStr>::Err: Debug,
+		NumberFor<Block>: FromStr,
+		<NumberFor<Block> as FromStr>::Err: Debug,
+		ExecDispatch: NativeExecutionDispatch + 'static,
+	{
+		match &self.command {
+			Command::OnRuntimeUpgrade(ref cmd) => {
+				on_runtime_upgrade::<Block, ExecDispatch>(self.shared.clone(), cmd.clone(), config).await
+			}
+			Command::OffchainWorker(cmd) => {
+				offchain_worker::<Block, ExecDispatch>(self.shared.clone(), cmd.clone(), config).await
+			}
+		}
 	}
 }
 
 impl CliConfiguration for TryRuntimeCmd {
 	fn shared_params(&self) -> &sc_cli::SharedParams {
-		&self.shared_params
+		&self.shared.shared_params
 	}
 
 	fn chain_id(&self, _is_dev: bool) -> sc_cli::Result<String> {
-		Ok(match self.shared_params.chain {
+		Ok(match self.shared.shared_params.chain {
 			Some(ref chain) => chain.clone(),
 			None => "dev".into(),
 		})
 	}
 }
+
+/// Extract `:code` from the given chain spec and return as `StorageData` along with the 
+/// corresponding `StorageKey`.
+fn extract_code(spec: Box<dyn ChainSpec>) -> sc_cli::Result<(StorageKey, StorageData)> {
+	let genesis_storage = spec.build_storage()?;
+	let code = StorageData(
+		genesis_storage
+			.top
+			.get(well_known_keys::CODE)
+			.expect("code key must exist in genesis storage; qed")
+			.to_vec(),
+	);
+	let code_key = StorageKey(well_known_keys::CODE.to_vec());
+
+	Ok((code_key, code))
+}

substrate/utils/frame/try-runtime/cli/src/parse.rs

@@ -0,0 +1,44 @@
+// 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.
+
+//! Utils for parsing user input
+
+pub(crate) fn hash(block_hash: &str) -> Result<String, String> {
+	let (block_hash, offset) = if block_hash.starts_with("0x") {
+		(&block_hash[2..], 2)
+	} else {
+		(block_hash, 0)
+	};
+
+	if let Some(pos) = block_hash.chars().position(|c| !c.is_ascii_hexdigit()) {
+		Err(format!(
+			"Expected block hash, found illegal hex character at position: {}",
+			offset + pos,
+		))
+	} else {
+		Ok(block_hash.into())
+	}
+}
+
+pub(crate) fn url(s: &str) -> Result<String, &'static str> {
+	if s.starts_with("ws://") || s.starts_with("wss://") {
+		// could use Url crate as well, but lets keep it simple for now.
+		Ok(s.to_string())
+	} else {
+		Err("not a valid WS(S) url: must start with 'ws://' or 'wss://'")
+	}
+}