Skip to content
This repository was archived by the owner on Nov 15, 2023. It is now read-only.

Commit 4bd673f

Browse files
chemecheme
committed
Do not add two time child trie with same root to sync reply.
1 parent 54b2966 commit 4bd673f

File tree

6 files changed

+149
-70
lines changed

6 files changed

+149
-70
lines changed

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

Lines changed: 43 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ pub struct StateSync<B: BlockT> {
3636
target_header: B::Header,
3737
target_root: B::Hash,
3838
last_key: SmallVec<[Vec<u8>; 2]>,
39-
state: HashMap<Vec<Vec<u8>>, Vec<(Vec<u8>, Vec<u8>)>>,
39+
state: HashMap<Vec<u8>, (Vec<(Vec<u8>, Vec<u8>)>, Vec<Vec<u8>>)>,
4040
complete: bool,
4141
client: Arc<dyn Client<B>>,
4242
imported_bytes: u64,
@@ -122,28 +122,38 @@ impl<B: BlockT> StateSync<B> {
122122
};
123123

124124
for values in values.0 {
125-
use std::collections::hash_map::Entry;
126-
let key_values = if values.parent_storages.len() == 0 {
127-
// skip all child key root (will be recalculated on import)
125+
let key_values = if values.state_root.is_empty() {
126+
// Read child trie roots.
128127
values.key_values.into_iter()
129-
.filter(|key_value| !well_known_keys::is_child_storage_key(key_value.0.as_slice()))
128+
.filter(|key_value| if well_known_keys::is_child_storage_key(key_value.0.as_slice()) {
129+
self.state.entry(key_value.1.clone())
130+
.or_default().1
131+
.push(key_value.0.clone());
132+
false
133+
} else {
134+
true
135+
})
130136
.collect()
131137
} else {
132138
values.key_values
133139
};
134-
match self.state.entry(values.parent_storages) {
135-
Entry::Occupied(mut entry) => {
136-
for (key, value) in key_values {
140+
let mut entry = self.state.entry(values.state_root).or_default();
141+
if entry.0.len() > 0 && entry.1.len() > 1 {
142+
// Already imported child_trie with same root.
143+
// Warning this will not work with parallel download.
144+
} else {
145+
if entry.0.is_empty() {
146+
for (key, _value) in key_values.iter() {
137147
self.imported_bytes += key.len() as u64;
138-
entry.get_mut().push((key, value))
139148
}
140-
},
141-
Entry::Vacant(entry) => {
142-
for (key, _value) in key_values.iter() {
149+
150+
entry.0 = key_values;
151+
} else {
152+
for (key, value) in key_values {
143153
self.imported_bytes += key.len() as u64;
154+
entry.0.push((key, value))
144155
}
145-
entry.insert(key_values);
146-
},
156+
}
147157
}
148158
}
149159
self.imported_bytes += proof_size;
@@ -174,13 +184,25 @@ impl<B: BlockT> StateSync<B> {
174184
}
175185
complete = false;
176186
}
177-
let is_top = state.parent_storages.len() == 0;
178-
let entry = self.state.entry(state.parent_storages).or_default();
179-
for StateEntry { key, value } in state.entries {
180-
// Skip all child key root (will be recalculated on import).
181-
if !(is_top && well_known_keys::is_child_storage_key(key.as_slice())) {
182-
self.imported_bytes += key.len() as u64;
183-
entry.push((key, value))
187+
let is_top = state.state_root.is_empty();
188+
let entry = self.state.entry(state.state_root).or_default();
189+
if entry.0.len() > 0 && entry.1.len() > 1 {
190+
// Already imported child trie with same root.
191+
} else {
192+
let mut child_roots = Vec::new();
193+
for StateEntry { key, value } in state.entries {
194+
// Skip all child key root (will be recalculated on import).
195+
if is_top && well_known_keys::is_child_storage_key(key.as_slice()) {
196+
child_roots.push((value, key));
197+
} else {
198+
self.imported_bytes += key.len() as u64;
199+
entry.0.push((key, value))
200+
}
201+
}
202+
for (root, storage_key) in child_roots {
203+
self.state.entry(root)
204+
.or_default().1
205+
.push(storage_key);
184206
}
185207
}
186208
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,9 @@ message StateResponse {
8888

8989
// A key value state.
9090
message KeyValueStateEntry {
91-
// Parent nested storage location, empty for top state.
92-
repeated bytes parent_storages = 1;
91+
// Root of for this level, empty length bytes
92+
// if top level.
93+
bytes state_root = 1;
9394
// A collection of keys-values.
9495
repeated StateEntry entries = 2;
9596
// Set to true when there are no more keys to return.

‎client/network/src/state_request_handler.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ impl<B: BlockT> StateRequestHandler<B> {
193193
)?;
194194
response.entries = entries.into_iter().map(|(state, complete)| {
195195
KeyValueStateEntry {
196-
parent_storages: state.parent_storages,
196+
state_root: state.state_root,
197197
entries: state.key_values.into_iter()
198198
.map(|(key, value)| StateEntry { key, value }).collect(),
199199
complete,

‎client/network/test/src/sync.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1103,6 +1103,10 @@ fn syncs_state() {
11031103
data: child_data.clone(),
11041104
child_info: sp_core::storage::ChildInfo::new_default(b"child1"),
11051105
};
1106+
let child3 = sp_core::storage::StorageChild {
1107+
data: child_data.clone(),
1108+
child_info: sp_core::storage::ChildInfo::new_default(b"child3"),
1109+
};
11061110
for i in 22u8..33 {
11071111
child_data.insert(vec![i; 5], vec![i; 33]);
11081112
}
@@ -1112,6 +1116,7 @@ fn syncs_state() {
11121116
};
11131117
genesis_storage.children_default.insert(child1.child_info.storage_key().to_vec(), child1);
11141118
genesis_storage.children_default.insert(child2.child_info.storage_key().to_vec(), child2);
1119+
genesis_storage.children_default.insert(child3.child_info.storage_key().to_vec(), child3);
11151120
let mut config_one = FullPeerConfig::default();
11161121
config_one.extra_storage = Some(genesis_storage.clone());
11171122
net.add_full_peer_with_config(config_one);

‎client/service/src/client/client.rs‎

Lines changed: 40 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -807,28 +807,27 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
807807
sp_consensus::StorageChanges::Import(changes) => {
808808
let mut storage = sp_storage::Storage::default();
809809
for state in changes.state.0.into_iter() {
810-
if state.parent_storages.len() == 0 {
810+
if state.parent_storage_keys.len() == 0 && state.state_root.len() == 0 {
811811
for (key, value) in state.key_values.into_iter() {
812812
storage.top.insert(key, value);
813813
}
814-
} else if state.parent_storages.len() == 1 {
815-
let storage_key = PrefixedStorageKey::new_ref(&state.parent_storages[0]);
816-
let storage_key = match ChildType::from_prefixed_key(&storage_key) {
817-
Some((ChildType::ParentKeyId, storage_key)) => storage_key,
818-
None => return Err(Error::Backend("Invalid child storage key.".to_string())),
819-
};
820-
let entry = storage.children_default.entry(storage_key.to_vec())
821-
.or_insert_with(|| StorageChild {
822-
data: Default::default(),
823-
child_info: ChildInfo::new_default(storage_key),
824-
});
825-
for (key, value) in state.key_values.into_iter() {
826-
entry.data.insert(key, value);
814+
} else {
815+
for parent_storage in state.parent_storage_keys {
816+
let storage_key = PrefixedStorageKey::new_ref(&parent_storage);
817+
let storage_key = match ChildType::from_prefixed_key(&storage_key) {
818+
Some((ChildType::ParentKeyId, storage_key)) => storage_key,
819+
None => return Err(Error::Backend("Invalid child storage key.".to_string())),
820+
};
821+
let entry = storage.children_default.entry(storage_key.to_vec())
822+
.or_insert_with(|| StorageChild {
823+
data: Default::default(),
824+
child_info: ChildInfo::new_default(storage_key),
825+
});
826+
for (key, value) in state.key_values.iter() {
827+
entry.data.insert(key.clone(), value.clone());
828+
}
827829
}
828830
}
829-
if state.parent_storages.len() > 1 {
830-
return Err(Error::InvalidState);
831-
}
832831
}
833832

834833
let state_root = operation.op.reset_storage(storage)?;
@@ -1397,24 +1396,32 @@ impl<B, E, Block, RA> ProofProvider<Block> for Client<B, E, Block, RA> where
13971396
}
13981397
};
13991398
let mut current_child = if start_key.len() == 2 {
1400-
Some(child_info(start_key.get(0).expect("checked len"))?)
1399+
let start_key = start_key.get(0).expect("checked len");
1400+
if let Some(child_root) = state.storage(&start_key)
1401+
.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))? {
1402+
Some((child_info(start_key)?, child_root))
1403+
} else {
1404+
return Err(Error::Backend("Invalid root start key.".to_string()));
1405+
}
14011406
} else {
14021407
None
14031408
};
14041409
let mut current_key = start_key.last().map(Clone::clone).unwrap_or(Vec::new());
14051410
let mut total_size = 0;
14061411
let mut result = vec![(KeyValueStorageLevel {
1407-
parent_storages: Vec::new(),
1412+
state_root: Vec::new(),
14081413
key_values: Vec::new(),
1414+
parent_storage_keys: Vec::new(),
14091415
}, false)];
14101416

1417+
let mut child_roots = HashSet::new();
14111418
loop {
14121419
let mut entries = Vec::new();
14131420
let mut complete = true;
14141421
let mut switch_child_key = None;
14151422
while let Some(next_key) = if let Some(child) = current_child.as_ref() {
14161423
state
1417-
.next_child_storage_key(child, &current_key)
1424+
.next_child_storage_key(&child.0, &current_key)
14181425
.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?
14191426
} else {
14201427
state
@@ -1423,7 +1430,7 @@ impl<B, E, Block, RA> ProofProvider<Block> for Client<B, E, Block, RA> where
14231430
} {
14241431
let value = if let Some(child) = current_child.as_ref() {
14251432
state
1426-
.child_storage(child, next_key.as_ref())
1433+
.child_storage(&child.0, next_key.as_ref())
14271434
.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?
14281435
.unwrap_or_default()
14291436
} else {
@@ -1438,24 +1445,29 @@ impl<B, E, Block, RA> ProofProvider<Block> for Client<B, E, Block, RA> where
14381445
break;
14391446
}
14401447
total_size += size;
1441-
entries.push((next_key.clone(), value));
14421448

14431449
if current_child.is_none()
14441450
&& sp_core::storage::well_known_keys::is_child_storage_key(next_key.as_slice()) {
1445-
switch_child_key = Some(next_key);
1446-
break;
1451+
if !child_roots.contains(value.as_slice()) {
1452+
child_roots.insert(value.clone());
1453+
switch_child_key = Some((next_key.clone(), value.clone()));
1454+
entries.push((next_key.clone(), value));
1455+
break;
1456+
}
14471457
}
1458+
entries.push((next_key.clone(), value));
14481459
current_key = next_key;
14491460
}
1450-
if let Some(child) = switch_child_key.take() {
1461+
if let Some((child, child_root)) = switch_child_key.take() {
14511462
result[0].0.key_values.extend(entries.into_iter());
1452-
current_child = Some(child_info(&child)?);
1463+
current_child = Some((child_info(&child)?, child_root));
14531464
current_key = Vec::new();
1454-
} else if let Some(child) = current_child.take() {
1465+
} else if let Some((child, child_root)) = current_child.take() {
14551466
current_key = child.into_prefixed_storage_key().into_inner();
14561467
result.push((KeyValueStorageLevel {
1457-
parent_storages: vec![current_key.clone()],
1468+
state_root: child_root,
14581469
key_values: entries,
1470+
parent_storage_keys: Vec::new(),
14591471
}, complete));
14601472
if !complete {
14611473
break;

0 commit comments

Comments
 (0)