🏗️ Vesting pallet block number provider migration

integration-tests/src/lib.rs

@@ -118,5 +118,7 @@ pub mod shortcuts {
 
 	pub type PolkadotSystem = <PolkaNet as Chain>::System;
 	pub type PolimecSystem = <PolimecNet as Chain>::System;
+
+	pub type PolimecParachainSystem = <PolimecNet as PolimecParaPallet>::ParachainSystem;
 }
 pub use shortcuts::*;

integration-tests/src/tests/vest.rs

@@ -14,21 +14,35 @@
 // You should have received a copy of the GNU General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{polimec::ED, *};
 /// Tests for the oracle pallet integration.
 /// Alice, Bob, Charlie are members of the OracleProvidersMembers.
 /// Only members should be able to feed data into the oracle.
-use frame_support::traits::fungible::Inspect;
-use frame_support::traits::fungible::Mutate;
+use crate::{polimec::ED, *};
+use cumulus_pallet_parachain_system::ValidationData;
+use cumulus_primitives_core::PersistedValidationData;
+use frame_support::traits::fungible::{Inspect, Mutate};
 use macros::generate_accounts;
 use pallet_funding::assert_close_enough;
 use pallet_vesting::VestingInfo;
 use polimec_runtime::{Balances, BlockchainOperationTreasury, ParachainStaking, RuntimeOrigin, Vesting, PLMC};
 use sp_core::crypto::get_public_from_string_or_panic;
-use sp_runtime::Perquintill;
+use sp_runtime::{Perquintill, SaturatedConversion};
 
 generate_accounts!(PEPE, CARLOS);
 
+fn set_relay_chain_block_number<T: cumulus_pallet_parachain_system::Config>(to: u32) {
+	let mut validation_data = ValidationData::<T>::get().unwrap_or_else(||
+		// PersistedValidationData does not impl default in non-std
+		PersistedValidationData {
+			parent_head: vec![].into(),
+			relay_parent_number: Default::default(),
+			max_pov_size: Default::default(),
+			relay_parent_storage_root: Default::default(),
+		});
+	validation_data.relay_parent_number = to.saturated_into();
+	ValidationData::<T>::put(validation_data)
+}
+
 #[test]
 fn base_vested_can_stake() {
 	PolimecNet::execute_with(|| {
@@ -89,10 +103,7 @@ fn base_can_withdraw_when_free_is_below_frozen_with_hold() {
 		// Vesting schedule for PEPE of 20k PLMC + ED, which should have start date before it is applied
 		let vesting_schedule = VestingInfo::new(2_020 * PLMC, 10 * PLMC, 0);
 
-		assert_eq!(Balances::free_balance(&CARLOS.into()), 0);
-		// We need some free balance at the time of the vested transfer
-		// Otherwise the user will never have free balance to pay for the "vest" extrinsic
-		PolimecSystem::set_block_number(1u32);
+		set_relay_chain_block_number::<PolimecRuntime>(1);
 
 		// The actual vested transfer
 		assert_ok!(Vesting::vested_transfer(
@@ -101,7 +112,7 @@ fn base_can_withdraw_when_free_is_below_frozen_with_hold() {
 			vesting_schedule
 		));
 
-		// Vested transfer didnt start with the full amount locked, since start date was befire execution
+		// Vested transfer didnt start with the full amount locked, since start date was before execution
 		assert_eq!(Balances::usable_balance(&CARLOS.into()), 10 * PLMC);
 
 		let carlos_acc: AccountId = CARLOS.into();
@@ -123,7 +134,7 @@ fn base_can_withdraw_when_free_is_below_frozen_with_hold() {
 		assert_eq!(Balances::usable_balance(&CARLOS.into()), 10 * PLMC);
 
 		// Be able to vest 10 more PLMC for this example description
-		PolimecSystem::set_block_number(2u32);
+		set_relay_chain_block_number::<PolimecRuntime>(2);
 
 		// This should pass if the fee is correctly deducted with the new fee struct
 		assert_ok!(Vesting::vest(RuntimeOrigin::signed(CARLOS.into())));

pallets/funding/src/benchmarking.rs

@@ -97,7 +97,7 @@ pub fn string_account<AccountId: Decode>(
 
 #[benchmarks(
 	where
-	T: Config + frame_system::Config<RuntimeEvent = <T as Config>::RuntimeEvent> + pallet_balances::Config<Balance = Balance> + cumulus_pallet_parachain_system::Config + sp_std::fmt::Debug,
+	T: Config + frame_system::Config<RuntimeEvent = <T as Config>::RuntimeEvent> + pallet_balances::Config<Balance = Balance> + cumulus_pallet_parachain_system::Config + core::fmt::Debug,
 	<T as Config>::RuntimeEvent: TryInto<Event<T>> + Parameter + Member,
 	<T as Config>::Price: From<u128>,
 	T::Hash: From<H256>,

runtimes/polimec/Cargo.toml

@@ -74,6 +74,7 @@ sp-api.workspace = true
 sp-block-builder.workspace = true
 sp-consensus-aura.workspace = true
 sp-core.workspace = true
+sp-io.workspace = true
 sp-inherents.workspace = true
 sp-offchain.workspace = true
 sp-runtime.workspace = true
@@ -177,6 +178,7 @@ std = [
 	"sp-core/std",
 	"sp-genesis-builder/std",
 	"sp-inherents/std",
+	"sp-io/std",
 	"sp-offchain/std",
 	"sp-runtime/std",
 	"sp-session/std",

runtimes/polimec/src/custom_migrations/mod.rs

@@ -18,3 +18,5 @@
 #![allow(clippy::all)]
 
 pub mod asset_id_migration;
+
+pub mod vesting;

runtimes/polimec/src/custom_migrations/vesting.rs

@@ -0,0 +1,149 @@
+use alloc::vec::Vec;
+use frame_support::{traits::Currency, BoundedVec};
+use frame_system::pallet_prelude::BlockNumberFor;
+use pallet_vesting::{MaxVestingSchedulesGet, VestingInfo};
+use parity_scale_codec::Encode;
+
+#[cfg(feature = "try-runtime")]
+use parity_scale_codec::Decode;
+#[cfg(feature = "try-runtime")]
+use sp_runtime::DispatchError;
+
+type VestingInfoOf<T> = VestingInfo<
+	<<T as pallet_vesting::Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance,
+	BlockNumberFor<T>,
+>;
+pub type Values<T> = BoundedVec<VestingInfoOf<T>, MaxVestingSchedulesGet<T>>;
+
+pub mod v1 {
+	use super::*;
+	use core::marker::PhantomData;
+	use frame_support::traits::OnRuntimeUpgrade;
+	use parachains_common::impls::AccountIdOf;
+	use sp_core::Get;
+	use sp_runtime::{traits::BlockNumberProvider, Saturating, Weight};
+
+	const LOG: &str = "pallet_vesting::migration::v1";
+
+	/// Stores the status of vesting pallet migration to async backing. If it is populated, the migration already happened.
+	pub const VESTING_ASYNC_BACKED_KEY: &[u8; 31] = b"vesting_async_backing_migration";
+
+	pub struct UncheckedMigrationToV1<T: pallet_vesting::Config>(PhantomData<T>);
+
+	impl<T: pallet_vesting::Config> OnRuntimeUpgrade for UncheckedMigrationToV1<T> {
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, DispatchError> {
+			let vesting_async_backed = sp_io::storage::get(VESTING_ASYNC_BACKED_KEY);
+
+			if vesting_async_backed.is_some() {
+				log::info!(target: LOG, "Skipping migration as vesting pallet is already migrated with async backing");
+				return Ok(Vec::new());
+			}
+
+			let migration_count = pallet_vesting::Vesting::<T>::iter().count() as u32;
+			log::info!(target: LOG, "Pre-upgrade: {} UserMigrations entries", migration_count);
+
+			let vestings = pallet_vesting::Vesting::<T>::iter().collect::<Vec<_>>();
+
+			Ok((migration_count, vestings).encode())
+		}
+
+		fn on_runtime_upgrade() -> Weight {
+			let vesting_async_backed = sp_io::storage::get(VESTING_ASYNC_BACKED_KEY);
+
+			if vesting_async_backed.is_some() {
+				log::info!(target: LOG, "Skipping migration as vesting pallet is already migrated with async backing");
+				return Weight::zero();
+			}
+
+			let mut items = 0u64;
+			let translate_vesting_info = |_: AccountIdOf<T>, vesting_info: Values<T>| -> Option<Values<T>> {
+				let migrated: Vec<_> = vesting_info
+					.iter()
+					.map(|vesting| {
+						items = items.saturating_add(1);
+
+						// adjust starting block to relay chain block number
+						let relay_chain_now = T::BlockNumberProvider::current_block_number();
+
+						let polimec_now = frame_system::Pallet::<T>::current_block_number();
+
+						let two = 2_u32.into();
+
+						let relay_chain_starting_block = if polimec_now < vesting.starting_block() {
+							let blocks_diff = vesting.starting_block().saturating_sub(polimec_now);
+							relay_chain_now.saturating_add(blocks_diff.saturating_mul(two))
+						} else {
+							let blocks_passed = polimec_now.saturating_sub(vesting.starting_block());
+							relay_chain_now.saturating_sub(blocks_passed.saturating_mul(two))
+						};
+
+						let adjusted_per_block = vesting.per_block().saturating_mul(2_u32.into());
+
+						VestingInfo::new(vesting.locked(), adjusted_per_block, relay_chain_starting_block)
+					})
+					.collect();
+
+				Values::<T>::try_from(migrated).ok()
+			};
+
+			log::info!(target: LOG, "Starting vesting time migration to V1");
+
+			pallet_vesting::Vesting::<T>::translate(translate_vesting_info);
+
+			log::info!(target: LOG, "Migrated {} vesting entries", items);
+
+			sp_io::storage::set(VESTING_ASYNC_BACKED_KEY, &().encode()[..]);
+
+			T::DbWeight::get().reads_writes(items, items)
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(pre_state: Vec<u8>) -> Result<(), DispatchError> {
+			let vesting_async_backed = sp_io::storage::get(VESTING_ASYNC_BACKED_KEY);
+			if vesting_async_backed.is_some() {
+				log::info!(target: LOG, "Skipping migration as vesting pallet is already migrated with async backing");
+				return Ok(());
+			}
+
+			let (pre_migration_count, pre_vestings): (u32, Vec<(AccountIdOf<T>, Values<T>)>) =
+				Decode::decode(&mut &pre_state[..]).expect("Failed to decode pre-migration state");
+
+			let post_migration_count = pallet_vesting::Vesting::<T>::iter().count() as u32;
+
+			if pre_migration_count != post_migration_count {
+				return Err("Migration count mismatch".into());
+			}
+
+			for (account, pre_vesting) in pre_vestings {
+				let post_vesting = pallet_vesting::Vesting::<T>::get(&account).unwrap_or_default();
+
+				// check that the starting block has been adjusted
+				let relay_chain_now = T::BlockNumberProvider::current_block_number();
+
+				for (pre_vesting_info, post_vesting_info) in pre_vesting.iter().zip(post_vesting.iter()) {
+					assert_ne!(
+						pre_vesting_info.starting_block(),
+						post_vesting_info.starting_block(),
+						"Starting block not adjusted"
+					);
+					assert!(
+						post_vesting_info.starting_block() <=
+							relay_chain_now.try_into().ok().expect("safe to convert; qed"),
+						"Starting block not adjusted correctly"
+					);
+
+					assert!(
+						post_vesting_info.per_block() ==
+							pre_vesting_info
+								.per_block()
+								.saturating_mul(2_u32.try_into().ok().expect("safe to convert; qed")),
+						"Per block not adjusted"
+					);
+				}
+			}
+
+			Ok(())
+		}
+	}
+}