fix(test-externalities): include memory db reference counts in snapshots (#14411)

* persist memory db reference counts in snapshots

* update proxy_test snapshot

* clippy

* comment

* comment

* add snapshot versioning

* update proxy_test

* compact snapshot version

* kick ci

* kick ci

* check snapshot version without extra struct

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]

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(),