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

substrate/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)),
-			);
+			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 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 let Some(e) = response.entries.last() {
-				self.last_key = e.key.clone();
+				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);
+					}
+				}
 			}
-			for StateEntry { key, value } in response.entries {
-				self.imported_bytes += (key.len() + value.len()) as u64;
-				self.state.push((key, value))
-			}
-			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 }
 	}
 }

substrate/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;

substrate/client/network/src/state_request_handler.rs

@@ -21,7 +21,7 @@ use crate::{
 	chain::Client,
 	config::ProtocolId,
 	request_responses::{IncomingRequest, OutgoingResponse, ProtocolConfig},
-	schema::v1::{StateEntry, StateRequest, StateResponse},
+	schema::v1::{KeyValueStateEntry, StateEntry, StateRequest, StateResponse},
 	PeerId, ReputationChange,
 };
 use codec::{Decode, Encode};
@@ -66,7 +66,7 @@ fn generate_protocol_name(protocol_id: &ProtocolId) -> String {
 	let mut s = String::new();
 	s.push_str("/");
 	s.push_str(protocol_id.as_ref());
-	s.push_str("/state/1");
+	s.push_str("/state/2");
 	s
 }
 
@@ -75,7 +75,7 @@ fn generate_protocol_name(protocol_id: &ProtocolId) -> String {
 struct SeenRequestsKey<B: BlockT> {
 	peer: PeerId,
 	block: B::Hash,
-	start: Vec<u8>,
+	start: Vec<Vec<u8>>,
 }
 
 #[allow(clippy::derive_hash_xor_eq)]
@@ -169,10 +169,10 @@ impl<B: BlockT> StateRequestHandler<B> {
 
 		trace!(
 			target: LOG_TARGET,
-			"Handling state request from {}: Block {:?}, Starting at {:?}, no_proof={}",
+			"Handling state request from {}: Block {:?}, Starting at {:x?}, no_proof={}",
 			peer,
 			request.block,
-			sp_core::hexdisplay::HexDisplay::from(&request.start),
+			&request.start,
 			request.no_proof,
 		);
 
@@ -180,36 +180,45 @@ impl<B: BlockT> StateRequestHandler<B> {
 			let mut response = StateResponse::default();
 
 			if !request.no_proof {
-				let (proof, count) = self.client.read_proof_collection(
+				let (proof, _count) = self.client.read_proof_collection(
 					&BlockId::hash(block),
-					&request.start,
+					request.start.as_slice(),
 					MAX_RESPONSE_BYTES,
 				)?;
 				response.proof = proof.encode();
-				if count == 0 {
-					response.complete = true;
-				}
 			} else {
 				let entries = self.client.storage_collection(
 					&BlockId::hash(block),
-					&request.start,
+					request.start.as_slice(),
 					MAX_RESPONSE_BYTES,
 				)?;
-				response.entries =
-					entries.into_iter().map(|(key, value)| StateEntry { key, value }).collect();
-				if response.entries.is_empty() {
-					response.complete = true;
-				}
+				response.entries = entries
+					.into_iter()
+					.map(|(state, complete)| KeyValueStateEntry {
+						state_root: state.state_root,
+						entries: state
+							.key_values
+							.into_iter()
+							.map(|(key, value)| StateEntry { key, value })
+							.collect(),
+						complete,
+					})
+					.collect();
 			}
 
 			trace!(
 				target: LOG_TARGET,
-				"StateResponse contains {} keys, {}, proof nodes, complete={}, from {:?} to {:?}",
+				"StateResponse contains {} keys, {}, proof nodes, from {:?} to {:?}",
 				response.entries.len(),
 				response.proof.len(),
-				response.complete,
-				response.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
-				response.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+				response.entries.get(0).and_then(|top| top
+					.entries
+					.first()
+					.map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key))),
+				response.entries.get(0).and_then(|top| top
+					.entries
+					.last()
+					.map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key))),
 			);
 			if let Some(value) = self.seen_requests.get_mut(&key) {
 				// If this is the first time we have processed this request, we need to change