Fast sync child trie support. (#9239)

* state machine proofs.

* initial implementation

* Remove todo.

* Extend test and fix import.

* fix no proof, with proof ko.

* fix start at logic.

* Restore response size.

* Rework comments.

* Add explicit ref

* Use compact proof.

* ref change

* elaborato on empty change set condition.

* KeyValueState renaming.

* Do not add two time child trie with same root to sync reply.

* rust format

* Fix merge.

* fix warnings and fmt

* fmt

* update protocol id to V2

client/api/src/backend.rs

@@ -42,7 +42,7 @@ use std::{
 	sync::Arc,
 };
 
-pub use sp_state_machine::Backend as StateBackend;
+pub use sp_state_machine::{Backend as StateBackend, KeyValueStates};
 use std::marker::PhantomData;
 
 /// Extracts the state backend type for the given backend.

client/api/src/lib.rs

@@ -39,7 +39,7 @@ pub use proof_provider::*;
 pub use sp_blockchain as blockchain;
 pub use sp_blockchain::HeaderBackend;
 
-pub use sp_state_machine::{ExecutionStrategy, StorageProof};
+pub use sp_state_machine::{CompactProof, ExecutionStrategy, StorageProof};
 pub use sp_storage::{ChildInfo, PrefixedStorageKey, StorageData, StorageKey};
 
 /// Usage Information Provider interface

client/api/src/proof_provider.rs

@@ -17,8 +17,9 @@
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
 //! Proof utilities
-use crate::{ChangesProof, StorageProof};
+use crate::{ChangesProof, CompactProof, StorageProof};
 use sp_runtime::{generic::BlockId, traits::Block as BlockT};
+use sp_state_machine::{KeyValueStates, KeyValueStorageLevel};
 use sp_storage::{ChildInfo, PrefixedStorageKey, StorageKey};
 
 /// Interface for providing block proving utilities.
@@ -71,31 +72,43 @@ pub trait ProofProvider<Block: BlockT> {
 		key: &StorageKey,
 	) -> sp_blockchain::Result<ChangesProof<Block::Header>>;
 
-	/// Given a `BlockId` iterate over all storage values starting at `start_key` exclusively,
-	/// building proofs until size limit is reached. Returns combined proof and the number of
-	/// collected keys.
+	/// Given a `BlockId` iterate over all storage values starting at `start_keys`.
+	/// Last `start_keys` element contains last accessed key value.
+	/// With multiple `start_keys`, first `start_keys` element is
+	/// the current storage key of of the last accessed child trie.
+	/// at last level the value to start at exclusively.
+	/// Proofs is build until size limit is reached and always include at
+	/// least one key following `start_keys`.
+	/// Returns combined proof and the numbers of collected keys.
 	fn read_proof_collection(
 		&self,
 		id: &BlockId<Block>,
-		start_key: &[u8],
+		start_keys: &[Vec<u8>],
 		size_limit: usize,
-	) -> sp_blockchain::Result<(StorageProof, u32)>;
+	) -> sp_blockchain::Result<(CompactProof, u32)>;
 
 	/// Given a `BlockId` iterate over all storage values starting at `start_key`.
 	/// Returns collected keys and values.
+	/// Returns the collected keys values content of the top trie followed by the
+	/// collected keys values of child tries.
+	/// Only child tries with their root part of the collected content or
+	/// related to `start_key` are attached.
+	/// For each collected state a boolean indicates if state reach
+	/// end.
 	fn storage_collection(
 		&self,
 		id: &BlockId<Block>,
-		start_key: &[u8],
+		start_key: &[Vec<u8>],
 		size_limit: usize,
-	) -> sp_blockchain::Result<Vec<(Vec<u8>, Vec<u8>)>>;
+	) -> sp_blockchain::Result<Vec<(KeyValueStorageLevel, bool)>>;
 
 	/// Verify read storage proof for a set of keys.
-	/// Returns collected key-value pairs and a flag indicating if iteration is complete.
+	/// Returns collected key-value pairs and a the nested state
+	/// depth of current iteration or 0 if completed.
 	fn verify_range_proof(
 		&self,
 		root: Block::Hash,
-		proof: StorageProof,
-		start_key: &[u8],
-	) -> sp_blockchain::Result<(Vec<(Vec<u8>, Vec<u8>)>, bool)>;
+		proof: CompactProof,
+		start_keys: &[Vec<u8>],
+	) -> sp_blockchain::Result<(KeyValueStates, usize)>;
 }

client/consensus/common/src/block_import.rs

@@ -133,7 +133,7 @@ pub struct ImportedState<B: BlockT> {
 	/// Target block hash.
 	pub block: B::Hash,
 	/// State keys and values.
-	pub state: Vec<(Vec<u8>, Vec<u8>)>,
+	pub state: sp_state_machine::KeyValueStates,
 }
 
 impl<B: BlockT> std::fmt::Debug for ImportedState<B> {

client/network/src/protocol/sync/state.rs

@@ -23,9 +23,11 @@ use crate::{
 };
 use codec::{Decode, Encode};
 use log::debug;
-use sc_client_api::StorageProof;
+use sc_client_api::CompactProof;
+use smallvec::SmallVec;
+use sp_core::storage::well_known_keys;
 use sp_runtime::traits::{Block as BlockT, Header, NumberFor};
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 /// State sync support.
 
@@ -35,8 +37,8 @@ pub struct StateSync<B: BlockT> {
 	target_block: B::Hash,
 	target_header: B::Header,
 	target_root: B::Hash,
-	last_key: Vec<u8>,
-	state: Vec<(Vec<u8>, Vec<u8>)>,
+	last_key: SmallVec<[Vec<u8>; 2]>,
+	state: HashMap<Vec<u8>, (Vec<(Vec<u8>, Vec<u8>)>, Vec<Vec<u8>>)>,
 	complete: bool,
 	client: Arc<dyn Client<B>>,
 	imported_bytes: u64,
@@ -61,8 +63,8 @@ impl<B: BlockT> StateSync<B> {
 			target_block: target.hash(),
 			target_root: target.state_root().clone(),
 			target_header: target,
-			last_key: Vec::default(),
-			state: Vec::default(),
+			last_key: SmallVec::default(),
+			state: HashMap::default(),
 			complete: false,
 			imported_bytes: 0,
 			skip_proof,
@@ -71,7 +73,7 @@ impl<B: BlockT> StateSync<B> {
 
 	///  Validate and import a state reponse.
 	pub fn import(&mut self, response: StateResponse) -> ImportResult<B> {
-		if response.entries.is_empty() && response.proof.is_empty() && !response.complete {
+		if response.entries.is_empty() && response.proof.is_empty() {
 			debug!(target: "sync", "Bad state response");
 			return ImportResult::BadResponse
 		}
@@ -82,56 +84,135 @@ impl<B: BlockT> StateSync<B> {
 		let complete = if !self.skip_proof {
 			debug!(target: "sync", "Importing state from {} trie nodes", response.proof.len());
 			let proof_size = response.proof.len() as u64;
-			let proof = match StorageProof::decode(&mut response.proof.as_ref()) {
+			let proof = match CompactProof::decode(&mut response.proof.as_ref()) {
 				Ok(proof) => proof,
 				Err(e) => {
 					debug!(target: "sync", "Error decoding proof: {:?}", e);
 					return ImportResult::BadResponse
 				},
 			};
-			let (values, complete) =
-				match self.client.verify_range_proof(self.target_root, proof, &self.last_key) {
-					Err(e) => {
-						debug!(target: "sync", "StateResponse failed proof verification: {:?}", e);
-						return ImportResult::BadResponse
-					},
-					Ok(values) => values,
-				};
+			let (values, completed) = match self.client.verify_range_proof(
+				self.target_root,
+				proof,
+				self.last_key.as_slice(),
+			) {
+				Err(e) => {
+					debug!(
+						target: "sync",
+						"StateResponse failed proof verification: {:?}",
+						e,
+					);
+					return ImportResult::BadResponse
+				},
+				Ok(values) => values,
+			};
 			debug!(target: "sync", "Imported with {} keys", values.len());
 
-			if let Some(last) = values.last().map(|(k, _)| k) {
-				self.last_key = last.clone();
-			}
+			let complete = completed == 0;
+			if !complete && !values.update_last_key(completed, &mut self.last_key) {
+				debug!(target: "sync", "Error updating key cursor, depth: {}", completed);
+			};
 
-			for (key, value) in values {
-				self.imported_bytes += key.len() as u64;
-				self.state.push((key, value))
+			for values in values.0 {
+				let key_values = if values.state_root.is_empty() {
+					// Read child trie roots.
+					values
+						.key_values
+						.into_iter()
+						.filter(|key_value| {
+							if well_known_keys::is_child_storage_key(key_value.0.as_slice()) {
+								self.state
+									.entry(key_value.1.clone())
+									.or_default()
+									.1
+									.push(key_value.0.clone());
+								false
+							} else {
+								true
+							}
+						})
+						.collect()
+				} else {
+					values.key_values
+				};
+				let mut entry = self.state.entry(values.state_root).or_default();
+				if entry.0.len() > 0 && entry.1.len() > 1 {
+					// Already imported child_trie with same root.
+					// Warning this will not work with parallel download.
+				} else {
+					if entry.0.is_empty() {
+						for (key, _value) in key_values.iter() {
+							self.imported_bytes += key.len() as u64;
+						}
+
+						entry.0 = key_values;
+					} else {
+						for (key, value) in key_values {
+							self.imported_bytes += key.len() as u64;
+							entry.0.push((key, value))
+						}
+					}
+				}
 			}
 			self.imported_bytes += proof_size;
 			complete
 		} else {
-			debug!(
-				target: "sync",
-				"Importing state from {:?} to {:?}",
-				response.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
-				response.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
-			);
-
-			if let Some(e) = response.entries.last() {
-				self.last_key = e.key.clone();
+			let mut complete = true;
+			// if the trie is a child trie and one of its parent trie is empty,
+			// the parent cursor stays valid.
+			// Empty parent trie content only happens when all the response content
+			// is part of a single child trie.
+			if self.last_key.len() == 2 && response.entries[0].entries.len() == 0 {
+				// Do not remove the parent trie position.
+				self.last_key.pop();
+			} else {
+				self.last_key.clear();
 			}
-			for StateEntry { key, value } in response.entries {
-				self.imported_bytes += (key.len() + value.len()) as u64;
-				self.state.push((key, value))
+			for state in response.entries {
+				debug!(
+					target: "sync",
+					"Importing state from {:?} to {:?}",
+					state.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+					state.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+				);
+
+				if !state.complete {
+					if let Some(e) = state.entries.last() {
+						self.last_key.push(e.key.clone());
+					}
+					complete = false;
+				}
+				let is_top = state.state_root.is_empty();
+				let entry = self.state.entry(state.state_root).or_default();
+				if entry.0.len() > 0 && entry.1.len() > 1 {
+					// Already imported child trie with same root.
+				} else {
+					let mut child_roots = Vec::new();
+					for StateEntry { key, value } in state.entries {
+						// Skip all child key root (will be recalculated on import).
+						if is_top && well_known_keys::is_child_storage_key(key.as_slice()) {
+							child_roots.push((value, key));
+						} else {
+							self.imported_bytes += key.len() as u64;
+							entry.0.push((key, value))
+						}
+					}
+					for (root, storage_key) in child_roots {
+						self.state.entry(root).or_default().1.push(storage_key);
+					}
+				}
 			}
-			response.complete
+			complete
 		};
 		if complete {
 			self.complete = true;
 			ImportResult::Import(
 				self.target_block,
 				self.target_header.clone(),
-				ImportedState { block: self.target_block, state: std::mem::take(&mut self.state) },
+				ImportedState {
+					block: self.target_block.clone(),
+					state: std::mem::take(&mut self.state).into(),
+				},
 			)
 		} else {
 			ImportResult::Continue
@@ -142,7 +223,7 @@ impl<B: BlockT> StateSync<B> {
 	pub fn next_request(&self) -> StateRequest {
 		StateRequest {
 			block: self.target_block.encode(),
-			start: self.last_key.clone(),
+			start: self.last_key.clone().into_vec(),
 			no_proof: self.skip_proof,
 		}
 	}
@@ -164,7 +245,8 @@ impl<B: BlockT> StateSync<B> {
 
 	/// Returns state sync estimated progress.
 	pub fn progress(&self) -> StateDownloadProgress {
-		let percent_done = (*self.last_key.get(0).unwrap_or(&0u8) as u32) * 100 / 256;
+		let cursor = *self.last_key.get(0).and_then(|last| last.get(0)).unwrap_or(&0u8);
+		let percent_done = cursor as u32 * 100 / 256;
 		StateDownloadProgress { percentage: percent_done, size: self.imported_bytes }
 	}
 }

client/network/src/schema/api.v1.proto

@@ -74,22 +74,32 @@ message BlockData {
 message StateRequest {
 	// Block header hash.
 	bytes block = 1;
-	// Start from this key. Equivalent to <empty bytes> if omitted.
-	bytes start = 2; // optional
+	// Start from this key.
+	// Multiple keys used for nested state start.
+	repeated bytes start = 2; // optional
 	// if 'true' indicates that response should contain raw key-values, rather than proof.
 	bool no_proof = 3;
 }
 
 message StateResponse {
-	// A collection of keys-values. Only populated if `no_proof` is `true`
-	repeated StateEntry entries = 1;
+	// A collection of keys-values states. Only populated if `no_proof` is `true`
+	repeated KeyValueStateEntry entries = 1;
 	// If `no_proof` is false in request, this contains proof nodes.
 	bytes proof = 2;
+}
+
+// A key value state.
+message KeyValueStateEntry {
+	// Root of for this level, empty length bytes
+	// if top level.
+	bytes state_root = 1;
+	// A collection of keys-values.
+	repeated StateEntry entries = 2;
 	// Set to true when there are no more keys to return.
 	bool complete = 3;
 }
 
-// A key-value pair
+// A key-value pair.
 message StateEntry {
 	bytes key = 1;
 	bytes value = 2;