State Trie Migration v1 (#2907)

* init * migration on idle * rustfmt * add a comment * change state_version * fix visibility * count get_next_key * Update runtime/moonbase/src/lib.rs Co-authored-by: Rodrigo Quelhas <22591718+RomarQ@users.noreply.github.com> * migrate code * fix * add pre/post upgrades * rename * disable * Revert "disable" This reverts commit 24e2eca. * change to state_version 0 * Revert "change to state_version 0" This reverts commit dffb686. * Update pallets/moonbeam-lazy-migrations/src/lib.rs Co-authored-by: Éloïs <c@elo.tf> * apply review suggestions * fix pov in tests * fmt & fix pov values * fix numbers * update & move zombie setup script into test/scripts * use branch build for RPC tests --------- Co-authored-by: Rodrigo Quelhas <22591718+RomarQ@users.noreply.github.com> Co-authored-by: Éloïs <c@elo.tf> Co-authored-by: noandrea <no.andrea@gmail.com> Co-authored-by: Pablo Labarta <pablitolabarta@gmail.com>
moonbeam-foundation · Oct 1, 2024 · 4c20e1e · 4c20e1e
1 parent 9854dad
commit 4c20e1e
Show file tree

Hide file tree

Showing 24 changed files with 759 additions and 49 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -295,7 +295,6 @@ cumulus-primitives-parachain-inherent = { git = "https://github.com/moonbeam-fou
 cumulus-primitives-timestamp = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
 cumulus-primitives-utility = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
 cumulus-test-relay-sproof-builder = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
-cumulus-primitives-storage-weight-reclaim = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
 parachain-info = { package = "staging-parachain-info", git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
 parachains-common = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false }
 
@@ -312,6 +311,7 @@ cumulus-relay-chain-inprocess-interface = { git = "https://github.com/moonbeam-f
 cumulus-relay-chain-interface = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0" }
 cumulus-relay-chain-minimal-node = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0" }
 cumulus-relay-chain-rpc-interface = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0" }
+cumulus-primitives-storage-weight-reclaim = { git = "https://github.com/moonbeam-foundation/polkadot-sdk", branch = "moonbeam-polkadot-v1.11.0", default-features = false}
 
 # Polkadot / XCM (wasm)
 orml-traits = { git = "https://github.com/moonbeam-foundation/open-runtime-module-library", branch = "moonbeam-polkadot-v1.11.0", default-features = false }

diff --git a/pallets/moonbeam-lazy-migrations/Cargo.toml b/pallets/moonbeam-lazy-migrations/Cargo.toml
@@ -1,37 +1,41 @@
 [package]
-authors = {workspace = true}
+authors = { workspace = true }
 description = "A pallet for performing migrations from extrinsics"
 edition = "2021"
 name = "pallet-moonbeam-lazy-migrations"
 version = "0.1.0"
 
 [dependencies]
-log = {workspace = true}
+log = { workspace = true }
 
 # Substrate
-frame-support = {workspace = true}
-frame-system = {workspace = true}
-pallet-scheduler = {workspace = true}
-pallet-balances = {workspace = true}
-parity-scale-codec = {workspace = true}
-scale-info = {workspace = true, features = ["derive"]}
-sp-core = {workspace = true}
-sp-io = {workspace = true}
-sp-runtime = {workspace = true}
-sp-std = {workspace = true}
+frame-support = { workspace = true }
+frame-system = { workspace = true }
+pallet-scheduler = { workspace = true }
+pallet-assets = { workspace = true }
+pallet-balances = { workspace = true }
+parity-scale-codec = { workspace = true }
+scale-info = { workspace = true, features = ["derive"] }
+sp-core = { workspace = true }
+sp-io = { workspace = true }
+sp-runtime = { workspace = true }
+sp-std = { workspace = true }
 
 # Frontier
-pallet-evm = {workspace = true, features = ["forbid-evm-reentrancy"]}
+pallet-evm = { workspace = true, features = ["forbid-evm-reentrancy"] }
+
+# Runtime Interfaces
+cumulus-primitives-storage-weight-reclaim = { workspace = true, default-features = false }
 
 # Benchmarks
-frame-benchmarking = {workspace = true, optional = true}
+frame-benchmarking = { workspace = true, optional = true }
 
 [dev-dependencies]
-frame-benchmarking = {workspace = true, features = ["std"]}
-pallet-balances = {workspace = true, features = ["std", "insecure_zero_ed"]}
-pallet-timestamp = {workspace = true, features = ["std"]}
-rlp = {workspace = true, features = ["std"]}
-sp-io = {workspace = true, features = ["std"]}
+frame-benchmarking = { workspace = true, features = ["std"] }
+pallet-balances = { workspace = true, features = ["std", "insecure_zero_ed"] }
+pallet-timestamp = { workspace = true, features = ["std"] }
+rlp = { workspace = true, features = ["std"] }
+sp-io = { workspace = true, features = ["std"] }
 
 [features]
 default = ["std"]
@@ -47,6 +51,8 @@ std = [
   "sp-std/std",
   "pallet-evm/std",
   "pallet-timestamp/std",
+  "pallet-assets/std",
+  "cumulus-primitives-storage-weight-reclaim/std",
   "rlp/std",
 ]
-try-runtime = ["frame-support/try-runtime"]
+try-runtime = ["frame-support/try-runtime"]
diff --git a/pallets/moonbeam-lazy-migrations/src/lib.rs b/pallets/moonbeam-lazy-migrations/src/lib.rs
@@ -38,6 +38,7 @@ const MAX_CONTRACT_CODE_SIZE: u64 = 25 * 1024;
 #[pallet]
 pub mod pallet {
 	use super::*;
+	use cumulus_primitives_storage_weight_reclaim::get_proof_size;
 	use frame_support::pallet_prelude::*;
 	use frame_system::pallet_prelude::*;
 	use sp_core::H160;
@@ -53,6 +54,26 @@ pub mod pallet {
 	/// The total number of suicided contracts that were removed
 	pub(crate) type SuicidedContractsRemoved<T: Config> = StorageValue<_, u32, ValueQuery>;
 
+	#[pallet::storage]
+	pub(crate) type StateMigrationStatusValue<T: Config> =
+		StorageValue<_, (StateMigrationStatus, u64), ValueQuery>;
+
+	pub(crate) type StorageKey = BoundedVec<u8, ConstU32<1_024>>;
+
+	#[derive(Clone, Encode, Decode, scale_info::TypeInfo, PartialEq, Eq, MaxEncodedLen, Debug)]
+	pub enum StateMigrationStatus {
+		NotStarted,
+		Started(StorageKey),
+		Error(BoundedVec<u8, ConstU32<1024>>),
+		Complete,
+	}
+
+	impl Default for StateMigrationStatus {
+		fn default() -> Self {
+			return StateMigrationStatus::NotStarted;
+		}
+	}
+
 	/// Configuration trait of this pallet.
 	#[pallet::config]
 	pub trait Config: frame_system::Config + pallet_evm::Config + pallet_balances::Config {
@@ -71,6 +92,233 @@ pub mod pallet {
 		ContractMetadataAlreadySet,
 		/// Contract not exist
 		ContractNotExist,
+		/// The key lengths exceeds the maximum allowed
+		KeyTooLong,
+	}
+
+	pub(crate) const MAX_ITEM_PROOF_SIZE: u64 = 30 * 1024; // 30 KB
+	pub(crate) const PROOF_SIZE_BUFFER: u64 = 100 * 1024; // 100 KB
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn on_idle(_n: BlockNumberFor<T>, remaining_weight: Weight) -> Weight {
+			let proof_size_before: u64 = get_proof_size().unwrap_or(0);
+			let res = Pallet::<T>::handle_migration(remaining_weight);
+			let proof_size_after: u64 = get_proof_size().unwrap_or(0);
+			let proof_size_diff = proof_size_after.saturating_sub(proof_size_before);
+
+			Weight::from_parts(0, proof_size_diff)
+				.saturating_add(T::DbWeight::get().reads_writes(res.reads, res.writes))
+		}
+	}
+
+	#[derive(Default, Clone, PartialEq, Eq, Encode, Decode, Debug)]
+	pub(crate) struct ReadWriteOps {
+		pub reads: u64,
+		pub writes: u64,
+	}
+
+	impl ReadWriteOps {
+		pub fn new() -> Self {
+			Self {
+				reads: 0,
+				writes: 0,
+			}
+		}
+
+		pub fn add_one_read(&mut self) {
+			self.reads += 1;
+		}
+
+		pub fn add_one_write(&mut self) {
+			self.writes += 1;
+		}
+
+		pub fn add_reads(&mut self, reads: u64) {
+			self.reads += reads;
+		}
+
+		pub fn add_writes(&mut self, writes: u64) {
+			self.writes += writes;
+		}
+	}
+
+	#[derive(Clone)]
+	struct StateMigrationResult {
+		last_key: Option<StorageKey>,
+		error: Option<&'static str>,
+		migrated: u64,
+		reads: u64,
+		writes: u64,
+	}
+
+	enum NextKeyResult {
+		NextKey(StorageKey),
+		NoMoreKeys,
+		Error(&'static str),
+	}
+
+	impl<T: Config> Pallet<T> {
+		/// Handle the migration of the storage keys, returns the number of read and write operations
+		pub(crate) fn handle_migration(remaining_weight: Weight) -> ReadWriteOps {
+			let mut read_write_ops = ReadWriteOps::new();
+
+			// maximum number of items that can be migrated in one block
+			let migration_limit = remaining_weight
+				.proof_size()
+				.saturating_sub(PROOF_SIZE_BUFFER)
+				.saturating_div(MAX_ITEM_PROOF_SIZE);
+
+			if migration_limit == 0 {
+				return read_write_ops;
+			}
+
+			let (status, mut migrated_keys) = StateMigrationStatusValue::<T>::get();
+			read_write_ops.add_one_read();
+
+			let next_key = match &status {
+				StateMigrationStatus::NotStarted => Default::default(),
+				StateMigrationStatus::Started(storage_key) => {
+					let (reads, next_key_result) = Pallet::<T>::get_next_key(storage_key);
+					read_write_ops.add_reads(reads);
+					match next_key_result {
+						NextKeyResult::NextKey(next_key) => next_key,
+						NextKeyResult::NoMoreKeys => {
+							StateMigrationStatusValue::<T>::put((
+								StateMigrationStatus::Complete,
+								migrated_keys,
+							));
+							read_write_ops.add_one_write();
+							return read_write_ops;
+						}
+						NextKeyResult::Error(e) => {
+							StateMigrationStatusValue::<T>::put((
+								StateMigrationStatus::Error(
+									e.as_bytes().to_vec().try_into().unwrap_or_default(),
+								),
+								migrated_keys,
+							));
+							read_write_ops.add_one_write();
+							return read_write_ops;
+						}
+					}
+				}
+				StateMigrationStatus::Complete | StateMigrationStatus::Error(_) => {
+					return read_write_ops;
+				}
+			};
+
+			let res = Pallet::<T>::migrate_keys(next_key, migration_limit);
+			migrated_keys += res.migrated;
+			read_write_ops.add_reads(res.reads);
+			read_write_ops.add_writes(res.writes);
+
+			match (res.last_key, res.error) {
+				(None, None) => {
+					StateMigrationStatusValue::<T>::put((
+						StateMigrationStatus::Complete,
+						migrated_keys,
+					));
+					read_write_ops.add_one_write();
+				}
+				// maybe we should store the previous key in the storage as well
+				(_, Some(e)) => {
+					StateMigrationStatusValue::<T>::put((
+						StateMigrationStatus::Error(
+							e.as_bytes().to_vec().try_into().unwrap_or_default(),
+						),
+						migrated_keys,
+					));
+					read_write_ops.add_one_write();
+				}
+				(Some(key), None) => {
+					StateMigrationStatusValue::<T>::put((
+						StateMigrationStatus::Started(key),
+						migrated_keys,
+					));
+					read_write_ops.add_one_write();
+				}
+			}
+
+			read_write_ops
+		}
+
+		/// Tries to get the next key in the storage, returns None if there are no more keys to migrate.
+		/// Returns an error if the key is too long.
+		fn get_next_key(key: &StorageKey) -> (u64, NextKeyResult) {
+			if let Some(next) = sp_io::storage::next_key(key) {
+				let next: Result<StorageKey, _> = next.try_into();
+				match next {
+					Ok(next_key) => {
+						if next_key.as_slice() == sp_core::storage::well_known_keys::CODE {
+							let (reads, next_key_res) = Pallet::<T>::get_next_key(&next_key);
+							return (1 + reads, next_key_res);
+						}
+						(1, NextKeyResult::NextKey(next_key))
+					}
+					Err(_) => (1, NextKeyResult::Error("Key too long")),
+				}
+			} else {
+				(1, NextKeyResult::NoMoreKeys)
+			}
+		}
+
+		/// Migrate maximum of `limit` keys starting from `start`, returns the next key to migrate
+		/// Returns None if there are no more keys to migrate.
+		/// Returns an error if an error occurred during migration.
+		fn migrate_keys(start: StorageKey, limit: u64) -> StateMigrationResult {
+			let mut key = start;
+			let mut migrated = 0;
+			let mut next_key_reads = 0;
+			let mut writes = 0;
+
+			while migrated < limit {
+				let data = sp_io::storage::get(&key);
+				if let Some(data) = data {
+					sp_io::storage::set(&key, &data);
+					writes += 1;
+				}
+
+				migrated += 1;
+
+				if migrated < limit {
+					let (reads, next_key_res) = Pallet::<T>::get_next_key(&key);
+					next_key_reads += reads;
+
+					match next_key_res {
+						NextKeyResult::NextKey(next_key) => {
+							key = next_key;
+						}
+						NextKeyResult::NoMoreKeys => {
+							return StateMigrationResult {
+								last_key: None,
+								error: None,
+								migrated,
+								reads: migrated + next_key_reads,
+								writes,
+							};
+						}
+						NextKeyResult::Error(e) => {
+							return StateMigrationResult {
+								last_key: Some(key),
+								error: Some(e),
+								migrated,
+								reads: migrated + next_key_reads,
+								writes,
+							};
+						}
+					};
+				}
+			}
+
+			StateMigrationResult {
+				last_key: Some(key),
+				error: None,
+				migrated,
+				reads: migrated + next_key_reads,
+				writes,
+			}
+		}
 	}
 
 	#[pallet::call]