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..7194195da55f7 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 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 + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// 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 + /// 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 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 + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// 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 + /// 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 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 + /// resultant `maybe_cursor` is `None`. + /// + /// # Limit + /// + /// Maximum amount of items to migrate. + /// + /// # Cursor + /// + /// 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 + /// 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.