remote-externalities: 'instant' snapshots, threading refactor, better progress logs

primitives/state-machine/src/testing.rs

@@ -27,7 +27,7 @@ use crate::{
 	StorageTransactionCache, StorageValue, TrieBackendBuilder,
 };
 
-use hash_db::Hasher;
+use hash_db::{HashDB, Hasher};
 use sp_core::{
 	offchain::testing::TestPersistentOffchainDB,
 	storage::{
@@ -160,6 +160,34 @@ where
 		self.extensions.register(ext);
 	}
 
+	/// Sets raw storage key/values and a root.
+	///
+	/// This can be used as a fast way to restore the storage state from a backup because the trie
+	/// does not need to be computed.
+	pub fn from_raw_snapshot(&mut self, raw_storage: Vec<(H::Out, Vec<u8>)>, storage_root: H::Out) {
+		for (k, v) in raw_storage {
+			self.backend.backend_storage_mut().emplace(k, hash_db::EMPTY_PREFIX, v);
+		}
+		self.backend.set_root(storage_root);
+	}
+
+	/// Drains the underlying raw storage key/values and returns the root hash.
+	///
+	/// Useful for backing up the storage in a format that can be quickly re-loaded.
+	///
+	/// Note: This DB will be inoperable after this call.
+	pub fn into_raw_snapshot(mut self) -> (Vec<(H::Out, Vec<u8>)>, H::Out) {
+		let raw_key_values = self
+			.backend
+			.backend_storage_mut()
+			.drain()
+			.into_iter()
+			.map(|(k, v)| (k, v.0))
+			.collect::<Vec<(H::Out, Vec<u8>)>>();
+
+		(raw_key_values, *self.backend.root())
+	}
+
 	/// Return a new backend with all pending changes.
 	///
 	/// In contrast to [`commit_all`](Self::commit_all) this will not panic if there are open
@@ -362,6 +390,46 @@ mod tests {
 		assert_eq!(H256::from_slice(ext.storage_root(Default::default()).as_slice()), root);
 	}
 
+	#[test]
+	fn raw_storage_drain_and_restore() {
+		// Create a TestExternalities with some data in it.
+		let mut original_ext =
+			TestExternalities::<BlakeTwo256>::from((Default::default(), Default::default()));
+		original_ext.insert(b"doe".to_vec(), b"reindeer".to_vec());
+		original_ext.insert(b"dog".to_vec(), b"puppy".to_vec());
+		original_ext.insert(b"dogglesworth".to_vec(), b"cat".to_vec());
+		let child_info = ChildInfo::new_default(&b"test_child"[..]);
+		original_ext.insert_child(child_info.clone(), b"cattytown".to_vec(), b"is_dark".to_vec());
+		original_ext.insert_child(child_info.clone(), b"doggytown".to_vec(), b"is_sunny".to_vec());
+
+		// Drain the raw storage and root.
+		let root = original_ext.backend.root().clone();
+		let (raw_storage, storage_root) = original_ext.into_raw_snapshot();
+
+		// Load the raw storage and root into a new TestExternalities.
+		let mut recovered_ext =
+			TestExternalities::<BlakeTwo256>::from((Default::default(), Default::default()));
+		recovered_ext.from_raw_snapshot(raw_storage, storage_root);
+
+		// Check the storage root is the same as the original
+		assert_eq!(root, *recovered_ext.backend.root());
+
+		// Check the original storage key/values were recovered correctly
+		assert_eq!(recovered_ext.backend.storage(b"doe").unwrap(), Some(b"reindeer".to_vec()));
+		assert_eq!(recovered_ext.backend.storage(b"dog").unwrap(), Some(b"puppy".to_vec()));
+		assert_eq!(recovered_ext.backend.storage(b"dogglesworth").unwrap(), Some(b"cat".to_vec()));
+
+		// Check the original child storage key/values were recovered correctly
+		assert_eq!(
+			recovered_ext.backend.child_storage(&child_info, b"cattytown").unwrap(),
+			Some(b"is_dark".to_vec())
+		);
+		assert_eq!(
+			recovered_ext.backend.child_storage(&child_info, b"doggytown").unwrap(),
+			Some(b"is_sunny".to_vec())
+		);
+	}
+
 	#[test]
 	fn set_and_retrieve_code() {
 		let mut ext = TestExternalities::<BlakeTwo256>::default();

primitives/state-machine/src/trie_backend.rs

@@ -242,11 +242,21 @@ where
 		&self.essence
 	}
 
+	/// Get backend storage reference.
+	pub fn backend_storage_mut(&mut self) -> &mut S {
+		self.essence.backend_storage_mut()
+	}
+
 	/// Get backend storage reference.
 	pub fn backend_storage(&self) -> &S {
 		self.essence.backend_storage()
 	}
 
+	/// Set trie root.
+	pub fn set_root(&mut self, root: H::Out) {
+		self.essence.set_root(root)
+	}
+
 	/// Get trie root.
 	pub fn root(&self) -> &H::Out {
 		self.essence.root()

utils/frame/remote-externalities/Cargo.toml

@@ -24,6 +24,8 @@ tokio = { version = "1.22.0", features = ["macros", "rt-multi-thread"] }
 substrate-rpc-client = { path = "../rpc/client" }
 futures = "0.3"
 async-recursion = "1.0.4"
+indicatif = "0.17.3"
+spinners = "4.1.0"
 
 [dev-dependencies]
 frame-support = { version = "4.0.0-dev", path = "../../../frame/support" }