1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
use jsonrpsee::{
core::{Error as JsonRpseeError, RpcResult},
proc_macros::rpc,
types::error::{CallError, ErrorCode, ErrorObject},
};
use sc_rpc_api::DenyUnsafe;
use serde::{Deserialize, Serialize};
use sp_runtime::{generic::BlockId, traits::Block as BlockT};
use std::sync::Arc;
use sp_core::{
storage::{ChildInfo, ChildType, PrefixedStorageKey},
Hasher,
};
use sp_state_machine::backend::AsTrieBackend;
use sp_trie::{
trie_types::{TrieDB, TrieDBBuilder},
KeySpacedDB, Trie,
};
use trie_db::{
node::{NodePlan, ValuePlan},
TrieDBNodeIterator,
};
fn count_migrate<'a, H: Hasher>(
storage: &'a dyn trie_db::HashDBRef<H, Vec<u8>>,
root: &'a H::Out,
) -> std::result::Result<(u64, TrieDB<'a, 'a, H>), String> {
let mut nb = 0u64;
let trie = TrieDBBuilder::new(storage, root).build();
let iter_node =
TrieDBNodeIterator::new(&trie).map_err(|e| format!("TrieDB node iterator error: {}", e))?;
for node in iter_node {
let node = node.map_err(|e| format!("TrieDB node iterator error: {}", e))?;
match node.2.node_plan() {
NodePlan::Leaf { value, .. } | NodePlan::NibbledBranch { value: Some(value), .. } =>
if let ValuePlan::Inline(range) = value {
if (range.end - range.start) as u32 >=
sp_core::storage::TRIE_VALUE_NODE_THRESHOLD
{
nb += 1;
}
},
_ => (),
}
}
Ok((nb, trie))
}
pub fn migration_status<H, B>(backend: &B) -> std::result::Result<(u64, u64), String>
where
H: Hasher,
H::Out: codec::Codec,
B: AsTrieBackend<H>,
{
let trie_backend = backend.as_trie_backend();
let essence = trie_backend.essence();
let (nb_to_migrate, trie) = count_migrate(essence, essence.root())?;
let mut nb_to_migrate_child = 0;
let mut child_roots: Vec<(ChildInfo, Vec<u8>)> = Vec::new();
for key_value in trie.iter().map_err(|e| format!("TrieDB node iterator error: {}", e))? {
let (key, value) = key_value.map_err(|e| format!("TrieDB node iterator error: {}", e))?;
if key[..].starts_with(sp_core::storage::well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX)
{
let prefixed_key = PrefixedStorageKey::new(key);
let (_type, unprefixed) = ChildType::from_prefixed_key(&prefixed_key).unwrap();
child_roots.push((ChildInfo::new_default(unprefixed), value));
}
}
for (child_info, root) in child_roots {
let mut child_root = H::Out::default();
let storage = KeySpacedDB::new(essence, child_info.keyspace());
child_root.as_mut()[..].copy_from_slice(&root[..]);
nb_to_migrate_child += count_migrate(&storage, &child_root)?.0;
}
Ok((nb_to_migrate, nb_to_migrate_child))
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct MigrationStatusResult {
top_remaining_to_migrate: u64,
child_remaining_to_migrate: u64,
}
#[rpc(server)]
pub trait StateMigrationApi<BlockHash> {
#[method(name = "state_trieMigrationStatus")]
fn call(&self, at: Option<BlockHash>) -> RpcResult<MigrationStatusResult>;
}
pub struct StateMigration<C, B, BA> {
client: Arc<C>,
backend: Arc<BA>,
deny_unsafe: DenyUnsafe,
_marker: std::marker::PhantomData<(B, BA)>,
}
impl<C, B, BA> StateMigration<C, B, BA> {
pub fn new(client: Arc<C>, backend: Arc<BA>, deny_unsafe: DenyUnsafe) -> Self {
StateMigration { client, backend, deny_unsafe, _marker: Default::default() }
}
}
impl<C, B, BA> StateMigrationApiServer<<B as BlockT>::Hash> for StateMigration<C, B, BA>
where
B: BlockT,
C: Send + Sync + 'static + sc_client_api::HeaderBackend<B>,
BA: 'static + sc_client_api::backend::Backend<B>,
{
fn call(&self, at: Option<<B as BlockT>::Hash>) -> RpcResult<MigrationStatusResult> {
self.deny_unsafe.check_if_safe()?;
let block_id = BlockId::hash(at.unwrap_or_else(|| self.client.info().best_hash));
let state = self.backend.state_at(block_id).map_err(error_into_rpc_err)?;
let (top, child) = migration_status(&state).map_err(error_into_rpc_err)?;
Ok(MigrationStatusResult {
top_remaining_to_migrate: top,
child_remaining_to_migrate: child,
})
}
}
fn error_into_rpc_err(err: impl std::fmt::Display) -> JsonRpseeError {
JsonRpseeError::Call(CallError::Custom(ErrorObject::owned(
ErrorCode::InternalError.code(),
"Error while checking migration state",
Some(err.to_string()),
)))
}