diff --git a/substrate/primitives/state-machine/src/testing.rs b/substrate/primitives/state-machine/src/testing.rs
index 78fec43cd7ec798c936fecddd69f058e888a62b2..eefd6b1fbde93844329b3eb80367bced3c292087 100644
--- a/substrate/primitives/state-machine/src/testing.rs
+++ b/substrate/primitives/state-machine/src/testing.rs
@@ -164,9 +164,17 @@ where
///
/// 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);
+ pub fn from_raw_snapshot(
+ &mut self,
+ raw_storage: Vec<(H::Out, (Vec<u8>, i32))>,
+ storage_root: H::Out,
+ ) {
+ for (k, (v, ref_count)) in raw_storage {
+ // Each time .emplace is called the internal MemoryDb ref count increments.
+ // Repeatedly call emplace to initialise the ref count to the correct value.
+ for _ in 0..ref_count {
+ self.backend.backend_storage_mut().emplace(k, hash_db::EMPTY_PREFIX, v.clone());
+ }
}
self.backend.set_root(storage_root);
}
@@ -176,14 +184,13 @@ where
/// 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) {
+ pub fn into_raw_snapshot(mut self) -> (Vec<(H::Out, (Vec<u8>, i32))>, 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>)>>();
+ .collect::<Vec<(H::Out, (Vec<u8>, i32))>>();
(raw_key_values, *self.backend.root())
}
@@ -402,6 +409,28 @@ mod tests {
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());
+ // Call emplace on one of the keys to increment the MemoryDb refcount, so we can check
+ // that it is intact in the recovered_ext.
+ let keys = original_ext.backend.backend_storage_mut().keys();
+ let expected_ref_count = 5;
+ let ref_count_key = keys.into_iter().next().unwrap().0;
+ for _ in 0..expected_ref_count - 1 {
+ original_ext.backend.backend_storage_mut().emplace(
+ ref_count_key,
+ hash_db::EMPTY_PREFIX,
+ // We can use anything for the 'value' because it does not affect behavior when
+ // emplacing an existing key.
+ (&[0u8; 32]).to_vec(),
+ );
+ }
+ let refcount = original_ext
+ .backend
+ .backend_storage()
+ .raw(&ref_count_key, hash_db::EMPTY_PREFIX)
+ .unwrap()
+ .1;
+ assert_eq!(refcount, expected_ref_count);
+
// Drain the raw storage and root.
let root = *original_ext.backend.root();
let (raw_storage, storage_root) = original_ext.into_raw_snapshot();
@@ -428,6 +457,15 @@ mod tests {
recovered_ext.backend.child_storage(&child_info, b"doggytown").unwrap(),
Some(b"is_sunny".to_vec())
);
+
+ // Check the refcount of the key with > 1 refcount is correct.
+ let refcount = recovered_ext
+ .backend
+ .backend_storage()
+ .raw(&ref_count_key, hash_db::EMPTY_PREFIX)
+ .unwrap()
+ .1;
+ assert_eq!(refcount, expected_ref_count);
}
#[test]
diff --git a/substrate/utils/frame/remote-externalities/src/lib.rs b/substrate/utils/frame/remote-externalities/src/lib.rs
index 42df99137b2ded632805236ff04e8ef20f7c88c6..7cd3d3750aebc72d3f720c257dbece9ce983c1ef 100644
--- a/substrate/utils/frame/remote-externalities/src/lib.rs
+++ b/substrate/utils/frame/remote-externalities/src/lib.rs
@@ -21,7 +21,7 @@
//! based chain, or a local state snapshot file.
use async_recursion::async_recursion;
-use codec::{Decode, Encode};
+use codec::{Compact, Decode, Encode};
use indicatif::{ProgressBar, ProgressStyle};
use jsonrpsee::{
core::params::ArrayParams,
@@ -54,18 +54,61 @@ use tokio_retry::{strategy::FixedInterval, Retry};
type KeyValue = (StorageKey, StorageData);
type TopKeyValues = Vec<KeyValue>;
type ChildKeyValues = Vec<(ChildInfo, Vec<KeyValue>)>;
+type SnapshotVersion = Compact<u16>;
const LOG_TARGET: &str = "remote-ext";
const DEFAULT_HTTP_ENDPOINT: &str = "https://rpc.polkadot.io:443";
+const SNAPSHOT_VERSION: SnapshotVersion = Compact(2);
+
/// The snapshot that we store on disk.
#[derive(Decode, Encode)]
struct Snapshot<B: BlockT> {
+ snapshot_version: SnapshotVersion,
state_version: StateVersion,
block_hash: B::Hash,
- raw_storage: Vec<(H256, Vec<u8>)>,
+ // <Vec<Key, (Value, MemoryDbRefCount)>>
+ raw_storage: Vec<(H256, (Vec<u8>, i32))>,
storage_root: H256,
}
+impl<B: BlockT> Snapshot<B> {
+ pub fn new(
+ state_version: StateVersion,
+ block_hash: B::Hash,
+ raw_storage: Vec<(H256, (Vec<u8>, i32))>,
+ storage_root: H256,
+ ) -> Self {
+ Self {
+ snapshot_version: SNAPSHOT_VERSION,
+ state_version,
+ block_hash,
+ raw_storage,
+ storage_root,
+ }
+ }
+
+ fn load(path: &PathBuf) -> Result<Snapshot<B>, &'static str> {
+ let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
+ // The first item in the SCALE encoded struct bytes is the snapshot version. We decode and
+ // check that first, before proceeding to decode the rest of the snapshot.
+ let maybe_version: Result<SnapshotVersion, _> = Decode::decode(&mut &*bytes);
+ match maybe_version {
+ Ok(snapshot_version) => {
+ if snapshot_version != SNAPSHOT_VERSION {
+ return Err(
+ "Unsupported snapshot version detected. Please create a new snapshot.",
+ )
+ }
+ match Decode::decode(&mut &*bytes) {
+ Ok(snapshot) => return Ok(snapshot),
+ Err(_) => Err("Decode failed"),
+ }
+ },
+ Err(_) => Err("Decode failed"),
+ }
+ }
+}
+
/// An externalities that acts exactly the same as [`sp_io::TestExternalities`] but has a few extra
/// bits and pieces to it, and can be loaded remotely.
pub struct RemoteExternalities<B: BlockT> {
@@ -908,15 +951,14 @@ where
// If we need to save a snapshot, save the raw storage and root hash to the snapshot.
if let Some(path) = self.as_online().state_snapshot.clone().map(|c| c.path) {
let (raw_storage, storage_root) = pending_ext.into_raw_snapshot();
- let snapshot = Snapshot::<B> {
+ let snapshot = Snapshot::<B>::new(
state_version,
- block_hash: self
- .as_online()
+ self.as_online()
.at
.expect("set to `Some` in `init_remote_client`; must be called before; qed"),
- raw_storage: raw_storage.clone(),
+ raw_storage.clone(),
storage_root,
- };
+ );
let encoded = snapshot.encode();
log::info!(
target: LOG_TARGET,
@@ -939,12 +981,6 @@ where
Ok(pending_ext)
}
- fn load_snapshot(&mut self, path: PathBuf) -> Result<Snapshot<B>, &'static str> {
- info!(target: LOG_TARGET, "loading data from snapshot {:?}", path);
- let bytes = fs::read(path).map_err(|_| "fs::read failed.")?;
- Decode::decode(&mut &*bytes).map_err(|_| "decode failed")
- }
-
async fn do_load_remote(&mut self) -> Result<RemoteExternalities<B>, &'static str> {
self.init_remote_client().await?;
let block_hash = self.as_online().at_expected();
@@ -958,8 +994,9 @@ where
) -> Result<RemoteExternalities<B>, &'static str> {
let mut sp = Spinner::with_timer(Spinners::Dots, "Loading snapshot...".into());
let start = Instant::now();
- let Snapshot { block_hash, state_version, raw_storage, storage_root } =
- self.load_snapshot(config.state_snapshot.path.clone())?;
+ info!(target: LOG_TARGET, "Loading snapshot from {:?}", &config.state_snapshot.path);
+ let Snapshot { snapshot_version: _, block_hash, state_version, raw_storage, storage_root } =
+ Snapshot::<B>::load(&config.state_snapshot.path)?;
let mut inner_ext = TestExternalities::new_with_code_and_state(
Default::default(),
diff --git a/substrate/utils/frame/remote-externalities/test_data/proxy_test b/substrate/utils/frame/remote-externalities/test_data/proxy_test
index f749531a8a9d767b14e1869bbfa16c390cacfa68..de8105ee152df9e6f0970243f7789d37c1461734 100644
Binary files a/substrate/utils/frame/remote-externalities/test_data/proxy_test and b/substrate/utils/frame/remote-externalities/test_data/proxy_test differ