fix: undecodable ForeignInvestmentInfo on Demo (#1819)

* feat: add storage key iteratos to migration utils

* fix: undecodable FI storage entry on demo

* fmt: taplo

* fix: clippy

pallets/foreign-investments/src/lib.rs

@@ -228,7 +228,7 @@ pub mod pallet {
 	/// NOTE: The storage is killed once the investment is fully collected, or
 	/// decreased.
 	#[pallet::storage]
-	pub(super) type ForeignInvestmentInfo<T: Config> = StorageDoubleMap<
+	pub type ForeignInvestmentInfo<T: Config> = StorageDoubleMap<
 		_,
 		Blake2_128Concat,
 		T::AccountId,

runtime/common/src/migrations/mod.rs

@@ -15,3 +15,58 @@
 pub mod increase_storage_version;
 pub mod nuke;
 pub mod precompile_account_codes;
+
+pub mod utils {
+	use frame_support::storage::unhashed;
+	use sp_arithmetic::traits::Saturating;
+
+	/// Iterates keys of storage and removes undecodable keys
+	///
+	/// WARNING: USE WITH EXTREME CAUTION! Ensure the cleanup can be performed
+	/// beforehand!
+	pub fn remove_undecodable_storage_keys<T: parity_scale_codec::Decode + Sized>(
+		prefix: [u8; 32],
+	) -> (u64, u64) {
+		let mut previous_key = prefix.clone().to_vec();
+		let mut reads: u64 = 1;
+		let mut writes: u64 = 0;
+		while let Some(next) =
+			sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix))
+		{
+			reads.saturating_accrue(1);
+			previous_key = next;
+			let maybe_value = unhashed::get::<T>(&previous_key);
+			match maybe_value {
+				Some(_) => continue,
+				None => {
+					log::debug!(
+						"Removing key which could not be decoded at {:?}",
+						previous_key
+					);
+					unhashed::kill(&previous_key);
+					writes.saturating_accrue(1);
+					continue;
+				}
+			}
+		}
+
+		(reads, writes)
+	}
+
+	/// Iterates all keys of a storage and returns the count.
+	///
+	/// NOTE: Necessary because `Storage::<T>::iter().count()` does not account
+	/// for keys which cannot be decoded.
+	pub fn count_storage_keys(prefix: [u8; 32]) -> u32 {
+		let mut previous_key = prefix.clone().to_vec();
+		let mut n: u32 = 0;
+		while let Some(next) =
+			sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix))
+		{
+			n.saturating_accrue(1);
+			previous_key = next;
+		}
+
+		n
+	}
+}

runtime/development/Cargo.toml

@@ -13,6 +13,7 @@ documentation.workspace = true
 getrandom = { workspace = true }
 hex = { workspace = true }
 hex-literal = { workspace = true, optional = true }
+log = { workspace = true }
 parity-scale-codec = { workspace = true }
 scale-info = { workspace = true }
 serde = { workspace = true, optional = true }
@@ -148,9 +149,9 @@ std = [
   "parity-scale-codec/std",
   "getrandom/std",
   "hex/std",
+  "log/std",
   "scale-info/std",
   "serde/std",
-
   # Substrate related
   "sp-api/std",
   "sp-runtime/std",
@@ -187,14 +188,12 @@ std = [
   "pallet-transaction-payment-rpc-runtime-api/std",
   "polkadot-runtime-common/std",
   "polkadot-parachain-primitives/std",
-
   # Locals
   "cfg-primitives/std",
   "cfg-traits/std",
   "cfg-types/std",
   "runtime-common/std",
   "liquidity-pools-gateway-routers/std",
-
   # Pallet list
   "axelar-gateway-precompile/std",
   "chainbridge/std",
@@ -273,7 +272,6 @@ runtime-benchmarks = [
   "frame-system-benchmarking/runtime-benchmarks",
   "frame-benchmarking/runtime-benchmarks",
   "cumulus-pallet-session-benchmarking/runtime-benchmarks",
-
   # Substrate related
   "sp-runtime/runtime-benchmarks",
   "sp-staking/runtime-benchmarks",
@@ -284,14 +282,12 @@ runtime-benchmarks = [
   "xcm-primitives/runtime-benchmarks",
   "polkadot-runtime-common/runtime-benchmarks",
   "polkadot-parachain-primitives/runtime-benchmarks",
-
   # Locals
   "cfg-primitives/runtime-benchmarks",
   "cfg-traits/runtime-benchmarks",
   "cfg-types/runtime-benchmarks",
   "runtime-common/runtime-benchmarks",
   "liquidity-pools-gateway-routers/runtime-benchmarks",
-
   # Pallet list
   "axelar-gateway-precompile/runtime-benchmarks",
   "chainbridge/runtime-benchmarks",
@@ -356,22 +352,19 @@ runtime-benchmarks = [
 try-runtime = [
   # Enabling optional
   "frame-try-runtime/try-runtime",
-
   # Substrate related
   "sp-runtime/try-runtime",
   "frame-support/try-runtime",
   "frame-system/try-runtime",
   "frame-executive/try-runtime",
   "fp-self-contained/try-runtime",
   "polkadot-runtime-common/try-runtime",
-
   # Locals
   "cfg-primitives/try-runtime",
   "cfg-traits/try-runtime",
   "cfg-types/try-runtime",
   "runtime-common/try-runtime",
   "liquidity-pools-gateway-routers/try-runtime",
-
   # Pallet list
   "axelar-gateway-precompile/try-runtime",
   "chainbridge/try-runtime",

runtime/development/src/lib.rs

@@ -156,7 +156,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
 	spec_name: create_runtime_str!("centrifuge-devel"),
 	impl_name: create_runtime_str!("centrifuge-devel"),
 	authoring_version: 1,
-	spec_version: 1046,
+	spec_version: 1047,
 	impl_version: 1,
 	apis: RUNTIME_API_VERSIONS,
 	transaction_version: 2,
@@ -2026,7 +2026,7 @@ pub type Executive = frame_executive::Executive<
 	frame_system::ChainContext<Runtime>,
 	Runtime,
 	AllPalletsWithSystem,
-	crate::migrations::UpgradeDevelopment1046,
+	crate::migrations::UpgradeDevelopment1047,
 >;
 
 // Frame Order in this block dictates the index of each one in the metadata

runtime/development/src/migrations.rs

@@ -12,5 +12,82 @@
 
 /// The migration set for Development & Demo.
 /// It includes all the migrations that have to be applied on that chain.
-pub type UpgradeDevelopment1046 =
-	pallet_collator_selection::migration::v1::MigrateToV1<crate::Runtime>;
+pub type UpgradeDevelopment1047 = (
+	pallet_collator_selection::migration::v1::MigrateToV1<crate::Runtime>,
+	cleanup_foreign_investments::Migration<crate::Runtime>,
+);
+
+mod cleanup_foreign_investments {
+	use cfg_types::tokens::CurrencyId;
+	use frame_support::{
+		storage::StoragePrefixedMap,
+		traits::{Get, OnRuntimeUpgrade},
+		weights::Weight,
+	};
+	#[cfg(feature = "try-runtime")]
+	use runtime_common::migrations::utils::count_storage_keys;
+	use runtime_common::migrations::utils::remove_undecodable_storage_keys;
+	#[cfg(feature = "try-runtime")]
+	use sp_runtime::DispatchError;
+	#[cfg(feature = "try-runtime")]
+	use sp_runtime::SaturatedConversion;
+
+	pub struct Migration<T>(sp_std::marker::PhantomData<T>);
+
+	const LOG_PREFIX: &str = "CleanupForeignInvestments";
+	impl<T> OnRuntimeUpgrade for Migration<T>
+	where
+		T: pallet_foreign_investments::Config + frame_system::Config,
+	{
+		fn on_runtime_upgrade() -> Weight {
+			log::info!("{LOG_PREFIX} Initiating removal of undecodable keys");
+			let (reads, writes) = remove_undecodable_storage_keys::<CurrencyId>(
+				pallet_foreign_investments::ForeignInvestmentInfo::<T>::final_prefix(),
+			);
+
+			log::info!("{LOG_PREFIX} Removed {writes} undecodable keys");
+
+			T::DbWeight::get().reads_writes(reads, writes)
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<sp_std::vec::Vec<u8>, DispatchError> {
+			let n: u32 = count_storage_keys(
+				pallet_foreign_investments::ForeignInvestmentInfo::<T>::final_prefix(),
+			);
+			let n_count: u32 = pallet_foreign_investments::ForeignInvestmentInfo::<T>::iter_keys()
+				.count()
+				.saturated_into();
+
+			if n == n_count {
+				log::info!(
+					"{LOG_PREFIX} Storage cleanup can be skipped because all keys can be decoded"
+				);
+			} else {
+				log::info!(
+					"{LOG_PREFIX} Failed to decode {} keys, cleanup necessary",
+					n.saturating_sub(n_count)
+				);
+			}
+
+			log::info!("{LOG_PREFIX} pre_upgrade done!",);
+
+			Ok(sp_std::vec![])
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_pre_state: sp_std::vec::Vec<u8>) -> Result<(), DispatchError> {
+			let n: u32 = count_storage_keys(
+				pallet_foreign_investments::ForeignInvestmentInfo::<T>::final_prefix(),
+			);
+			let n_count: u32 = pallet_foreign_investments::ForeignInvestmentInfo::<T>::iter_keys()
+				.count()
+				.saturated_into();
+			assert_eq!(n, n_count);
+
+			log::info!("{LOG_PREFIX} post_upgrade done with {n} remaining storage keys!",);
+
+			Ok(())
+		}
+	}
+}