remote-externalities: store block header in snapshot

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

@@ -36,7 +36,7 @@ use sp_core::{
 	},
 };
 use sp_runtime::{
-	traits::{Block as BlockT, Hash, HashingFor},
+	traits::{Block as BlockT, HashingFor},
 	StateVersion,
 };
 use sp_state_machine::TestExternalities;
@@ -58,37 +58,37 @@ 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(3);
+const DEFAULT_HTTP_ENDPOINT: &str = "https://polkadot-try-runtime-node.parity-chains.parity.io:443";
+const SNAPSHOT_VERSION: SnapshotVersion = Compact(4);
 
 /// The snapshot that we store on disk.
 #[derive(Decode, Encode)]
-struct Snapshot<H> {
+struct Snapshot<B: BlockT> {
 	snapshot_version: SnapshotVersion,
 	state_version: StateVersion,
-	block_hash: H,
 	// <Vec<Key, (Value, MemoryDbRefCount)>>
 	raw_storage: Vec<(Vec<u8>, (Vec<u8>, i32))>,
-	storage_root: H,
+	storage_root: B::Hash,
+	header: B::Header,
 }
 
-impl<H: Decode> Snapshot<H> {
+impl<B: BlockT> Snapshot<B> {
 	pub fn new(
 		state_version: StateVersion,
-		block_hash: H,
 		raw_storage: Vec<(Vec<u8>, (Vec<u8>, i32))>,
-		storage_root: H,
+		storage_root: B::Hash,
+		header: B::Header,
 	) -> Self {
 		Self {
 			snapshot_version: SNAPSHOT_VERSION,
 			state_version,
-			block_hash,
 			raw_storage,
 			storage_root,
+			header,
 		}
 	}
 
-	fn load(path: &PathBuf) -> Result<Snapshot<H>, &'static str> {
+	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.
@@ -105,21 +105,21 @@ impl<H: Decode> Snapshot<H> {
 
 /// 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<H: Hash> {
+pub struct RemoteExternalities<B: BlockT + DeserializeOwned> {
 	/// The inner externalities.
-	pub inner_ext: TestExternalities<H>,
-	/// The block hash with which we created this externality env.
-	pub block_hash: H::Out,
+	pub inner_ext: TestExternalities<HashingFor<B>>,
+	/// The block header which we created this externality env.
+	pub header: B::Header,
 }
 
-impl<H: Hash> Deref for RemoteExternalities<H> {
-	type Target = TestExternalities<H>;
+impl<B: BlockT + DeserializeOwned> Deref for RemoteExternalities<B> {
+	type Target = TestExternalities<HashingFor<B>>;
 	fn deref(&self) -> &Self::Target {
 		&self.inner_ext
 	}
 }
 
-impl<H: Hash> DerefMut for RemoteExternalities<H> {
+impl<B: BlockT + DeserializeOwned> DerefMut for RemoteExternalities<B> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
 		&mut self.inner_ext
 	}
@@ -1030,6 +1030,21 @@ where
 		Ok(())
 	}
 
+	async fn load_header(&self) -> Result<B::Header, &'static str> {
+		let retry_strategy =
+			FixedInterval::new(Self::KEYS_PAGE_RETRY_INTERVAL).take(Self::MAX_RETRIES);
+		let get_header_closure = || {
+			ChainApi::<(), _, B::Header, ()>::header(
+				self.as_online().rpc_client(),
+				Some(self.as_online().at_expected()),
+			)
+		};
+		Ok(Retry::spawn(retry_strategy, get_header_closure)
+			.await
+			.map_err(|_| "Failed to fetch header for block from network")?
+			.ok_or("Network returned None block header")?)
+	}
+
 	/// Load the data from a remote server. The main code path is calling into `load_top_remote` and
 	/// `load_child_remote`.
 	///
@@ -1058,13 +1073,11 @@ 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::Hash>::new(
+			let snapshot = Snapshot::<B>::new(
 				state_version,
-				self.as_online()
-					.at
-					.expect("set to `Some` in `init_remote_client`; must be called before; qed"),
 				raw_storage.clone(),
 				storage_root,
+				self.load_header().await?,
 			);
 			let encoded = snapshot.encode();
 			log::info!(
@@ -1086,22 +1099,21 @@ where
 		Ok(pending_ext)
 	}
 
-	async fn do_load_remote(&mut self) -> Result<RemoteExternalities<HashingFor<B>>, &'static str> {
+	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();
 		let inner_ext = self.load_remote_and_maybe_save().await?;
-		Ok(RemoteExternalities { block_hash, inner_ext })
+		Ok(RemoteExternalities { header: self.load_header().await?, inner_ext })
 	}
 
 	fn do_load_offline(
 		&mut self,
 		config: OfflineConfig,
-	) -> Result<RemoteExternalities<HashingFor<B>>, &'static str> {
+	) -> Result<RemoteExternalities<B>, &'static str> {
 		let mut sp = Spinner::with_timer(Spinners::Dots, "Loading snapshot...".into());
 		let start = Instant::now();
 		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::Hash>::load(&config.state_snapshot.path)?;
+		let Snapshot { snapshot_version: _, header, state_version, raw_storage, storage_root } =
+			Snapshot::<B>::load(&config.state_snapshot.path)?;
 
 		let inner_ext = TestExternalities::from_raw_snapshot(
 			raw_storage,
@@ -1110,12 +1122,10 @@ where
 		);
 		sp.stop_with_message(format!("✅ Loaded snapshot ({:.2}s)", start.elapsed().as_secs_f32()));
 
-		Ok(RemoteExternalities { inner_ext, block_hash })
+		Ok(RemoteExternalities { inner_ext, header })
 	}
 
-	pub(crate) async fn pre_build(
-		mut self,
-	) -> Result<RemoteExternalities<HashingFor<B>>, &'static str> {
+	pub(crate) async fn pre_build(mut self) -> Result<RemoteExternalities<B>, &'static str> {
 		let mut ext = match self.mode.clone() {
 			Mode::Offline(config) => self.do_load_offline(config)?,
 			Mode::Online(_) => self.do_load_remote().await?,
@@ -1191,7 +1201,7 @@ where
 		self
 	}
 
-	pub async fn build(self) -> Result<RemoteExternalities<HashingFor<B>>, &'static str> {
+	pub async fn build(self) -> Result<RemoteExternalities<B>, &'static str> {
 		let mut ext = self.pre_build().await?;
 		ext.commit_all().unwrap();
 
@@ -1226,7 +1236,7 @@ mod tests {
 		init_logger();
 		Builder::<Block>::new()
 			.mode(Mode::Offline(OfflineConfig {
-				state_snapshot: SnapshotConfig::new("test_data/proxy_test"),
+				state_snapshot: SnapshotConfig::new("test_data/test.snap"),
 			}))
 			.build()
 			.await
@@ -1241,7 +1251,7 @@ mod tests {
 		// get the first key from the snapshot file.
 		let some_key = Builder::<Block>::new()
 			.mode(Mode::Offline(OfflineConfig {
-				state_snapshot: SnapshotConfig::new("test_data/proxy_test"),
+				state_snapshot: SnapshotConfig::new("test_data/test.snap"),
 			}))
 			.build()
 			.await
@@ -1255,7 +1265,7 @@ mod tests {
 
 		Builder::<Block>::new()
 			.mode(Mode::Offline(OfflineConfig {
-				state_snapshot: SnapshotConfig::new("test_data/proxy_test"),
+				state_snapshot: SnapshotConfig::new("test_data/test.snap"),
 			}))
 			.blacklist_hashed_key(&some_key)
 			.build()
@@ -1341,7 +1351,7 @@ mod remote_tests {
 			.await
 			.unwrap();
 
-		assert_eq!(ext.block_hash, cached_ext.block_hash);
+		assert_eq!(ext.header.hash(), cached_ext.header.hash());
 	}
 
 	#[tokio::test]