State Trie Migration v1

Cargo.toml

@@ -305,6 +305,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 }

pallets/moonbeam-lazy-migrations/Cargo.toml

@@ -1,38 +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-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}
+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"]
@@ -48,6 +51,7 @@ std = [
   "sp-std/std",
   "pallet-evm/std",
   "pallet-timestamp/std",
+  "cumulus-primitives-storage-weight-reclaim/std",
   "rlp/std",
 ]
-try-runtime = ["frame-support/try-runtime"]
+try-runtime = ["frame-support/try-runtime"]

pallets/moonbeam-lazy-migrations/src/lib.rs

@@ -36,6 +36,7 @@ pub use pallet::*;
 #[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;
@@ -51,6 +52,26 @@ pub mod pallet {
 	/// The total number of suicided contracts that were removed
 	pub(crate) type SuicidedContractsRemoved<T: Config> = StorageValue<_, u32, ValueQuery>;
 
+	#[pallet::storage]
+	/// The total number of suicided contracts that were removed
+	pub(crate) type StateMigrationStatus<T: Config> = StorageValue<_, MigrationStatus, ValueQuery>;
+
+	pub(crate) type StorageKey = BoundedVec<u8, ConstU32<512>>;
+
+	#[derive(Clone, Encode, Decode, scale_info::TypeInfo, PartialEq, Eq, MaxEncodedLen)]
+	pub enum MigrationStatus {
+		NotStarted,
+		Started(StorageKey),
+		Error(BoundedVec<u8, ConstU32<1024>>),
+		Complete,
+	}
+
+	impl Default for MigrationStatus {
+		fn default() -> Self {
+			return MigrationStatus::NotStarted;
+		}
+	}
+
 	/// Configuration trait of this pallet.
 	#[pallet::config]
 	pub trait Config: frame_system::Config + pallet_evm::Config + pallet_balances::Config {
@@ -65,6 +86,209 @@ pub mod pallet {
 		AddressesLengthCannotBeZero,
 		/// The contract is not corrupted (Still exist or properly suicided)
 		ContractNotCorrupted,
+		/// The key lengths exceeds the maximum allowed
+		KeyTooLong,
+	}
+
+	const MAX_ITEM_PROOF_SIZE: u64 = 30 * 1024; // 30 KB
+	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)
+				// For now the DbWeight is only recording the ref_time and not account for
+				// the proof_size.
+				.saturating_add(T::DbWeight::get().reads_writes(res.reads, res.writes))
+		}
+	}
+
+	#[derive(Default, Clone)]
+	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 MigrationResult {
+		last_key: Option<StorageKey>,
+		error: Option<&'static str>,
+		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
+		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);
+
+			let status = StateMigrationStatus::<T>::get();
+			read_write_ops.add_one_read();
+
+			let next_key = match &status {
+				MigrationStatus::NotStarted => Default::default(),
+				MigrationStatus::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 => {
+							StateMigrationStatus::<T>::put(MigrationStatus::Complete);
+							read_write_ops.add_one_write();
+							return read_write_ops;
+						}
+						NextKeyResult::Error(e) => {
+							StateMigrationStatus::<T>::put(MigrationStatus::Error(
+								e.as_bytes().to_vec().try_into().unwrap_or_default(),
+							));
+							read_write_ops.add_one_write();
+							return read_write_ops;
+						}
+					}
+				}
+				MigrationStatus::Complete | MigrationStatus::Error(_) => {
+					return read_write_ops;
+				}
+			};
+
+			let res = Pallet::<T>::migrate_keys(next_key, migration_limit);
+			read_write_ops.add_reads(res.reads);
+			read_write_ops.add_writes(res.writes);
+
+			match (res.last_key, res.error) {
+				(None, None) => {
+					StateMigrationStatus::<T>::put(MigrationStatus::Complete);
+					read_write_ops.add_one_write();
+				}
+				// maybe we should store the previous key in the storage as well
+				(_, Some(e)) => {
+					StateMigrationStatus::<T>::put(MigrationStatus::Error(
+						e.as_bytes().to_vec().try_into().unwrap_or_default(),
+					));
+					read_write_ops.add_one_write();
+				}
+				(Some(key), None) => {
+					StateMigrationStatus::<T>::put(MigrationStatus::Started(key));
+					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(key) => {
+						if key.as_slice() == sp_core::storage::well_known_keys::CODE {
+							let (reads, next_key_res) = Pallet::<T>::get_next_key(&key);
+							return (1 + reads, next_key_res);
+						}
+						(1, NextKeyResult::NextKey(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
+		/// if there are more keys 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) -> MigrationResult {
+			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;
+
+				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 MigrationResult {
+							last_key: None,
+							error: None,
+							reads: migrated,
+							writes,
+						};
+					}
+					NextKeyResult::Error(e) => {
+						return MigrationResult {
+							last_key: Some(key),
+							error: Some(e),
+							reads: migrated,
+							writes,
+						};
+					}
+				};
+			}
+
+			MigrationResult {
+				last_key: Some(key),
+				error: None,
+				reads: migrated + next_key_reads,
+				writes,
+			}
+		}
 	}
 
 	#[pallet::call]

runtime/common/Cargo.toml

@@ -56,6 +56,7 @@ pallet-xcm = { workspace = true }
 sp-api = { workspace = true }
 sp-consensus-slots = { workspace = true }
 sp-core = { workspace = true }
+sp-io = { workspace = true }
 sp-runtime = { workspace = true }
 sp-std = { workspace = true }
 sp-genesis-builder = { workspace = true }
@@ -85,6 +86,9 @@ parity-scale-codec = { workspace = true }
 
 account = { workspace = true }
 
+# Runtime Interfaces
+cumulus-primitives-storage-weight-reclaim = { workspace = true, default-features = false }
+
 [features]
 std = [
 	"cumulus-pallet-parachain-system/std",
@@ -111,9 +115,11 @@ std = [
 	"precompile-utils/std",
 	"sp-consensus-slots/std",
 	"sp-core/std",
+	"sp-io/std",
 	"sp-runtime/std",
 	"sp-std/std",
 	"sp-genesis-builder/std",
+	"cumulus-primitives-storage-weight-reclaim/std",
 	"xcm-executor/std",
 	"xcm-fee-payment-runtime-api/std",
 	"xcm/std",