Refactor session away from needless double_maps (paritytech#5202)

* Meat and bones.

* Fix migration

* Update frame/balances/src/migration.rs

Co-Authored-By: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

Co-authored-by: Kian Paimani <5588131+kianenigma@users.noreply.github.com>

frame/balances/src/migration.rs

@@ -1,3 +1,21 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Substrate.
+
+// Substrate is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Substrate is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Temporary migrations of the balances module.
+
 use super::*;
 
 pub fn on_runtime_upgrade<T: Trait<I>, I: Instance>() {

frame/session/src/lib.rs

@@ -342,8 +342,6 @@ pub trait Trait: frame_system::Trait {
 	type DisabledValidatorsThreshold: Get<Perbill>;
 }
 
-const DEDUP_KEY_PREFIX: &[u8] = b":session:keys";
-
 decl_storage! {
 	trait Store for Module<T: Trait> as Session {
 		/// The current set of validators.
@@ -366,20 +364,10 @@ decl_storage! {
 		DisabledValidators get(fn disabled_validators): Vec<u32>;
 
 		/// The next session keys for a validator.
-		///
-		/// The first key is always `DEDUP_KEY_PREFIX` to have all the data in the same branch of
-		/// the trie. Having all data in the same branch should prevent slowing down other queries.
-		// TODO: Migrate to a normal map now https://github.com/paritytech/substrate/issues/4917
-		NextKeys: double_map hasher(twox_64_concat) Vec<u8>, hasher(blake2_256) T::ValidatorId
-			=> Option<T::Keys>;
+		NextKeys: map hasher(blake2_256) T::ValidatorId => Option<T::Keys>;
 
-		/// The owner of a key. The second key is the `KeyTypeId` + the encoded key.
-		///
-		/// The first key is always `DEDUP_KEY_PREFIX` to have all the data in the same branch of
-		/// the trie. Having all data in the same branch should prevent slowing down other queries.
-		// TODO: Migrate to a normal map now https://github.com/paritytech/substrate/issues/4917
-		KeyOwner: double_map hasher(twox_64_concat) Vec<u8>, hasher(blake2_256) (KeyTypeId, Vec<u8>)
-			=> Option<T::ValidatorId>;
+		/// The owner of a key. The key is the `KeyTypeId` + the encoded key.
+		KeyOwner: map hasher(blake2_256) (KeyTypeId, Vec<u8>) => Option<T::ValidatorId>;
 	}
 	add_extra_genesis {
 		config(keys): Vec<(T::AccountId, T::ValidatorId, T::Keys)>;
@@ -460,10 +448,6 @@ decl_error! {
 
 decl_module! {
 	pub struct Module<T: Trait> for enum Call where origin: T::Origin {
-		/// Used as first key for `NextKeys` and `KeyOwner` to put all the data into the same branch
-		/// of the trie.
-		const DEDUP_KEY_PREFIX: &[u8] = DEDUP_KEY_PREFIX;
-
 		type Error = Error<T>;
 
 		fn deposit_event() = default;
@@ -514,10 +498,36 @@ decl_module! {
 				Self::rotate_session();
 			}
 		}
+
+		/// Called when the runtime is upgraded.
+		fn on_runtime_upgrade() {
+			Self::migrate();
+		}
 	}
 }
 
 impl<T: Trait> Module<T> {
+	/// Move keys from NextKeys and KeyOwner, if any exist.
+	fn migrate() {
+		use frame_support::storage::migration::{put_storage_value, StorageIterator};
+		sp_runtime::print("Migrating session's double-maps...");
+
+		let prefix = {
+			const DEDUP_KEY_PREFIX: &[u8] = b":session:keys";
+			let encoded_prefix_key_hash = codec::Encode::encode(&DEDUP_KEY_PREFIX);
+			let mut h = sp_io::hashing::twox_64(&encoded_prefix_key_hash[..]).to_vec();
+			h.extend(&encoded_prefix_key_hash[..]);
+			h
+		};
+
+		for (hash, value) in StorageIterator::<T::Keys>::with_suffix(b"Session", b"NextKeys", &prefix[..]).drain() {
+			put_storage_value(b"Session", b"NextKeys", &hash, value);
+		}
+		for (hash, value) in StorageIterator::<T::ValidatorId>::with_suffix(b"Session", b"KeyOwner", &prefix[..]).drain() {
+			put_storage_value(b"Session", b"KeyOwner", &hash, value);
+		}
+	}
+
 	/// Move on to next session. Register new validator set and session keys. Changes
 	/// to the validator set have a session of delay to take effect. This allows for
 	/// equivocation punishment after a fork.
@@ -704,27 +714,27 @@ impl<T: Trait> Module<T> {
 	}
 
 	fn load_keys(v: &T::ValidatorId) -> Option<T::Keys> {
-		<NextKeys<T>>::get(DEDUP_KEY_PREFIX, v)
+		<NextKeys<T>>::get(v)
 	}
 
 	fn take_keys(v: &T::ValidatorId) -> Option<T::Keys> {
-		<NextKeys<T>>::take(DEDUP_KEY_PREFIX, v)
+		<NextKeys<T>>::take(v)
 	}
 
 	fn put_keys(v: &T::ValidatorId, keys: &T::Keys) {
-		<NextKeys<T>>::insert(DEDUP_KEY_PREFIX, v, keys);
+		<NextKeys<T>>::insert(v, keys);
 	}
 
 	fn key_owner(id: KeyTypeId, key_data: &[u8]) -> Option<T::ValidatorId> {
-		<KeyOwner<T>>::get(DEDUP_KEY_PREFIX, (id, key_data))
+		<KeyOwner<T>>::get((id, key_data))
 	}
 
 	fn put_key_owner(id: KeyTypeId, key_data: &[u8], v: &T::ValidatorId) {
-		<KeyOwner<T>>::insert(DEDUP_KEY_PREFIX, (id, key_data), v)
+		<KeyOwner<T>>::insert((id, key_data), v)
 	}
 
 	fn clear_key_owner(id: KeyTypeId, key_data: &[u8]) {
-		<KeyOwner<T>>::remove(DEDUP_KEY_PREFIX, (id, key_data));
+		<KeyOwner<T>>::remove((id, key_data));
 	}
 }
 

frame/support/src/storage/migration.rs

@@ -22,7 +22,7 @@ use crate::{StorageHasher, Twox128};
 
 /// Utility to iterate through raw items in storage.
 pub struct StorageIterator<T> {
-	prefix: [u8; 32],
+	prefix: Vec<u8>,
 	previous_key: Vec<u8>,
 	drain: bool,
 	_phantom: ::sp_std::marker::PhantomData<T>,
@@ -31,11 +31,19 @@ pub struct StorageIterator<T> {
 impl<T> StorageIterator<T> {
 	/// Construct iterator to iterate over map items in `module` for the map called `item`.
 	pub fn new(module: &[u8], item: &[u8]) -> Self {
-		let mut prefix = [0u8; 32];
-		prefix[0..16].copy_from_slice(&Twox128::hash(module));
-		prefix[16..32].copy_from_slice(&Twox128::hash(item));
-		Self { prefix, previous_key: prefix[..].to_vec(), drain: false, _phantom: Default::default() }
+		Self::with_suffix(module, item, &[][..])
 	}
+
+	/// Construct iterator to iterate over map items in `module` for the map called `item`.
+	pub fn with_suffix(module: &[u8], item: &[u8], suffix: &[u8]) -> Self {
+		let mut prefix = Vec::new();
+		prefix.extend_from_slice(&Twox128::hash(module));
+		prefix.extend_from_slice(&Twox128::hash(item));
+		prefix.extend_from_slice(suffix);
+		let previous_key = prefix.clone();
+		Self { prefix, previous_key, drain: false, _phantom: Default::default() }
+	}
+
 	/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
 	pub fn drain(mut self) -> Self {
 		self.drain = true;
@@ -59,7 +67,7 @@ impl<T: Decode + Sized> Iterator for StorageIterator<T> {
 							if self.drain {
 								frame_support::storage::unhashed::kill(&next);
 							}
-							Some((self.previous_key[32..].to_vec(), value))
+							Some((self.previous_key[self.prefix.len()..].to_vec(), value))
 						}
 						None => continue,
 					}