OnRuntimeUpgrade for Collective v5 migration

Signed-off-by: muraca <mmuraca247@gmail.com>
paritytech · Sep 16, 2023 · c987be4 · c987be4
1 parent c59b240
commit c987be4
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 46 deletions.
diff --git a/substrate/frame/collective/src/migrations/v5.rs b/substrate/frame/collective/src/migrations/v5.rs
@@ -19,8 +19,8 @@ use crate::{Config, Pallet};
 use codec::{Decode, Encode};
 use frame_support::{
 	pallet_prelude::{OptionQuery, StorageVersion, TypeInfo, ValueQuery},
-	sp_runtime::{RuntimeDebug, Saturating},
-	traits::{Get, Hash as PreimageHash, QueryPreimage, StorePreimage},
+	sp_runtime::RuntimeDebug,
+	traits::{Get, Hash as PreimageHash, OnRuntimeUpgrade, QueryPreimage, StorePreimage},
 	weights::Weight,
 	BoundedVec, Identity,
 };
@@ -71,55 +71,130 @@ pub struct OldVotes<AccountId, BlockNumber> {
 	pub end: BlockNumber,
 }
 
-pub fn migrate<T: Config<I>, I: 'static>() -> Weight {
-	let storage_version = StorageVersion::get::<Pallet<T, I>>();
-	log::info!(
-		target: "runtime::collective",
-		"Running migration for collective with storage version {:?}",
-		storage_version,
-	);
-
-	if storage_version <= 4 {
-		let mut count = 0u64;
-
-		for (hash, proposal) in ProposalOf::<T, I>::drain() {
-			let new_hash = <T as Config<I>>::Preimages::note(proposal.encode().into()).unwrap();
-			assert_eq!(new_hash, PreimageHash::from_slice(hash.as_ref()));
-			count.saturating_inc();
-			assert!(<T as Config<I>>::Preimages::is_requested(&new_hash));
+pub struct MigrateToV5<T, I>(sp_std::marker::PhantomData<(T, I)>);
+impl<T: Config<I>, I: 'static> OnRuntimeUpgrade for MigrateToV5<T, I> {
+	#[cfg(feature = "try-runtime")]
+	fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
+		log::info!("pre-migration collective v5");
+		frame_support::ensure!(
+			StorageVersion::get::<Pallet<T, I>>() < 5,
+			"collective already migrated"
+		);
+		let count = ProposalOf::<T, I>::iter().count();
+		if Proposals::<T, I>::get().len() != count {
+			log::info!("collective proposals count inconsistency");
+		}
+		if Members::<T, I>::get().len() > <T as Config<I>>::MaxMembers::get() as usize {
+			log::info!("collective members exceeds MaxMembers");
 		}
 
-		Proposals::<T, I>::kill();
+		Ok((count as u32).encode())
+	}
+
+	fn on_runtime_upgrade() -> Weight {
+		let storage_version = StorageVersion::get::<Pallet<T, I>>();
+		log::info!(
+			target: "runtime::collective",
+			"Running migration for collective with storage version {:?}",
+			storage_version,
+		);
+
+		if storage_version <= 4 {
+			let mut weight = T::DbWeight::get().reads(1);
+
+			// ProposalOf
+			for (hash, proposal) in ProposalOf::<T, I>::drain() {
+				weight = weight.saturating_add(T::DbWeight::get().reads_writes(1, 1));
+				match <T as Config<I>>::Preimages::note(proposal.encode().into()) {
+					Ok(new_hash) => {
+						weight = weight.saturating_add(T::DbWeight::get().writes(1));
+						if new_hash != PreimageHash::from_slice(hash.as_ref()) {
+							log::info!(
+								target: "runtime::collective",
+								"Preimage hash mismatch for proposal {:?}. Expected {:?}, got {:?}",
+								hash,
+								new_hash,
+								PreimageHash::from_slice(hash.as_ref())
+							);
+						}
+						if !<T as Config<I>>::Preimages::is_requested(&new_hash) {
+							log::info!(
+								target: "runtime::collective",
+								"Preimage for proposal {:?} was not requested",
+								hash,
+							);
+						}
+					},
+					Err(_) => {
+						log::info!(
+							target: "runtime::collective",
+							"Failed to note preimage for proposal {:?}",
+							hash,
+						);
+					},
+				}
+			}
+
+			// Proposals
+			Proposals::<T, I>::kill();
+			weight = weight.saturating_add(T::DbWeight::get().writes(1));
 
-		crate::Voting::<T, I>::translate::<OldVotes<T::AccountId, BlockNumberFor<T>>, _>(
-			|_, vote| {
-				count.saturating_inc();
-				Some(
+			// Voting
+			for (hash, vote) in Voting::<T, I>::drain() {
+				weight = weight.saturating_add(T::DbWeight::get().reads_writes(1, 2));
+				crate::Voting::<T, I>::insert(
+					PreimageHash::from_slice(hash.as_ref()),
 					crate::Votes::<T::AccountId, BlockNumberFor<T>, <T as Config<I>>::MaxMembers> {
 						index: vote.index,
 						threshold: vote.threshold,
-						ayes: vote
-							.ayes
-							.try_into()
-							.expect("runtime::collective migration failed, ayes overflow"),
-						nays: vote
-							.nays
-							.try_into()
-							.expect("runtime::collective migration failed, nays overflow"),
+						// the following operations are safe since the bound was previously enforced
+						// by runtime code
+						ayes: BoundedVec::truncate_from(vote.ayes),
+						nays: BoundedVec::truncate_from(vote.nays),
 						end: vote.end,
 					},
-				)
-			},
-		);
+				);
+			}
 
-		StorageVersion::new(5).put::<Pallet<T, I>>();
-		T::DbWeight::get().reads_writes(count, count + 2)
-	} else {
-		log::warn!(
-			target: "runtime::collective",
-			"Attempted to apply migration to V5 but failed because storage version is {:?}",
-			storage_version,
+			// Members
+			crate::Members::<T, I>::put(BoundedVec::truncate_from(Members::<T, I>::get()));
+			weight = weight.saturating_add(T::DbWeight::get().reads_writes(1, 1));
+
+			// StorageVersion
+			StorageVersion::new(5).put::<Pallet<T, I>>();
+			weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+			weight
+		} else {
+			log::info!(
+				target: "runtime::collective",
+				"Attempted to apply migration to V5 but failed because storage version is {:?}",
+				storage_version,
+			);
+			T::DbWeight::get().reads(1)
+		}
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(state: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
+		use frame_support::ensure;
+		log::info!("post-migration collective v5");
+		ensure!(StorageVersion::get::<Pallet<T, I>>() == 5, "collective migration to v5 failed");
+
+		match u32::decode(&mut state.as_slice()) {
+			Ok(count) => ensure!(
+				crate::Voting::<T, I>::count() == count,
+				"collective new proposal count should equal the old count",
+			),
+			Err(_) => return Err("the state parameter is not correct".into()),
+		}
+
+		ensure!(
+			ProposalOf::<T, I>::iter().count() == 0,
+			"collective v4 ProposalOf should be empty"
 		);
-		Weight::zero()
+		ensure!(Proposals::<T, I>::get().len() == 0, "collective v4 Proposal should be empty");
+
+		Pallet::<T, I>::do_try_state()
 	}
 }
diff --git a/substrate/frame/collective/src/tests.rs b/substrate/frame/collective/src/tests.rs
@@ -1545,10 +1545,12 @@ fn migration_v4() {
 	});
 }
 
+#[cfg(all(feature = "try-runtime", test))]
 #[test]
 fn migration_v5() {
 	ExtBuilder::default().build_and_execute(|| {
-		StorageVersion::new(4).put::<Collective>();
+		use frame_support::traits::OnRuntimeUpgrade;
+		StorageVersion::new(4).put::<Pallet<Test, ()>>();
 
 		// Setup the members.
 		let members: Vec<<Test as frame_system::Config>::AccountId> =
@@ -1563,6 +1565,7 @@ fn migration_v5() {
 			proposal_hash
 		]);
 		crate::migrations::v5::ProposalOf::<Test, ()>::insert(proposal_hash, proposal);
+		ProposalCount::<Test, ()>::put(1);
 
 		let vote = crate::migrations::v5::OldVotes {
 			index: 0,
@@ -1576,9 +1579,7 @@ fn migration_v5() {
 		crate::migrations::v5::Voting::<Test, ()>::insert(proposal_hash, vote.clone());
 
 		// Run migration.
-		crate::migrations::v5::migrate::<Test, ()>();
-		// Check that the storage version is updated.
-		assert_eq!(StorageVersion::get::<Pallet<Test, ()>>(), 5);
+		assert_ok!(crate::migrations::v5::MigrateToV5::<Test, ()>::try_on_runtime_upgrade(true));
 
 		// Check that the vote is present and bounded
 		assert_eq!(