Improve remote-externalities (#8397)

* make builder generic to allow using different hash types

* expose "cache", "block_number" and "modules" as cli options for live state

* Change Builder to be generic over Block instead of Hash
add rpc method to get hash from block number
allow passing of block numbers and hashes

* fix live tests

* fix formatting in utils/frame/remote-externalities/src/lib.rs

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

* change cli to only accept block hashes
break up lines that were too long
use starts_with instead of match s.get
use unwrap_or_default instead of unwrap_or(Vec::new())

* improve error message

* fix indentation

* replace Block with sp_runtime::testing::Block

* Move cache test out of remote-test feature tests
Add cache file (contains only "Proxy" module) for local test

* simplify match expression to and_then

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

* Combine the two cfg attributes into one

Co-authored-by: David <dvdplm@gmail.com>

* Restrict visibility of test_prelude use statements to crate level

* Fix usage of and_then

* Rename cache to snapshot

* Remove fully qualified path for Debug

* Refine naming. snapshot to state_snapshot

* Remove unnecessary comment

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

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

substrate/Cargo.lock

@@ -6564,7 +6564,6 @@ dependencies = [
 name = "remote-externalities"
 version = "0.9.0"
 dependencies = [
- "async-std",
  "env_logger 0.8.3",
  "hex-literal",
  "jsonrpsee-http-client",
@@ -6574,6 +6573,8 @@ dependencies = [
  "parity-scale-codec",
  "sp-core",
  "sp-io",
+ "sp-runtime",
+ "tokio 0.2.25",
 ]
 
 [[package]]

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

@@ -25,9 +25,10 @@ codec = { package = "parity-scale-codec", version = "2.0.0" }
 
 sp-io = { version = "3.0.0", path = "../../../primitives/io" }
 sp-core = { version = "3.0.0", path = "../../../primitives/core" }
+sp-runtime = { version = "3.0.0", path = "../../../primitives/runtime" }
 
 [dev-dependencies]
-async-std = { version = "1.6.5", features = ["attributes"] }
+tokio = { version = "0.2", features = ["macros"] }
 
 [features]
 remote-test = []

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

@@ -18,7 +18,7 @@
 //! # Remote Externalities
 //!
 //! An equivalent of `sp_io::TestExternalities` that can load its state from a remote substrate
-//! based chain, or a local cache file.
+//! based chain, or a local state snapshot file.
 //!
 //! #### Runtime to Test Against
 //!
@@ -106,7 +106,7 @@ use std::{
 	path::{Path, PathBuf},
 };
 use log::*;
-use sp_core::{hashing::twox_128};
+use sp_core::hashing::twox_128;
 pub use sp_io::TestExternalities;
 use sp_core::{
 	hexdisplay::HexDisplay,
@@ -115,62 +115,62 @@ use sp_core::{
 use codec::{Encode, Decode};
 use jsonrpsee_http_client::{HttpClient, HttpConfig};
 
+use sp_runtime::traits::Block as BlockT;
+
 type KeyPair = (StorageKey, StorageData);
-type Hash = sp_core::H256;
-// TODO: make these two generic.
 
 const LOG_TARGET: &str = "remote-ext";
 const TARGET: &str = "http://localhost:9933";
 
 jsonrpsee_proc_macros::rpc_client_api! {
-	RpcApi {
+	RpcApi<B: BlockT> {
 		#[rpc(method = "state_getPairs", positional_params)]
-		fn storage_pairs(prefix: StorageKey, hash: Option<Hash>) -> Vec<(StorageKey, StorageData)>;
+		fn storage_pairs(prefix: StorageKey, hash: Option<B::Hash>) -> Vec<(StorageKey, StorageData)>;
 		#[rpc(method = "chain_getFinalizedHead")]
-		fn finalized_head() -> Hash;
+		fn finalized_head() -> B::Hash;
 	}
 }
 
 /// The execution mode.
 #[derive(Clone)]
-pub enum Mode {
+pub enum Mode<B: BlockT> {
 	/// Online.
-	Online(OnlineConfig),
-	/// Offline. Uses a cached file and needs not any client config.
+	Online(OnlineConfig<B>),
+	/// Offline. Uses a state snapshot file and needs not any client config.
 	Offline(OfflineConfig),
 }
 
 /// configuration of the online execution.
 ///
-/// A cache config must be present.
+/// A state snapshot config must be present.
 #[derive(Clone)]
 pub struct OfflineConfig {
-	/// The configuration of the cache file to use. It must be present.
-	pub cache: CacheConfig,
+	/// The configuration of the state snapshot file to use. It must be present.
+	pub state_snapshot: SnapshotConfig,
 }
 
 /// Configuration of the online execution.
 ///
-/// A cache config may be present and will be written to in that case.
+/// A state snapshot config may be present and will be written to in that case.
 #[derive(Clone)]
-pub struct OnlineConfig {
+pub struct OnlineConfig<B: BlockT> {
 	/// The HTTP uri to use.
 	pub uri: String,
 	/// The block number at which to connect. Will be latest finalized head if not provided.
-	pub at: Option<Hash>,
-	/// An optional cache file to WRITE to, not for reading. Not cached if set to `None`.
-	pub cache: Option<CacheConfig>,
+	pub at: Option<B::Hash>,
+	/// An optional state snapshot file to WRITE to, not for reading. Not written if set to `None`.
+	pub state_snapshot: Option<SnapshotConfig>,
 	/// The modules to scrape. If empty, entire chain state will be scraped.
 	pub modules: Vec<String>,
 }
 
-impl Default for OnlineConfig {
+impl<B: BlockT> Default for OnlineConfig<B> {
 	fn default() -> Self {
-		Self { uri: TARGET.to_owned(), at: None, cache: None, modules: Default::default() }
+		Self { uri: TARGET.to_owned(), at: None, state_snapshot: None, modules: Default::default() }
 	}
 }
 
-impl OnlineConfig {
+impl<B: BlockT> OnlineConfig<B> {
 	/// Return a new http rpc client.
 	fn rpc(&self) -> HttpClient {
 		HttpClient::new(&self.uri, HttpConfig { max_request_body_size: u32::MAX })
@@ -178,9 +178,9 @@ impl OnlineConfig {
 	}
 }
 
-/// Configuration of the cache.
+/// Configuration of the state snapshot.
 #[derive(Clone)]
-pub struct CacheConfig {
+pub struct SnapshotConfig {
 	// TODO: I could mix these two into one filed, but I think separate is better bc one can be
 	// configurable while one not.
 	/// File name.
@@ -189,43 +189,43 @@ pub struct CacheConfig {
 	pub directory: String,
 }
 
-impl Default for CacheConfig {
+impl Default for SnapshotConfig {
 	fn default() -> Self {
-		Self { name: "CACHE".into(), directory: ".".into() }
+		Self { name: "SNAPSHOT".into(), directory: ".".into() }
 	}
 }
 
-impl CacheConfig {
+impl SnapshotConfig {
 	fn path(&self) -> PathBuf {
 		Path::new(&self.directory).join(self.name.clone())
 	}
 }
 
 /// Builder for remote-externalities.
-pub struct Builder {
+pub struct Builder<B: BlockT> {
 	inject: Vec<KeyPair>,
-	mode: Mode,
+	mode: Mode<B>,
 }
 
-impl Default for Builder {
+impl<B: BlockT> Default for Builder<B> {
 	fn default() -> Self {
 		Self {
 			inject: Default::default(),
-			mode: Mode::Online(OnlineConfig::default())
+			mode: Mode::Online(OnlineConfig::default()),
 		}
 	}
 }
 
 // Mode methods
-impl Builder {
-	fn as_online(&self) -> &OnlineConfig {
+impl<B: BlockT> Builder<B> {
+	fn as_online(&self) -> &OnlineConfig<B> {
 		match &self.mode {
 			Mode::Online(config) => &config,
 			_ => panic!("Unexpected mode: Online"),
 		}
 	}
 
-	fn as_online_mut(&mut self) -> &mut OnlineConfig {
+	fn as_online_mut(&mut self) -> &mut OnlineConfig<B> {
 		match &mut self.mode {
 			Mode::Online(config) => config,
 			_ => panic!("Unexpected mode: Online"),
@@ -234,13 +234,13 @@ impl Builder {
 }
 
 // RPC methods
-impl Builder {
-	async fn rpc_get_head(&self) -> Result<Hash, &'static str> {
+impl<B: BlockT> Builder<B> {
+	async fn rpc_get_head(&self) -> Result<B::Hash, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: finalized_head");
-		RpcApi::finalized_head(&self.as_online().rpc()).await.map_err(|e| {
+		RpcApi::<B>::finalized_head(&self.as_online().rpc()).await.map_err(|e| {
 			error!("Error = {:?}", e);
 			"rpc finalized_head failed."
-			})
+		})
 	}
 
 	/// Relay the request to `state_getPairs` rpc endpoint.
@@ -249,28 +249,28 @@ impl Builder {
 	async fn rpc_get_pairs(
 		&self,
 		prefix: StorageKey,
-		at: Hash,
+		at: B::Hash,
 	) -> Result<Vec<KeyPair>, &'static str> {
 		trace!(target: LOG_TARGET, "rpc: storage_pairs: {:?} / {:?}", prefix, at);
-		RpcApi::storage_pairs(&self.as_online().rpc(), prefix, Some(at)).await.map_err(|e| {
+		RpcApi::<B>::storage_pairs(&self.as_online().rpc(), prefix, Some(at)).await.map_err(|e| {
 			error!("Error = {:?}", e);
 			"rpc storage_pairs failed"
-			})
+		})
 	}
 }
 
 // Internal methods
-impl Builder {
-	/// Save the given data as cache.
-	fn save_cache(&self, data: &[KeyPair], path: &Path) -> Result<(), &'static str> {
-		info!(target: LOG_TARGET, "writing to cache file {:?}", path);
+impl<B: BlockT> Builder<B> {
+	/// Save the given data as state snapshot.
+	fn save_state_snapshot(&self, data: &[KeyPair], path: &Path) -> Result<(), &'static str> {
+		info!(target: LOG_TARGET, "writing to state snapshot file {:?}", path);
 		fs::write(path, data.encode()).map_err(|_| "fs::write failed.")?;
 		Ok(())
 	}
 
-	/// initialize `Self` from cache. Panics if the file does not exist.
-	fn load_cache(&self, path: &Path) -> Result<Vec<KeyPair>, &'static str> {
-		info!(target: LOG_TARGET, "scraping keypairs from cache {:?}", path,);
+	/// initialize `Self` from state snapshot. Panics if the file does not exist.
+	fn load_state_snapshot(&self, path: &Path) -> Result<Vec<KeyPair>, &'static str> {
+		info!(target: LOG_TARGET, "scraping keypairs from state snapshot {:?}", path,);
 		let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
 		Decode::decode(&mut &*bytes).map_err(|_| "decode failed")
 	}
@@ -319,12 +319,12 @@ impl Builder {
 
 	async fn pre_build(mut self) -> Result<Vec<KeyPair>, &'static str> {
 		let mut base_kv = match self.mode.clone() {
-			Mode::Offline(config) => self.load_cache(&config.cache.path())?,
+			Mode::Offline(config) => self.load_state_snapshot(&config.state_snapshot.path())?,
 			Mode::Online(config) => {
 				self.init_remote_client().await?;
 				let kp = self.load_remote().await?;
-				if let Some(c) = config.cache {
-					self.save_cache(&kp, &c.path())?;
+				if let Some(c) = config.state_snapshot {
+					self.save_state_snapshot(&kp, &c.path())?;
 				}
 				kp
 			}
@@ -341,7 +341,7 @@ impl Builder {
 }
 
 // Public methods
-impl Builder {
+impl<B: BlockT> Builder<B> {
 	/// Create a new builder.
 	pub fn new() -> Self {
 		Default::default()
@@ -355,8 +355,8 @@ impl Builder {
 		self
 	}
 
-	/// Configure a cache to be used.
-	pub fn mode(mut self, mode: Mode) -> Self {
+	/// Configure a state snapshot to be used.
+	pub fn mode(mut self, mode: Mode<B>) -> Self {
 		self.mode = mode;
 		self
 	}
@@ -375,62 +375,75 @@ impl Builder {
 	}
 }
 
-#[cfg(feature = "remote-test")]
 #[cfg(test)]
-mod tests {
-	use super::*;
+mod test_prelude {
+	pub(crate) use super::*;
+	pub(crate) use sp_runtime::testing::{H256 as Hash, Block as RawBlock, ExtrinsicWrapper};
 
-	fn init_logger() {
+	pub(crate) type Block = RawBlock<ExtrinsicWrapper<Hash>>;
+
+	pub(crate) fn init_logger() {
 		let _ = env_logger::Builder::from_default_env()
 			.format_module_path(false)
 			.format_level(true)
 			.try_init();
 	}
+}
 
-	#[async_std::test]
+#[cfg(test)]
+mod tests {
+	use super::test_prelude::*;
+
+	#[tokio::test]
+	async fn can_load_state_snapshot() {
+		init_logger();
+		Builder::<Block>::new()
+			.mode(Mode::Offline(OfflineConfig {
+				state_snapshot: SnapshotConfig { name: "test_data/proxy_test".into(), ..Default::default() },
+			}))
+			.build()
+			.await
+			.expect("Can't read state snapshot file")
+			.execute_with(|| {});
+	}
+}
+
+#[cfg(all(test, feature = "remote-test"))]
+mod remote_tests {
+	use super::test_prelude::*;
+
+	#[tokio::test]
 	async fn can_build_one_pallet() {
 		init_logger();
-		Builder::new()
+		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
 				modules: vec!["Proxy".into()],
 				..Default::default()
 			}))
 			.build()
 			.await
-			.unwrap()
+			.expect("Can't reach the remote node. Is it running?")
 			.execute_with(|| {});
 	}
 
-	#[async_std::test]
-	async fn can_load_cache() {
+	#[tokio::test]
+	async fn can_create_state_snapshot() {
 		init_logger();
-		Builder::new()
-			.mode(Mode::Offline(OfflineConfig {
-				cache: CacheConfig { name: "proxy_test".into(), ..Default::default() },
-			}))
-			.build()
-			.await
-			.unwrap()
-			.execute_with(|| {});
-	}
-
-	#[async_std::test]
-	async fn can_create_cache() {
-		init_logger();
-		Builder::new()
+		Builder::<Block>::new()
 			.mode(Mode::Online(OnlineConfig {
-				cache: Some(CacheConfig {
-					name: "test_cache_to_remove.bin".into(),
+				state_snapshot: Some(SnapshotConfig {
+					name: "test_snapshot_to_remove.bin".into(),
 					..Default::default()
 				}),
 				..Default::default()
 			}))
 			.build()
 			.await
+			.expect("Can't reach the remote node. Is it running?")
 			.unwrap()
 			.execute_with(|| {});
 
-		let to_delete = std::fs::read_dir(CacheConfig::default().directory)
+		let to_delete = std::fs::read_dir(SnapshotConfig::default().directory)
 			.unwrap()
 			.into_iter()
 			.map(|d| d.unwrap())
@@ -444,9 +457,13 @@ mod tests {
 		}
 	}
 
-	#[async_std::test]
+	#[tokio::test]
 	async fn can_build_all() {
 		init_logger();
-		Builder::new().build().await.unwrap().execute_with(|| {});
+		Builder::<Block>::new()
+			.build()
+			.await
+			.expect("Can't reach the remote node. Is it running?")
+			.execute_with(|| {});
 	}
 }

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

@@ -18,7 +18,7 @@
 //! `Structopt`-ready struct for `try-runtime`.
 
 use parity_scale_codec::Decode;
-use std::{fmt::Debug, str::FromStr};
+use std::{fmt::Debug, path::PathBuf, str::FromStr};
 use sc_service::Configuration;
 use sc_cli::{CliConfiguration, ExecutionStrategy, WasmExecutionMethod};
 use sc_executor::NativeExecutor;
@@ -37,10 +37,6 @@ pub struct TryRuntimeCmd {
 	#[structopt(flatten)]
 	pub shared_params: sc_cli::SharedParams,
 
-	/// The state to use to run the migration. Should be a valid FILE or HTTP URI.
-	#[structopt(short, long, default_value = "http://localhost:9933")]
-	pub state: State,
-
 	/// The execution strategy that should be used for benchmarks
 	#[structopt(
 		long = "execution",
@@ -60,32 +56,90 @@ pub struct TryRuntimeCmd {
 		default_value = "Interpreted"
 	)]
 	pub wasm_method: WasmExecutionMethod,
+
+	/// The state to use to run the migration.
+	#[structopt(subcommand)]
+	pub state: State,
 }
 
 /// The state to use for a migration dry-run.
-#[derive(Debug)]
+#[derive(Debug, structopt::StructOpt)]
 pub enum State {
-	/// A snapshot. Inner value is a file path.
-	Snap(String),
+	/// Use a state snapshot as state to run the migration.
+	Snap {
+		#[structopt(flatten)]
+		snapshot_path: SnapshotPath,
+	},
 
-	/// A live chain. Inner value is the HTTP uri.
-	Live(String),
+	/// Use a live chain to run the migration.
+	Live {
+		/// An optional state snapshot file to WRITE to. Not written if set to `None`.
+		#[structopt(short, long)]
+		snapshot_path: Option<SnapshotPath>,
+
+		/// 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))]
+		block_at: Option<String>,
+
+		/// The modules to scrape. If empty, entire chain state will be scraped.
+		#[structopt(short, long, require_delimiter = true)]
+		modules: Option<Vec<String>>,
+
+		/// The url to connect to.
+		#[structopt(default_value = "http://localhost:9933", parse(try_from_str = parse_url))]
+		url: String,
+	},
 }
 
-impl FromStr for State {
+fn parse_hash(block_number: &str) -> Result<String, String> {
+	let block_number = if block_number.starts_with("0x") {
+		&block_number[2..]
+	} else {
+		block_number
+	};
+
+	if let Some(pos) = block_number.chars().position(|c| !c.is_ascii_hexdigit()) {
+		Err(format!(
+			"Expected block hash, found illegal hex character at position: {}",
+			2 + pos,
+		))
+	} else {
+		Ok(block_number.into())
+	}
+}
+
+fn parse_url(s: &str) -> Result<String, &'static str> {
+	if s.starts_with("http://") {
+		// could use Url crate as well, but lets keep it simple for now.
+		Ok(s.to_string())
+	} else {
+		Err("not a valid HTTP url: must start with 'http://'")
+	}
+}
+
+#[derive(Debug, structopt::StructOpt)]
+pub struct SnapshotPath {
+	/// The directory of the state snapshot.
+	#[structopt(short, long, default_value = ".")]
+	directory: String,
+
+	/// The file name of the state snapshot.
+	#[structopt(default_value = "SNAPSHOT")]
+	file_name: String,
+}
+
+impl FromStr for SnapshotPath {
 	type Err = &'static str;
 	fn from_str(s: &str) -> Result<Self, Self::Err> {
-		match s.get(..7) {
-			// could use Url crate as well, but lets keep it simple for now.
-			Some("http://") => Ok(State::Live(s.to_string())),
-			Some("file://") => s
-				.split("//")
-				.collect::<Vec<_>>()
-				.get(1)
-				.map(|s| State::Snap(s.to_string()))
-				.ok_or("invalid file URI"),
-			_ => Err("invalid format. Must be a valid HTTP or File URI"),
-		}
+		let p: PathBuf = s.parse().map_err(|_| "invalid path")?;
+		let parent = p.parent();
+		let file_name = p.file_name();
+
+		file_name.and_then(|file_name| Some(Self {
+			directory: parent.map(|p| p.to_string_lossy().into()).unwrap_or(".".to_string()),
+			file_name: file_name.to_string_lossy().into()
+		})).ok_or("invalid path")
 	}
 }
 
@@ -93,6 +147,10 @@ impl TryRuntimeCmd {
 	pub async fn run<B, ExecDispatch>(&self, config: Configuration) -> sc_cli::Result<()>
 	where
 		B: BlockT,
+		B::Hash: FromStr,
+		<B::Hash as FromStr>::Err: Debug,
+		NumberFor<B>: FromStr,
+		<NumberFor<B> as FromStr>::Err: Debug,
 		ExecDispatch: NativeExecutionDispatch + 'static,
 	{
 		let spec = config.chain_spec;
@@ -121,13 +179,33 @@ impl TryRuntimeCmd {
 		);
 
 		let ext = {
-			use remote_externalities::{Builder, Mode, CacheConfig, OfflineConfig, OnlineConfig};
+			use remote_externalities::{Builder, Mode, SnapshotConfig, OfflineConfig, OnlineConfig};
 			let builder = match &self.state {
-				State::Snap(file_path) => Builder::new().mode(Mode::Offline(OfflineConfig {
-					cache: CacheConfig { name: file_path.into(), ..Default::default() },
-				})),
-				State::Live(http_uri) => Builder::new().mode(Mode::Online(OnlineConfig {
-					uri: http_uri.into(),
+				State::Snap { snapshot_path } => {
+					let SnapshotPath { directory, file_name } = snapshot_path;
+					Builder::<B>::new().mode(Mode::Offline(OfflineConfig {
+						state_snapshot: SnapshotConfig {
+							name: file_name.into(),
+							directory: directory.into(),
+						},
+					}))
+				},
+				State::Live {
+					url,
+					snapshot_path,
+					block_at,
+					modules
+				} => Builder::<B>::new().mode(Mode::Online(OnlineConfig {
+					uri: url.into(),
+					state_snapshot: snapshot_path.as_ref().map(|c| SnapshotConfig {
+						name: c.file_name.clone(),
+						directory: c.directory.clone(),
+					}),
+					modules: modules.clone().unwrap_or_default(),
+					at: match block_at {
+						Some(b) => Some(b.parse().map_err(|e| format!("Could not parse hash: {:?}", e))?),
+						None => None,
+					},
 					..Default::default()
 				})),
 			};