From 0194da81419d8c0729ac75209ec1fffcfdc0510c Mon Sep 17 00:00:00 2001 From: Daniel Shiposha Date: Fri, 11 Nov 2022 16:01:30 +0000 Subject: [PATCH 1/3] feat: partial_translate for map storages --- .../src/storage/generator/double_map.rs | 174 +++++++++++++++++ frame/support/src/storage/generator/map.rs | 158 +++++++++++++++ frame/support/src/storage/generator/nmap.rs | 183 ++++++++++++++++++ frame/support/src/storage/mod.rs | 75 +++++++ 4 files changed, 590 insertions(+) diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index c95dcee9d7e5c..6f0fe06ab7c23 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -500,6 +500,60 @@ where } } } + + fn partial_translate(mut max: usize, cursor: Option<&[u8]>, mut f: F) -> Option> + where + O: Decode, + F: FnMut(K1, K2, O) -> Option, + { + let prefix = G::prefix_hash(); + let mut previous_key = match cursor { + Some(cursor) => cursor.into(), + None => prefix.clone(), + }; + + while let Some(next) = + sp_io::storage::next_key(&previous_key).filter(|n| max > 0 && n.starts_with(&prefix)) + { + previous_key = next; + let value = match unhashed::get::(&previous_key) { + Some(value) => value, + None => { + log::error!("Invalid partial translate: fail to decode old value"); + continue + }, + }; + + let mut key_material = G::Hasher1::reverse(&previous_key[prefix.len()..]); + let key1 = match K1::decode(&mut key_material) { + Ok(key1) => key1, + Err(_) => { + log::error!("Invalid translate: fail to decode key1"); + continue + }, + }; + + let mut key2_material = G::Hasher2::reverse(key_material); + let key2 = match K2::decode(&mut key2_material) { + Ok(key2) => key2, + Err(_) => { + log::error!("Invalid translate: fail to decode key2"); + continue + }, + }; + + match f(key1, key2, value) { + Some(new) => unhashed::put::(&previous_key, &new), + None => unhashed::kill(&previous_key), + } + + max = max.saturating_sub(1); + } + + sp_io::storage::next_key(&previous_key) + .filter(|n| n.starts_with(&prefix)) + .map(|_| previous_key) + } } /// Test iterators for StorageDoubleMap @@ -679,3 +733,123 @@ mod test_iterators { }) } } + +#[cfg(test)] +mod test_partial_translate { + use crate::storage::{generator::StorageDoubleMap, unhashed, IterableStorageDoubleMap}; + + #[test] + fn partial_translate_works() { + sp_io::TestExternalities::default().execute_with(|| { + use crate::hash::Identity; + #[crate::storage_alias] + type MyMap = StorageDoubleMap; + + let max_entries: usize = 50; + let quarter = max_entries / 4; + let prefix = MyMap::prefix_hash(); + + // Init the storage with `u64` values + for i in 0..max_entries { + let final_key = MyMap::storage_double_map_final_key(i as u64, i as u64); + unhashed::put(&final_key, &(i as u64 * 10)); + } + + { + // The migration shouldn't execute + let cursor = + MyMap::partial_translate(0, None, |_, _, v: u64| Some(v as f64 + 1.0)).unwrap(); + + // Check nothing is migrated + let mut rest_cursor = prefix.clone(); + for i in 0..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // The cursor should point to the initial location since we passed `None` + assert_eq!(cursor, prefix); + } + + { + // Migrate the first quarter + let cursor = + MyMap::partial_translate(quarter, None, |_, _, v: u64| Some(v as f64 + 1.0)) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is migrated + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64, i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The rest is untouched + for i in quarter..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the second quarter, remove even key + let cursor = MyMap::partial_translate(quarter, Some(&cursor), |k1, k2, v: u64| { + ((k1 + k2) % 2 != 0).then(|| v as f64 + 1.0) + }) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64, i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The second quarter is migrated + for i in quarter..(2 * quarter) { + assert_eq!( + MyMap::get(i as u64, i as u64), + ((i + i) % 2 != 0).then(|| (i * 10) as f64 + 1.0) + ); + } + + // The rest is untouched + for i in (2 * quarter)..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the rest using `limit` greater or equal to the number of the rest items. + let cursor = + MyMap::partial_translate(max_entries, Some(&cursor), |_, _, v: u64| { + Some(v as f64 + 1.0) + }); + + assert!(cursor.is_none()); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64, i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The second quarter is untouched + for i in quarter..(2 * quarter) { + assert_eq!( + MyMap::get(i as u64, i as u64), + ((i + i) % 2 != 0).then(|| (i * 10) as f64 + 1.0) + ); + } + + // The rest is migrated + for i in (2 * quarter)..max_entries { + assert_eq!(MyMap::get(i as u64, i as u64), Some((i * 10) as f64 + 1.0)); + } + } + }); + } +} diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs index f6c8eaa270bb3..32105c7414fd1 100644 --- a/frame/support/src/storage/generator/map.rs +++ b/frame/support/src/storage/generator/map.rs @@ -207,6 +207,51 @@ where } } } + + fn partial_translate(mut max: usize, cursor: Option<&[u8]>, mut f: F) -> Option> + where + O: Decode, + F: FnMut(K, O) -> Option, + { + let prefix = G::prefix_hash(); + let mut previous_key = match cursor { + Some(cursor) => cursor.into(), + None => prefix.clone(), + }; + + while let Some(next) = + sp_io::storage::next_key(&previous_key).filter(|n| max > 0 && n.starts_with(&prefix)) + { + previous_key = next; + let value = match unhashed::get::(&previous_key) { + Some(value) => value, + None => { + log::error!("Invalid partial translate: fail to decode old value"); + continue + }, + }; + + let mut key_material = G::Hasher::reverse(&previous_key[prefix.len()..]); + let key = match K::decode(&mut key_material) { + Ok(key) => key, + Err(_) => { + log::error!("Invalid partial translate: fail to decode key"); + continue + }, + }; + + match f(key, value) { + Some(new) => unhashed::put::(&previous_key, &new), + None => unhashed::kill(&previous_key), + } + + max = max.saturating_sub(1); + } + + sp_io::storage::next_key(&previous_key) + .filter(|n| n.starts_with(&prefix)) + .map(|_| previous_key) + } } impl> storage::StorageMap for G { @@ -458,3 +503,116 @@ mod test_iterators { }) } } + +#[cfg(test)] +mod test_partial_translate { + use crate::storage::{generator::StorageMap, unhashed, IterableStorageMap}; + + #[test] + fn partial_translate_works() { + sp_io::TestExternalities::default().execute_with(|| { + use crate::hash::Identity; + #[crate::storage_alias] + type MyMap = StorageMap; + + let max_entries: usize = 50; + let quarter = max_entries / 4; + let prefix = MyMap::prefix_hash(); + + // Init the storage with `u64` values + for i in 0..max_entries { + let final_key = MyMap::storage_map_final_key(i as u64); + unhashed::put(&final_key, &(i as u64 * 10)); + } + + { + // The migration shouldn't execute + let cursor = + MyMap::partial_translate(0, None, |_, v: u64| Some(v as f64 + 1.0)).unwrap(); + + // Check nothing is migrated + let mut rest_cursor = prefix.clone(); + for i in 0..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // The cursor should point to the initial location since we passed `None` + assert_eq!(cursor, prefix); + } + + { + // Migrate the first quarter + let cursor = + MyMap::partial_translate(quarter, None, |_, v: u64| Some(v as f64 + 1.0)) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is migrated + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The rest is untouched + for i in quarter..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the second quarter, remove even key + let cursor = MyMap::partial_translate(quarter, Some(&cursor), |k, v: u64| { + (k % 2 != 0).then(|| v as f64 + 1.0) + }) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The second quarter is migrated + for i in quarter..(2 * quarter) { + assert_eq!(MyMap::get(i as u64), (i % 2 != 0).then(|| (i * 10) as f64 + 1.0)); + } + + // The rest is untouched + for i in (2 * quarter)..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the rest using `limit` greater or equal to the number of the rest items. + let cursor = MyMap::partial_translate(max_entries, Some(&cursor), |_, v: u64| { + Some(v as f64 + 1.0) + }); + + assert!(cursor.is_none()); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!(MyMap::get(i as u64), Some((i * 10) as f64 + 1.0)); + } + + // The second quarter is untouched + for i in quarter..(2 * quarter) { + assert_eq!(MyMap::get(i as u64), (i % 2 != 0).then(|| (i * 10) as f64 + 1.0)); + } + + // The rest is migrated + for i in (2 * quarter)..max_entries { + assert_eq!(MyMap::get(i as u64), Some((i * 10) as f64 + 1.0)); + } + } + }); + } +} diff --git a/frame/support/src/storage/generator/nmap.rs b/frame/support/src/storage/generator/nmap.rs index 79f3d72044e28..6db81254d2abd 100755 --- a/frame/support/src/storage/generator/nmap.rs +++ b/frame/support/src/storage/generator/nmap.rs @@ -448,6 +448,50 @@ impl> } } } + + fn partial_translate(mut max: usize, cursor: Option<&[u8]>, mut f: F) -> Option> + where + O: Decode, + F: FnMut(K::Key, O) -> Option, + { + let prefix = G::prefix_hash(); + let mut previous_key = match cursor { + Some(cursor) => cursor.into(), + None => prefix.clone(), + }; + + while let Some(next) = + sp_io::storage::next_key(&previous_key).filter(|n| max > 0 && n.starts_with(&prefix)) + { + previous_key = next; + let value = match unhashed::get::(&previous_key) { + Some(value) => value, + None => { + log::error!("Invalid partial translate: fail to decode old value"); + continue + }, + }; + + let final_key = match K::decode_final_key(&previous_key[prefix.len()..]) { + Ok((final_key, _)) => final_key, + Err(_) => { + log::error!("Invalid translate: fail to decode key"); + continue + }, + }; + + match f(final_key, value) { + Some(new) => unhashed::put::(&previous_key, &new), + None => unhashed::kill(&previous_key), + } + + max = max.saturating_sub(1); + } + + sp_io::storage::next_key(&previous_key) + .filter(|n| n.starts_with(&prefix)) + .map(|_| previous_key) + } } /// Test iterators for StorageNMap @@ -656,3 +700,142 @@ mod test_iterators { }) } } + +#[cfg(test)] +mod test_partial_translate { + use crate::storage::{generator::StorageNMap, unhashed, IterableStorageNMap, Key as NMapKey}; + + #[test] + fn partial_translate_works() { + sp_io::TestExternalities::default().execute_with(|| { + use crate::hash::Identity; + #[crate::storage_alias] + type MyMap = StorageNMap< + MyModule, + (NMapKey, NMapKey, NMapKey), + f64, + >; + + type Key = (NMapKey, NMapKey, NMapKey); + + let max_entries: usize = 50; + let quarter = max_entries / 4; + let prefix = MyMap::prefix_hash(); + + // Init the storage with `u64` values + for i in 0..max_entries { + let final_key = + MyMap::storage_n_map_final_key::((i as u64, i as u64, i as u64)); + unhashed::put(&final_key, &(i as u64 * 10)); + } + + { + // The migration shouldn't execute + let cursor = + MyMap::partial_translate(0, None, |_, v: u64| Some(v as f64 + 1.0)).unwrap(); + + // Check nothing is migrated + let mut rest_cursor = prefix.clone(); + for i in 0..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // The cursor should point to the initial location since we passed `None` + assert_eq!(cursor, prefix); + } + + { + // Migrate the first quarter + let cursor = + MyMap::partial_translate(quarter, None, |_, v: u64| Some(v as f64 + 1.0)) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is migrated + for i in 0..quarter { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + Some((i * 10) as f64 + 1.0) + ); + } + + // The rest is untouched + for i in quarter..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the second quarter, remove even key + let cursor = + MyMap::partial_translate(quarter, Some(&cursor), |(k1, k2, k3), v: u64| { + ((k1 + k2 + k3) % 2 != 0).then(|| v as f64 + 1.0) + }) + .unwrap(); + let mut rest_cursor = cursor.clone(); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + Some((i * 10) as f64 + 1.0) + ); + } + + // The second quarter is migrated + for i in quarter..(2 * quarter) { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + ((i + i + i) % 2 != 0).then(|| (i * 10) as f64 + 1.0) + ); + } + + // The rest is untouched + for i in (2 * quarter)..max_entries { + let cursor_key = sp_io::storage::next_key(&rest_cursor).unwrap(); + + assert_eq!(unhashed::get(&cursor_key), Some(i as u64 * 10)); + + rest_cursor = cursor_key; + } + + // Migrate the rest using `limit` greater or equal to the number of the rest items. + let cursor = MyMap::partial_translate(max_entries, Some(&cursor), |_, v: u64| { + Some(v as f64 + 1.0) + }); + + assert!(cursor.is_none()); + + // The first quarter is untouched + for i in 0..quarter { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + Some((i * 10) as f64 + 1.0) + ); + } + + // The second quarter is untouched + for i in quarter..(2 * quarter) { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + ((i + i + i) % 2 != 0).then(|| (i * 10) as f64 + 1.0) + ); + } + + // The rest is migrated + for i in (2 * quarter)..max_entries { + assert_eq!( + MyMap::get((i as u64, i as u64, i as u64)), + Some((i * 10) as f64 + 1.0) + ); + } + } + }); + } +} diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 333f4382557b1..bb30e778d06de 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -295,6 +295,31 @@ pub trait IterableStorageMap: StorageMap { /// /// NOTE: If a value fail to decode because storage is corrupted then it is skipped. fn translate Option>(f: F); + + /// Partially translate items from the map. + /// + /// Returns the current cursor. + /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// + /// NOTE: After the initial call, it is important that no further items + /// are inserted into the map. If so, then the map may not be fully translated when the + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// passed once (in the initial call). + /// Subsequent calls operating on the same map + /// should always pass `Some`, and this should be equal to the + /// previous call result's cursor. + fn partial_translate(limit: usize, maybe_cursor: Option<&[u8]>, f: F) -> Option> + where + O: Decode, + F: FnMut(K, O) -> Option; } /// A strongly-typed double map in storage whose secondary keys and values can be iterated over. @@ -370,6 +395,31 @@ pub trait IterableStorageDoubleMap: /// /// NOTE: If a value fail to decode because storage is corrupted then it is skipped. fn translate Option>(f: F); + + /// Partially translate items from the map. + /// + /// Returns the current cursor. + /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// + /// NOTE: After the initial call, it is important that no further items + /// are inserted into the map. If so, then the map may not be fully translated when the + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// passed once (in the initial call). + /// Subsequent calls operating on the same map + /// should always pass `Some`, and this should be equal to the + /// previous call result's cursor. + fn partial_translate(limit: usize, maybe_cursor: Option<&[u8]>, f: F) -> Option> + where + O: Decode, + F: FnMut(K1, K2, O) -> Option; } /// A strongly-typed map with arbitrary number of keys in storage whose keys and values can be @@ -450,6 +500,31 @@ pub trait IterableStorageNMap: StorageN /// /// NOTE: If a value fail to decode because storage is corrupted then it is skipped. fn translate Option>(f: F); + + /// Partially translate items from the map. + /// + /// Returns the current cursor. + /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// + /// NOTE: After the initial call, it is important that no further items + /// are inserted into the map. If so, then the map may not be fully translated when the + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// passed once (in the initial call). + /// Subsequent calls operating on the same map + /// should always pass `Some`, and this should be equal to the + /// previous call result's cursor. + fn partial_translate(limit: usize, maybe_cursor: Option<&[u8]>, f: F) -> Option> + where + O: Decode, + F: FnMut(K::Key, O) -> Option; } /// An implementation of a map with a two keys. From 953b24cfd1638f8358973a17833a3e3eac67b04a Mon Sep 17 00:00:00 2001 From: Daniel Shiposha Date: Mon, 14 Nov 2022 12:16:22 +0000 Subject: [PATCH 2/3] fix: partial_translate doc comment --- frame/support/src/storage/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index bb30e778d06de..d405b8be47555 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -311,7 +311,7 @@ pub trait IterableStorageMap: StorageMap { /// /// # Cursor /// - /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// A *cursor* may be passed into this operation with `maybe_cursor`. `None` should only be /// passed once (in the initial call). /// Subsequent calls operating on the same map /// should always pass `Some`, and this should be equal to the @@ -411,7 +411,7 @@ pub trait IterableStorageDoubleMap: /// /// # Cursor /// - /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// A *cursor* may be passed into this operation with `maybe_cursor`. `None` should only be /// passed once (in the initial call). /// Subsequent calls operating on the same map /// should always pass `Some`, and this should be equal to the @@ -516,7 +516,7 @@ pub trait IterableStorageNMap: StorageN /// /// # Cursor /// - /// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be + /// A *cursor* may be passed into this operation with `maybe_cursor`. `None` should only be /// passed once (in the initial call). /// Subsequent calls operating on the same map /// should always pass `Some`, and this should be equal to the From 515e1a95cd846017f03872bf38002eb755eefbde Mon Sep 17 00:00:00 2001 From: Daniel Shiposha Date: Mon, 14 Nov 2022 12:23:59 +0000 Subject: [PATCH 3/3] fix: partial_translate doc comment --- frame/support/src/storage/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index d405b8be47555..7194195da55f7 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -299,7 +299,7 @@ pub trait IterableStorageMap: StorageMap { /// Partially translate items from the map. /// /// Returns the current cursor. - /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// Once the resultant cursor is `None`, then no further items remain to be translated. /// /// NOTE: After the initial call, it is important that no further items /// are inserted into the map. If so, then the map may not be fully translated when the @@ -399,7 +399,7 @@ pub trait IterableStorageDoubleMap: /// Partially translate items from the map. /// /// Returns the current cursor. - /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// Once the resultant cursor is `None`, then no further items remain to be translated. /// /// NOTE: After the initial call, it is important that no further items /// are inserted into the map. If so, then the map may not be fully translated when the @@ -504,7 +504,7 @@ pub trait IterableStorageNMap: StorageN /// Partially translate items from the map. /// /// Returns the current cursor. - /// Once the resultant cursor is `None`, then no further items remain to be deleted. + /// Once the resultant cursor is `None`, then no further items remain to be translated. /// /// NOTE: After the initial call, it is important that no further items /// are inserted into the map. If so, then the map may not be fully translated when the