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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
use super::LegacyRpcMethods;
use crate::backend::StorageResponse;
use crate::backend::utils::retry;
use crate::config::{Config, HashFor, RpcConfigFor};
use crate::error::BackendError;
use futures::{Future, FutureExt, Stream, StreamExt};
use std::collections::VecDeque;
use std::pin::Pin;
use std::task::{Context, Poll};
/// This provides a stream of values given some prefix `key`. It
/// internally manages pagination and such.
#[allow(clippy::type_complexity)]
pub struct StorageFetchDescendantKeysStream<T: Config> {
methods: LegacyRpcMethods<RpcConfigFor<T>>,
key: Vec<u8>,
at: HashFor<T>,
// How many entries to ask for each time.
storage_page_size: u32,
// What key do we start paginating from? None = from the beginning.
pagination_start_key: Option<Vec<u8>>,
// Keys, future and cached:
keys_fut:
Option<Pin<Box<dyn Future<Output = Result<Vec<Vec<u8>>, BackendError>> + Send + 'static>>>,
// Set to true when we're done:
done: bool,
}
impl<T: Config> StorageFetchDescendantKeysStream<T> {
/// Fetch descendant keys.
pub fn new(
methods: LegacyRpcMethods<RpcConfigFor<T>>,
key: Vec<u8>,
at: HashFor<T>,
storage_page_size: u32,
) -> Self {
StorageFetchDescendantKeysStream {
methods,
key,
at,
storage_page_size,
pagination_start_key: None,
keys_fut: None,
done: false,
}
}
}
impl<T: Config> std::marker::Unpin for StorageFetchDescendantKeysStream<T> {}
impl<T: Config> Stream for StorageFetchDescendantKeysStream<T> {
type Item = Result<Vec<Vec<u8>>, BackendError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut();
loop {
// We're already done.
if this.done {
return Poll::Ready(None);
}
// Poll future to fetch next keys.
if let Some(mut keys_fut) = this.keys_fut.take() {
match keys_fut.poll_unpin(cx) {
Poll::Ready(Ok(mut keys)) => {
if this.pagination_start_key.is_some()
&& keys.first() == this.pagination_start_key.as_ref()
{
// Currently, Smoldot returns the "start key" as the first key in the input
// (see https://github.com/smol-dot/smoldot/issues/1692), whereas Substrate doesn't.
// We don't expect the start key to be returned either (since it was the last key of prev
// iteration), so remove it if we see it. This `remove()` method isn't very efficient but
// this will be a non issue with the RPC V2 APIs or if Smoldot aligns with Substrate anyway.
keys.remove(0);
}
if keys.is_empty() {
// No keys left; we're done!
this.done = true;
return Poll::Ready(None);
}
// The last key is where we want to paginate from next time.
this.pagination_start_key = keys.last().cloned();
// return all of the keys from this run.
return Poll::Ready(Some(Ok(keys)));
}
Poll::Ready(Err(e)) => {
if e.is_disconnected_will_reconnect() {
// Loop around and try again. No more keys_fut as it was taken,
// so we'll ask for the keys again from the last good pagination_start_key.
continue;
}
// Error getting keys? Return it.
return Poll::Ready(Some(Err(e)));
}
Poll::Pending => {
this.keys_fut = Some(keys_fut);
return Poll::Pending;
}
}
}
// Else, we don't have a fut to get keys yet so start one going.
let methods = this.methods.clone();
let key = this.key.clone();
let at = this.at;
let storage_page_size = this.storage_page_size;
let pagination_start_key = this.pagination_start_key.clone();
let keys_fut = async move {
let keys = methods
.state_get_keys_paged(
&key,
storage_page_size,
pagination_start_key.as_deref(),
Some(at),
)
.await?;
Ok(keys)
};
this.keys_fut = Some(Box::pin(keys_fut));
}
}
}
/// This provides a stream of values given some stream of keys.
#[allow(clippy::type_complexity)]
pub struct StorageFetchDescendantValuesStream<T: Config> {
// Stream of keys.
keys_stream: StorageFetchDescendantKeysStream<T>,
// Keys back from the stream which we are currently trying to fetch results for:
keys: Vec<Vec<u8>>,
// A future which will resolve to the resulting values:
results_fut: Option<
Pin<
Box<
dyn Future<Output = Result<Option<VecDeque<(Vec<u8>, Vec<u8>)>>, BackendError>>
+ Send
+ 'static,
>,
>,
>,
// Once we get values back we put them here and hand them back one by one to the caller.
results: VecDeque<(Vec<u8>, Vec<u8>)>,
}
impl<T: Config> StorageFetchDescendantValuesStream<T> {
/// Fetch descendant values.
pub fn new(
methods: LegacyRpcMethods<RpcConfigFor<T>>,
key: Vec<u8>,
at: HashFor<T>,
storage_page_size: u32,
) -> Self {
StorageFetchDescendantValuesStream {
keys_stream: StorageFetchDescendantKeysStream {
methods,
key,
at,
storage_page_size,
pagination_start_key: None,
keys_fut: None,
done: false,
},
keys: Default::default(),
results_fut: None,
results: Default::default(),
}
}
}
impl<T: Config> Stream for StorageFetchDescendantValuesStream<T> {
type Item = Result<StorageResponse, BackendError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut();
loop {
// If we have results back, return them one by one
if let Some((key, value)) = this.results.pop_front() {
let res = StorageResponse { key, value };
return Poll::Ready(Some(Ok(res)));
}
// If we're waiting on the next results then poll that future:
if let Some(mut results_fut) = this.results_fut.take() {
match results_fut.poll_unpin(cx) {
Poll::Ready(Ok(Some(results))) => {
// Clear keys once result comes back.
this.keys = Vec::new();
this.results = results;
continue;
}
Poll::Ready(Ok(None)) => {
// Clear keys once result comes back.
this.keys = Vec::new();
// But no results back for these keys we we just skip them.
continue;
}
Poll::Ready(Err(e)) => {
if e.is_disconnected_will_reconnect() {
// Don't replace the `results_fut` since we got disconnected, and loop around.
// This will cause us to try re-fetching results for the current keys.
continue;
}
return Poll::Ready(Some(Err(e)));
}
Poll::Pending => {
this.results_fut = Some(results_fut);
return Poll::Pending;
}
}
}
// If we have keys ready to fetch results for, then line up a results future to get them.
// The keys stream handles disconnections internally for us.
if !this.keys.is_empty() {
let methods = this.keys_stream.methods.clone();
let at = this.keys_stream.at;
let keys = this.keys.clone();
let results_fut = async move {
let keys = keys.iter().map(|k| &**k);
let values = retry(|| async {
let res = methods
.state_query_storage_at(keys.clone(), Some(at))
.await?;
Ok(res)
})
.await?;
let values: VecDeque<_> = values
.into_iter()
.flat_map(|v| {
v.changes.into_iter().filter_map(|(k, v)| {
let v = v?;
Some((k.0, v.0))
})
})
.collect();
Ok(Some(values))
};
this.results_fut = Some(Box::pin(results_fut));
continue;
}
// We have no keys yet so wait for those first.
match this.keys_stream.poll_next_unpin(cx) {
Poll::Ready(Some(Ok(keys))) => {
this.keys = keys;
continue;
}
Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => return Poll::Pending,
}
}
}
}