1use std::{
22 collections::{hash_map::Entry, HashMap},
23 sync::Arc,
24};
25
26use itertools::Itertools;
27use sc_client_api::{Backend, ChildInfo, StorageKey, StorageProvider};
28use sp_runtime::traits::Block as BlockT;
29
30use super::error::Error as ArchiveError;
31use crate::{
32 archive::archive::LOG_TARGET,
33 common::{
34 events::{
35 ArchiveStorageDiffEvent, ArchiveStorageDiffItem, ArchiveStorageDiffOperationType,
36 ArchiveStorageDiffResult, ArchiveStorageDiffType, StorageResult,
37 },
38 storage::Storage,
39 },
40};
41use tokio::sync::mpsc;
42
43pub fn parse_hex_param(param: String) -> Result<Vec<u8>, ArchiveError> {
47 if param.is_empty() {
49 return Ok(Default::default())
50 }
51
52 array_bytes::hex2bytes(¶m).map_err(|_| ArchiveError::InvalidParam(param))
53}
54
55#[derive(Debug, PartialEq, Clone)]
56pub struct DiffDetails {
57 key: StorageKey,
58 return_type: ArchiveStorageDiffType,
59 child_trie_key: Option<ChildInfo>,
60 child_trie_key_string: Option<String>,
61}
62
63#[derive(Debug, PartialEq, Clone, Copy)]
65enum FetchStorageType {
66 Value,
68 Hash,
70 Both,
72}
73
74#[derive(Debug, PartialEq, Clone)]
76enum FetchedStorage {
77 Value(StorageResult),
79 Hash(StorageResult),
81 Both { value: StorageResult, hash: StorageResult },
83}
84
85pub struct ArchiveStorageDiff<Client, Block, BE> {
86 client: Storage<Client, Block, BE>,
87}
88
89impl<Client, Block, BE> ArchiveStorageDiff<Client, Block, BE> {
90 pub fn new(client: Arc<Client>) -> Self {
91 Self { client: Storage::new(client) }
92 }
93}
94
95impl<Client, Block, BE> ArchiveStorageDiff<Client, Block, BE>
96where
97 Block: BlockT + 'static,
98 BE: Backend<Block> + 'static,
99 Client: StorageProvider<Block, BE> + Send + Sync + 'static,
100{
101 fn fetch_storage(
103 &self,
104 hash: Block::Hash,
105 key: StorageKey,
106 maybe_child_trie: Option<ChildInfo>,
107 ty: FetchStorageType,
108 ) -> Result<Option<FetchedStorage>, String> {
109 match ty {
110 FetchStorageType::Value => {
111 let result = self.client.query_value(hash, &key, maybe_child_trie.as_ref())?;
112
113 Ok(result.map(FetchedStorage::Value))
114 },
115
116 FetchStorageType::Hash => {
117 let result = self.client.query_hash(hash, &key, maybe_child_trie.as_ref())?;
118
119 Ok(result.map(FetchedStorage::Hash))
120 },
121
122 FetchStorageType::Both => {
123 let Some(value) = self.client.query_value(hash, &key, maybe_child_trie.as_ref())?
124 else {
125 return Ok(None);
126 };
127
128 let Some(hash) = self.client.query_hash(hash, &key, maybe_child_trie.as_ref())?
129 else {
130 return Ok(None);
131 };
132
133 Ok(Some(FetchedStorage::Both { value, hash }))
134 },
135 }
136 }
137
138 fn belongs_to_query(key: &StorageKey, items: &[DiffDetails]) -> Option<FetchStorageType> {
147 if items.is_empty() {
149 return Some(FetchStorageType::Value)
150 }
151
152 let mut value = false;
153 let mut hash = false;
154
155 for item in items {
156 if key.as_ref().starts_with(&item.key.as_ref()) {
157 match item.return_type {
158 ArchiveStorageDiffType::Value => value = true,
159 ArchiveStorageDiffType::Hash => hash = true,
160 }
161 }
162 }
163
164 match (value, hash) {
165 (true, true) => Some(FetchStorageType::Both),
166 (true, false) => Some(FetchStorageType::Value),
167 (false, true) => Some(FetchStorageType::Hash),
168 (false, false) => None,
169 }
170 }
171
172 fn send_result(
176 tx: &mpsc::Sender<ArchiveStorageDiffEvent>,
177 result: FetchedStorage,
178 operation_type: ArchiveStorageDiffOperationType,
179 child_trie_key: Option<String>,
180 ) -> bool {
181 let items = match result {
182 FetchedStorage::Value(storage_result) | FetchedStorage::Hash(storage_result) =>
183 vec![storage_result],
184 FetchedStorage::Both { value, hash } => vec![value, hash],
185 };
186
187 for item in items {
188 let res = ArchiveStorageDiffEvent::StorageDiff(ArchiveStorageDiffResult {
189 key: item.key,
190 result: item.result,
191 operation_type,
192 child_trie_key: child_trie_key.clone(),
193 });
194 if tx.blocking_send(res).is_err() {
195 return false
196 }
197 }
198
199 true
200 }
201
202 fn handle_trie_queries_inner(
203 &self,
204 hash: Block::Hash,
205 previous_hash: Block::Hash,
206 items: Vec<DiffDetails>,
207 tx: &mpsc::Sender<ArchiveStorageDiffEvent>,
208 ) -> Result<(), String> {
209 let maybe_child_trie = items.first().and_then(|item| item.child_trie_key.clone());
211 let maybe_child_trie_str =
212 items.first().and_then(|item| item.child_trie_key_string.clone());
213
214 let keys_iter = self.client.raw_keys_iter(hash, maybe_child_trie.clone())?;
218 let previous_keys_iter =
219 self.client.raw_keys_iter(previous_hash, maybe_child_trie.clone())?;
220
221 let mut diff_iter = lexicographic_diff(keys_iter, previous_keys_iter);
222
223 while let Some(item) = diff_iter.next() {
224 let (operation_type, key) = match item {
225 Diff::Added(key) => (ArchiveStorageDiffOperationType::Added, key),
226 Diff::Deleted(key) => (ArchiveStorageDiffOperationType::Deleted, key),
227 Diff::Equal(key) => (ArchiveStorageDiffOperationType::Modified, key),
228 };
229
230 let Some(fetch_type) = Self::belongs_to_query(&key, &items) else {
231 continue;
233 };
234
235 let maybe_result = match operation_type {
236 ArchiveStorageDiffOperationType::Added =>
237 self.fetch_storage(hash, key.clone(), maybe_child_trie.clone(), fetch_type)?,
238 ArchiveStorageDiffOperationType::Deleted => self.fetch_storage(
239 previous_hash,
240 key.clone(),
241 maybe_child_trie.clone(),
242 fetch_type,
243 )?,
244 ArchiveStorageDiffOperationType::Modified => {
245 let Some(storage_result) = self.fetch_storage(
246 hash,
247 key.clone(),
248 maybe_child_trie.clone(),
249 fetch_type,
250 )?
251 else {
252 continue
253 };
254
255 let Some(previous_storage_result) = self.fetch_storage(
256 previous_hash,
257 key.clone(),
258 maybe_child_trie.clone(),
259 fetch_type,
260 )?
261 else {
262 continue
263 };
264
265 if storage_result == previous_storage_result {
267 continue
268 }
269
270 Some(storage_result)
271 },
272 };
273
274 if let Some(storage_result) = maybe_result {
275 if !Self::send_result(
276 &tx,
277 storage_result,
278 operation_type,
279 maybe_child_trie_str.clone(),
280 ) {
281 return Ok(())
282 }
283 }
284 }
285
286 Ok(())
287 }
288
289 pub async fn handle_trie_queries(
293 &self,
294 hash: Block::Hash,
295 items: Vec<ArchiveStorageDiffItem<String>>,
296 previous_hash: Block::Hash,
297 tx: mpsc::Sender<ArchiveStorageDiffEvent>,
298 ) -> Result<(), tokio::task::JoinError> {
299 let this = ArchiveStorageDiff { client: self.client.clone() };
300
301 tokio::task::spawn_blocking(move || {
302 let mut trie_items = match deduplicate_storage_diff_items(items) {
304 Ok(items) => items,
305 Err(error) => {
306 let _ = tx.blocking_send(ArchiveStorageDiffEvent::err(error.to_string()));
307 return
308 },
309 };
310 if trie_items.is_empty() {
312 trie_items.push(Vec::new());
313 }
314 log::trace!(target: LOG_TARGET, "Storage diff deduplicated items: {:?}", trie_items);
315
316 for items in trie_items {
317 log::trace!(
318 target: LOG_TARGET,
319 "handle_trie_queries: hash={:?}, previous_hash={:?}, items={:?}",
320 hash,
321 previous_hash,
322 items
323 );
324
325 let result = this.handle_trie_queries_inner(hash, previous_hash, items, &tx);
326
327 if let Err(error) = result {
328 log::trace!(
329 target: LOG_TARGET,
330 "handle_trie_queries: sending error={:?}",
331 error,
332 );
333
334 let _ = tx.blocking_send(ArchiveStorageDiffEvent::err(error));
335
336 return
337 } else {
338 log::trace!(
339 target: LOG_TARGET,
340 "handle_trie_queries: sending storage diff done",
341 );
342 }
343 }
344
345 let _ = tx.blocking_send(ArchiveStorageDiffEvent::StorageDiffDone);
346 })
347 .await?;
348
349 Ok(())
350 }
351}
352
353#[derive(Debug, PartialEq)]
355enum Diff<T> {
356 Added(T),
357 Deleted(T),
358 Equal(T),
359}
360
361fn lexicographic_diff<T, LeftIter, RightIter>(
363 mut left: LeftIter,
364 mut right: RightIter,
365) -> impl Iterator<Item = Diff<T>>
366where
367 T: Ord,
368 LeftIter: Iterator<Item = T>,
369 RightIter: Iterator<Item = T>,
370{
371 let mut a = left.next();
372 let mut b = right.next();
373
374 core::iter::from_fn(move || match (a.take(), b.take()) {
375 (Some(a_value), Some(b_value)) =>
376 if a_value < b_value {
377 b = Some(b_value);
378 a = left.next();
379
380 Some(Diff::Added(a_value))
381 } else if a_value > b_value {
382 a = Some(a_value);
383 b = right.next();
384
385 Some(Diff::Deleted(b_value))
386 } else {
387 a = left.next();
388 b = right.next();
389
390 Some(Diff::Equal(a_value))
391 },
392 (Some(a_value), None) => {
393 a = left.next();
394 Some(Diff::Added(a_value))
395 },
396 (None, Some(b_value)) => {
397 b = right.next();
398 Some(Diff::Deleted(b_value))
399 },
400 (None, None) => None,
401 })
402}
403
404fn deduplicate_storage_diff_items(
408 items: Vec<ArchiveStorageDiffItem<String>>,
409) -> Result<Vec<Vec<DiffDetails>>, ArchiveError> {
410 let mut deduplicated: HashMap<Option<ChildInfo>, Vec<DiffDetails>> = HashMap::new();
411
412 for diff_item in items {
413 let key = StorageKey(parse_hex_param(diff_item.key)?);
415 let child_trie_key_string = diff_item.child_trie_key.clone();
416 let child_trie_key = diff_item
417 .child_trie_key
418 .map(|child_trie_key| parse_hex_param(child_trie_key))
419 .transpose()?
420 .map(ChildInfo::new_default_from_vec);
421
422 let diff_item = DiffDetails {
423 key,
424 return_type: diff_item.return_type,
425 child_trie_key: child_trie_key.clone(),
426 child_trie_key_string,
427 };
428
429 match deduplicated.entry(child_trie_key.clone()) {
430 Entry::Occupied(mut entry) => {
431 let mut should_insert = true;
432
433 for existing in entry.get() {
434 if existing.return_type != diff_item.return_type {
436 continue
437 }
438 if existing.key == diff_item.key {
440 should_insert = false;
441 break
442 }
443
444 if diff_item.key.as_ref().starts_with(&existing.key.as_ref()) {
448 should_insert = false;
449 break
450 }
451
452 if existing.key.as_ref().starts_with(&diff_item.key.as_ref()) {
455 let to_remove = existing.clone();
456 entry.get_mut().retain(|item| item != &to_remove);
457 break;
458 }
459 }
460
461 if should_insert {
462 entry.get_mut().push(diff_item);
463 }
464 },
465 Entry::Vacant(entry) => {
466 entry.insert(vec![diff_item]);
467 },
468 }
469 }
470
471 Ok(deduplicated
472 .into_iter()
473 .sorted_by_key(|(child_trie_key, _)| child_trie_key.clone())
474 .map(|(_, values)| values)
475 .collect())
476}
477
478#[cfg(test)]
479mod tests {
480 use super::*;
481
482 #[test]
483 fn dedup_empty() {
484 let items = vec![];
485 let result = deduplicate_storage_diff_items(items).unwrap();
486 assert!(result.is_empty());
487 }
488
489 #[test]
490 fn dedup_single() {
491 let items = vec![ArchiveStorageDiffItem {
492 key: "0x01".into(),
493 return_type: ArchiveStorageDiffType::Value,
494 child_trie_key: None,
495 }];
496 let result = deduplicate_storage_diff_items(items).unwrap();
497 assert_eq!(result.len(), 1);
498 assert_eq!(result[0].len(), 1);
499
500 let expected = DiffDetails {
501 key: StorageKey(vec![1]),
502 return_type: ArchiveStorageDiffType::Value,
503 child_trie_key: None,
504 child_trie_key_string: None,
505 };
506 assert_eq!(result[0][0], expected);
507 }
508
509 #[test]
510 fn dedup_with_different_keys() {
511 let items = vec![
512 ArchiveStorageDiffItem {
513 key: "0x01".into(),
514 return_type: ArchiveStorageDiffType::Value,
515 child_trie_key: None,
516 },
517 ArchiveStorageDiffItem {
518 key: "0x02".into(),
519 return_type: ArchiveStorageDiffType::Value,
520 child_trie_key: None,
521 },
522 ];
523 let result = deduplicate_storage_diff_items(items).unwrap();
524 assert_eq!(result.len(), 1);
525 assert_eq!(result[0].len(), 2);
526
527 let expected = vec![
528 DiffDetails {
529 key: StorageKey(vec![1]),
530 return_type: ArchiveStorageDiffType::Value,
531 child_trie_key: None,
532 child_trie_key_string: None,
533 },
534 DiffDetails {
535 key: StorageKey(vec![2]),
536 return_type: ArchiveStorageDiffType::Value,
537 child_trie_key: None,
538 child_trie_key_string: None,
539 },
540 ];
541 assert_eq!(result[0], expected);
542 }
543
544 #[test]
545 fn dedup_with_same_keys() {
546 let items = vec![
548 ArchiveStorageDiffItem {
549 key: "0x01".into(),
550 return_type: ArchiveStorageDiffType::Value,
551 child_trie_key: None,
552 },
553 ArchiveStorageDiffItem {
554 key: "0x01".into(),
555 return_type: ArchiveStorageDiffType::Value,
556 child_trie_key: None,
557 },
558 ];
559 let result = deduplicate_storage_diff_items(items).unwrap();
560 assert_eq!(result.len(), 1);
561 assert_eq!(result[0].len(), 1);
562
563 let expected = vec![DiffDetails {
564 key: StorageKey(vec![1]),
565 return_type: ArchiveStorageDiffType::Value,
566 child_trie_key: None,
567 child_trie_key_string: None,
568 }];
569 assert_eq!(result[0], expected);
570 }
571
572 #[test]
573 fn dedup_with_same_prefix() {
574 let items = vec![
576 ArchiveStorageDiffItem {
577 key: "0x01".into(),
578 return_type: ArchiveStorageDiffType::Value,
579 child_trie_key: None,
580 },
581 ArchiveStorageDiffItem {
582 key: "0x01ff".into(),
583 return_type: ArchiveStorageDiffType::Value,
584 child_trie_key: None,
585 },
586 ];
587 let result = deduplicate_storage_diff_items(items).unwrap();
588 assert_eq!(result.len(), 1);
589 assert_eq!(result[0].len(), 1);
590
591 let expected = vec![DiffDetails {
592 key: StorageKey(vec![1]),
593 return_type: ArchiveStorageDiffType::Value,
594 child_trie_key: None,
595 child_trie_key_string: None,
596 }];
597 assert_eq!(result[0], expected);
598 }
599
600 #[test]
601 fn dedup_with_different_return_types() {
602 let items = vec![
603 ArchiveStorageDiffItem {
604 key: "0x01".into(),
605 return_type: ArchiveStorageDiffType::Value,
606 child_trie_key: None,
607 },
608 ArchiveStorageDiffItem {
609 key: "0x01".into(),
610 return_type: ArchiveStorageDiffType::Hash,
611 child_trie_key: None,
612 },
613 ];
614 let result = deduplicate_storage_diff_items(items).unwrap();
615 assert_eq!(result.len(), 1);
616 assert_eq!(result[0].len(), 2);
617
618 let expected = vec![
619 DiffDetails {
620 key: StorageKey(vec![1]),
621 return_type: ArchiveStorageDiffType::Value,
622 child_trie_key: None,
623 child_trie_key_string: None,
624 },
625 DiffDetails {
626 key: StorageKey(vec![1]),
627 return_type: ArchiveStorageDiffType::Hash,
628 child_trie_key: None,
629 child_trie_key_string: None,
630 },
631 ];
632 assert_eq!(result[0], expected);
633 }
634
635 #[test]
636 fn dedup_with_different_child_tries() {
637 let items = vec![
638 ArchiveStorageDiffItem {
639 key: "0x01".into(),
640 return_type: ArchiveStorageDiffType::Value,
641 child_trie_key: Some("0x01".into()),
642 },
643 ArchiveStorageDiffItem {
644 key: "0x01".into(),
645 return_type: ArchiveStorageDiffType::Value,
646 child_trie_key: Some("0x02".into()),
647 },
648 ];
649 let result = deduplicate_storage_diff_items(items).unwrap();
650 assert_eq!(result.len(), 2);
651 assert_eq!(result[0].len(), 1);
652 assert_eq!(result[1].len(), 1);
653
654 let expected = vec![
655 vec![DiffDetails {
656 key: StorageKey(vec![1]),
657 return_type: ArchiveStorageDiffType::Value,
658 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![1])),
659 child_trie_key_string: Some("0x01".into()),
660 }],
661 vec![DiffDetails {
662 key: StorageKey(vec![1]),
663 return_type: ArchiveStorageDiffType::Value,
664 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![2])),
665 child_trie_key_string: Some("0x02".into()),
666 }],
667 ];
668 assert_eq!(result, expected);
669 }
670
671 #[test]
672 fn dedup_with_same_child_tries() {
673 let items = vec![
674 ArchiveStorageDiffItem {
675 key: "0x01".into(),
676 return_type: ArchiveStorageDiffType::Value,
677 child_trie_key: Some("0x01".into()),
678 },
679 ArchiveStorageDiffItem {
680 key: "0x01".into(),
681 return_type: ArchiveStorageDiffType::Value,
682 child_trie_key: Some("0x01".into()),
683 },
684 ];
685 let result = deduplicate_storage_diff_items(items).unwrap();
686 assert_eq!(result.len(), 1);
687 assert_eq!(result[0].len(), 1);
688
689 let expected = vec![DiffDetails {
690 key: StorageKey(vec![1]),
691 return_type: ArchiveStorageDiffType::Value,
692 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![1])),
693 child_trie_key_string: Some("0x01".into()),
694 }];
695 assert_eq!(result[0], expected);
696 }
697
698 #[test]
699 fn dedup_with_shorter_key_reverse_order() {
700 let items = vec![
701 ArchiveStorageDiffItem {
702 key: "0x01ff".into(),
703 return_type: ArchiveStorageDiffType::Value,
704 child_trie_key: None,
705 },
706 ArchiveStorageDiffItem {
707 key: "0x01".into(),
708 return_type: ArchiveStorageDiffType::Value,
709 child_trie_key: None,
710 },
711 ];
712 let result = deduplicate_storage_diff_items(items).unwrap();
713 assert_eq!(result.len(), 1);
714 assert_eq!(result[0].len(), 1);
715
716 let expected = vec![DiffDetails {
717 key: StorageKey(vec![1]),
718 return_type: ArchiveStorageDiffType::Value,
719 child_trie_key: None,
720 child_trie_key_string: None,
721 }];
722 assert_eq!(result[0], expected);
723 }
724
725 #[test]
726 fn dedup_multiple_child_tries() {
727 let items = vec![
728 ArchiveStorageDiffItem {
729 key: "0x02".into(),
730 return_type: ArchiveStorageDiffType::Value,
731 child_trie_key: None,
732 },
733 ArchiveStorageDiffItem {
734 key: "0x01".into(),
735 return_type: ArchiveStorageDiffType::Value,
736 child_trie_key: Some("0x01".into()),
737 },
738 ArchiveStorageDiffItem {
739 key: "0x02".into(),
740 return_type: ArchiveStorageDiffType::Hash,
741 child_trie_key: Some("0x01".into()),
742 },
743 ArchiveStorageDiffItem {
744 key: "0x01".into(),
745 return_type: ArchiveStorageDiffType::Value,
746 child_trie_key: Some("0x02".into()),
747 },
748 ArchiveStorageDiffItem {
749 key: "0x01".into(),
750 return_type: ArchiveStorageDiffType::Hash,
751 child_trie_key: Some("0x02".into()),
752 },
753 ArchiveStorageDiffItem {
754 key: "0x01ff".into(),
755 return_type: ArchiveStorageDiffType::Value,
756 child_trie_key: Some("0x02".into()),
757 },
758 ];
759
760 let result = deduplicate_storage_diff_items(items).unwrap();
761
762 let expected = vec![
763 vec![DiffDetails {
764 key: StorageKey(vec![2]),
765 return_type: ArchiveStorageDiffType::Value,
766 child_trie_key: None,
767 child_trie_key_string: None,
768 }],
769 vec![
770 DiffDetails {
771 key: StorageKey(vec![1]),
772 return_type: ArchiveStorageDiffType::Value,
773 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![1])),
774 child_trie_key_string: Some("0x01".into()),
775 },
776 DiffDetails {
777 key: StorageKey(vec![2]),
778 return_type: ArchiveStorageDiffType::Hash,
779 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![1])),
780 child_trie_key_string: Some("0x01".into()),
781 },
782 ],
783 vec![
784 DiffDetails {
785 key: StorageKey(vec![1]),
786 return_type: ArchiveStorageDiffType::Value,
787 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![2])),
788 child_trie_key_string: Some("0x02".into()),
789 },
790 DiffDetails {
791 key: StorageKey(vec![1]),
792 return_type: ArchiveStorageDiffType::Hash,
793 child_trie_key: Some(ChildInfo::new_default_from_vec(vec![2])),
794 child_trie_key_string: Some("0x02".into()),
795 },
796 ],
797 ];
798
799 assert_eq!(result, expected);
800 }
801
802 #[test]
803 fn test_lexicographic_diff() {
804 let left = vec![1, 2, 3, 4, 5];
805 let right = vec![2, 3, 4, 5, 6];
806
807 let diff = lexicographic_diff(left.into_iter(), right.into_iter()).collect::<Vec<_>>();
808 let expected = vec![
809 Diff::Added(1),
810 Diff::Equal(2),
811 Diff::Equal(3),
812 Diff::Equal(4),
813 Diff::Equal(5),
814 Diff::Deleted(6),
815 ];
816 assert_eq!(diff, expected);
817 }
818
819 #[test]
820 fn test_lexicographic_diff_one_side_empty() {
821 let left = vec![];
822 let right = vec![1, 2, 3, 4, 5, 6];
823
824 let diff = lexicographic_diff(left.into_iter(), right.into_iter()).collect::<Vec<_>>();
825 let expected = vec![
826 Diff::Deleted(1),
827 Diff::Deleted(2),
828 Diff::Deleted(3),
829 Diff::Deleted(4),
830 Diff::Deleted(5),
831 Diff::Deleted(6),
832 ];
833 assert_eq!(diff, expected);
834
835 let left = vec![1, 2, 3, 4, 5, 6];
836 let right = vec![];
837
838 let diff = lexicographic_diff(left.into_iter(), right.into_iter()).collect::<Vec<_>>();
839 let expected = vec![
840 Diff::Added(1),
841 Diff::Added(2),
842 Diff::Added(3),
843 Diff::Added(4),
844 Diff::Added(5),
845 Diff::Added(6),
846 ];
847 assert_eq!(diff, expected);
848 }
849}