Skip to content

Commit 84ca86b

Browse files
dastansamdastansam
committed
Linear relaase pallet migration
1 parent 7c4f030 commit 84ca86b

File tree

3 files changed

+221
-80
lines changed

3 files changed

+221
-80
lines changed

‎pallets/linear-release/src/lib.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ use frame_support::{
3333
traits::{
3434
fungible::{BalancedHold, Inspect, InspectHold, Mutate, MutateHold},
3535
tokens::{Balance, Precision},
36-
Get, WithdrawReasons,
36+
Get, StorageVersion, WithdrawReasons,
3737
},
3838
};
3939
use frame_system::pallet_prelude::*;
@@ -110,6 +110,9 @@ impl<T: Config> Get<u32> for MaxVestingSchedulesGet<T> {
110110
}
111111
}
112112

113+
/// Current storage version
114+
pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
115+
113116
/// Enable `dev_mode` for this pallet.
114117
#[frame_support::pallet(dev_mode)]
115118
pub mod pallet {
@@ -172,6 +175,7 @@ pub mod pallet {
172175
// Simple declaration of the `Pallet` type. It is placeholder we use to implement traits and
173176
// method.
174177
#[pallet::pallet]
178+
#[pallet::storage_version(STORAGE_VERSION)]
175179
pub struct Pallet<T>(_);
176180

177181
#[pallet::event]

‎runtimes/polimec/src/custom_migrations/linear_release.rs‎

Lines changed: 214 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -1,108 +1,244 @@
1-
use crate::{Balance, BlockNumber, Runtime, RuntimeHoldReason};
1+
use crate::{Runtime, RuntimeHoldReason};
22
use alloc::vec::Vec;
33
use frame_support::{storage::storage_prefix, traits::OnRuntimeUpgrade, BoundedVec};
4-
use pallet_linear_release::{MaxVestingSchedulesGet, VestingInfo};
4+
use pallet_linear_release::{MaxVestingSchedulesGet, VestingInfo, VestingInfoOf};
55

66
#[cfg(feature = "try-runtime")]
77
use frame_support::{ensure, migrations::VersionedPostUpgradeData, traits::GetStorageVersion};
88
#[cfg(feature = "try-runtime")]
99
use itertools::Itertools;
1010
#[cfg(feature = "try-runtime")]
11-
use parity_scale_codec::{DecodeAll, Encode};
11+
use parity_scale_codec::{Decode, DecodeAll, Encode};
12+
#[cfg(feature = "try-runtime")]
13+
use sp_runtime::DispatchError;
1214

13-
pub type Values = BoundedVec<VestingInfo<Balance, BlockNumber>, MaxVestingSchedulesGet<Runtime>>;
15+
pub type Values<T> = BoundedVec<VestingInfoOf<T>, MaxVestingSchedulesGet<T>>;
1416

15-
pub struct LinearReleaseVestingMigration;
16-
impl OnRuntimeUpgrade for LinearReleaseVestingMigration {
17-
#[cfg(feature = "try-runtime")]
18-
fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
19-
use crate::LinearRelease;
17+
pub mod v1 {
18+
use super::*;
19+
use core::marker::PhantomData;
20+
use frame_support::{migrations::VersionedMigration, traits::UncheckedOnRuntimeUpgrade};
21+
use frame_system::pallet_prelude::BlockNumberFor;
22+
use pallet_linear_release::{AccountIdOf, ReasonOf};
23+
use sp_core::Get;
24+
use sp_runtime::{traits::BlockNumberProvider, Saturating, Weight};
2025

21-
let funding_on_chain_version = LinearRelease::on_chain_storage_version();
22-
if funding_on_chain_version == 0 {
23-
let storage = pallet_linear_release::Vesting::<Runtime>::iter().collect_vec();
24-
ensure!(storage.len() == 15, "LinearReleaseVestingMigration: Invalid storage length in pre_upgrade");
25-
Ok(VersionedPostUpgradeData::MigrationExecuted(Vec::new()).encode())
26-
} else {
27-
Ok(VersionedPostUpgradeData::Noop.encode())
26+
const LOG: &str = "linear_release::migration::v1";
27+
pub struct LinearReleaseVestingInfoMigration;
28+
29+
pub struct UncheckedMigrationToV1<T: pallet_linear_release::Config + cumulus_pallet_parachain_system::Config>(
30+
PhantomData<T>,
31+
);
32+
33+
impl<T: pallet_linear_release::Config + cumulus_pallet_parachain_system::Config> UncheckedOnRuntimeUpgrade
34+
for UncheckedMigrationToV1<T>
35+
{
36+
#[cfg(feature = "try-runtime")]
37+
fn pre_upgrade() -> Result<Vec<u8>, DispatchError> {
38+
let migration_count = pallet_linear_release::Vesting::<T>::iter().count() as u32;
39+
log::info!(target: LOG, "Pre-upgrade: {} UserMigrations entries", migration_count);
40+
41+
let vestings = pallet_linear_release::Vesting::<T>::iter().collect::<Vec<_>>();
42+
43+
Ok((migration_count, vestings).encode())
2844
}
29-
}
3045

31-
fn on_runtime_upgrade() -> frame_support::weights::Weight {
32-
let db_weight = <Runtime as frame_system::Config>::DbWeight::get();
46+
fn on_runtime_upgrade() -> Weight {
47+
let mut items = 0u64;
48+
let translate_vesting_info =
49+
|_: AccountIdOf<T>, _: ReasonOf<T>, vesting_info: Values<T>| -> Option<Values<T>> {
50+
let migrated: Vec<_> = vesting_info
51+
.iter()
52+
.map(|vesting| {
53+
items = items.saturating_add(1);
54+
55+
// adjust starting block to relay chain block number
56+
let relay_chain_now: BlockNumberFor<T> = T::BlockNumberProvider::current_block_number()
57+
.try_into()
58+
.ok()
59+
.expect("BlockNumber conversion failed");
60+
61+
let polimec_now = frame_system::Pallet::<T>::current_block_number();
62+
let blocks_passed = polimec_now.saturating_sub(vesting.starting_block);
63+
let relay_chain_starting_block = relay_chain_now.saturating_sub(
64+
blocks_passed.saturating_mul(2_u32.try_into().ok().expect("safe to convert; qed")),
65+
);
66+
67+
let adjusted_per_block =
68+
vesting.per_block.saturating_mul(2_u32.try_into().ok().expect("safe to convert; qed"));
69+
70+
VestingInfo {
71+
locked: vesting.locked,
72+
per_block: adjusted_per_block,
73+
starting_block: relay_chain_starting_block,
74+
}
75+
})
76+
.collect();
77+
78+
let bounded = Values::<T>::try_from(migrated)
79+
.unwrap_or_else(|_| panic!("MaxVestingSchedulesGet exceeded during migration"));
3380

34-
// Step 1: Get account count for weight calculation
35-
let account_count = pallet_linear_release::Vesting::<Runtime>::iter().count();
81+
Some(bounded)
82+
};
3683

37-
// Initial weight: account_count reads + 1 clear_prefix operation
38-
let mut total_weight = db_weight.reads_writes(account_count as u64, 1); // For clear_prefix
84+
log::info!(target: LOG, "Starting linear release vesting time migration to V1");
3985

40-
// Step 2: Collect all accounts and their hold amounts
41-
let mut account_holds = Vec::with_capacity(account_count);
86+
pallet_linear_release::Vesting::<T>::translate(translate_vesting_info);
4287

43-
for (account, reason, vesting_info) in pallet_linear_release::Vesting::<Runtime>::iter() {
44-
if !vesting_info.is_empty() {
45-
log::info!(
46-
"Found vesting for account: {:?}, reason: {:?}, schedule count: {:?}",
47-
account,
48-
reason,
49-
vesting_info.len()
50-
);
51-
account_holds.push((account, reason, vesting_info[0]));
88+
log::info!(target: LOG, "Migrated {} linear release vesting entries", items);
89+
90+
T::DbWeight::get().reads_writes(items, items)
91+
}
92+
93+
#[cfg(feature = "try-runtime")]
94+
fn post_upgrade(pre_state: Vec<u8>) -> Result<(), DispatchError> {
95+
let (pre_migration_count, pre_vestings): (u32, Vec<((AccountIdOf<T>, ReasonOf<T>), Values<T>)>) =
96+
Decode::decode(&mut &pre_state[..]).expect("Failed to decode pre-migration state");
97+
98+
let post_migration_count = pallet_linear_release::Vesting::<T>::iter().count() as u32;
99+
100+
if pre_migration_count != post_migration_count {
101+
return Err("Migration count mismatch".into());
102+
}
103+
104+
for ((account, reason), pre_vesting) in pre_vestings {
105+
let post_vesting = pallet_linear_release::Vesting::<T>::get(&account, &reason).unwrap_or_default();
106+
107+
// check that the starting block has been adjusted
108+
let relay_chain_now =
109+
cumulus_pallet_parachain_system::RelaychainDataProvider::<T>::current_block_number();
110+
111+
for (pre_vesting_info, post_vesting_info) in pre_vesting.iter().zip(post_vesting.iter()) {
112+
assert_ne!(
113+
pre_vesting_info.starting_block, post_vesting_info.starting_block,
114+
"Starting block not adjusted"
115+
);
116+
assert!(
117+
post_vesting_info.starting_block <=
118+
relay_chain_now.try_into().ok().expect("safe to convert; qed"),
119+
"Starting block not adjusted correctly"
120+
);
121+
122+
assert!(
123+
post_vesting_info.per_block ==
124+
pre_vesting_info
125+
.per_block
126+
.saturating_mul(2_u32.try_into().ok().expect("safe to convert; qed")),
127+
"Per block not adjusted"
128+
);
129+
}
130+
}
131+
132+
Ok(())
133+
}
134+
}
135+
136+
pub type LinearReleaseVestingMigrationV1<T> = VersionedMigration<
137+
0,
138+
1,
139+
UncheckedMigrationToV1<T>,
140+
pallet_linear_release::Pallet<T>,
141+
<T as frame_system::Config>::DbWeight,
142+
>;
143+
}
144+
145+
pub mod unversioned {
146+
use super::*;
147+
148+
pub struct LinearReleaseVestingMigration;
149+
impl OnRuntimeUpgrade for LinearReleaseVestingMigration {
150+
#[cfg(feature = "try-runtime")]
151+
fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
152+
use crate::LinearRelease;
153+
154+
let funding_on_chain_version = LinearRelease::on_chain_storage_version();
155+
if funding_on_chain_version == 0 {
156+
let storage = pallet_linear_release::Vesting::<Runtime>::iter().collect_vec();
157+
ensure!(storage.len() == 15, "LinearReleaseVestingMigration: Invalid storage length in pre_upgrade");
158+
Ok(VersionedPostUpgradeData::MigrationExecuted(Vec::new()).encode())
52159
} else {
53-
log::warn!("Empty vesting info found for account: {:?}, reason: {:?}", account, reason);
160+
Ok(VersionedPostUpgradeData::Noop.encode())
54161
}
55162
}
56163

57-
// Step 3: Clear all corrupted vesting entries
58-
let pallet_prefix = storage_prefix(b"LinearRelease", b"Vesting");
59-
let removed_keys = frame_support::storage::unhashed::clear_prefix(&pallet_prefix, None, None);
60-
log::info!("Cleared {:#?} vesting storage keys", removed_keys.deconstruct());
61-
62-
// Step 4: Create fresh vesting entries for all accounts with holds
63-
let mut success_count = 0u64;
64-
let mut failure_count = 0u64;
65-
66-
for (account, _corrupted_reason, vesting_info) in account_holds {
67-
// Create a BoundedVec with this schedule - reuse original VestingInfo value, but set the k2 to the new pallet_funding HoldReason::Participation.
68-
let mut schedules = BoundedVec::<_, MaxVestingSchedulesGet<Runtime>>::default();
69-
70-
match schedules.try_push(vesting_info) {
71-
Ok(_) => {
72-
pallet_linear_release::Vesting::<Runtime>::insert(
73-
&account,
74-
&RuntimeHoldReason::Funding(pallet_funding::HoldReason::Participation),
75-
schedules,
164+
fn on_runtime_upgrade() -> frame_support::weights::Weight {
165+
let db_weight = <Runtime as frame_system::Config>::DbWeight::get();
166+
167+
// Step 1: Get account count for weight calculation
168+
let account_count = pallet_linear_release::Vesting::<Runtime>::iter().count();
169+
170+
// Initial weight: account_count reads + 1 clear_prefix operation
171+
let mut total_weight = db_weight.reads_writes(account_count as u64, 1); // For clear_prefix
172+
173+
// Step 2: Collect all accounts and their hold amounts
174+
let mut account_holds = Vec::with_capacity(account_count);
175+
176+
for (account, reason, vesting_info) in pallet_linear_release::Vesting::<Runtime>::iter() {
177+
if !vesting_info.is_empty() {
178+
log::info!(
179+
"Found vesting for account: {:?}, reason: {:?}, schedule count: {:?}",
180+
account,
181+
reason,
182+
vesting_info.len()
76183
);
77-
success_count += 1;
78-
total_weight = total_weight.saturating_add(db_weight.writes(1));
79-
},
80-
Err(_) => {
81-
log::error!("Failed to add vesting schedule to BoundedVec for account {:?}", account);
82-
failure_count += 1;
83-
},
184+
account_holds.push((account, reason, vesting_info[0]));
185+
} else {
186+
log::warn!("Empty vesting info found for account: {:?}, reason: {:?}", account, reason);
187+
}
84188
}
85-
}
86189

87-
log::info!(
88-
"Migration complete. Successfully created {} new vesting entries. Failed: {}",
89-
success_count,
90-
failure_count
91-
);
190+
// Step 3: Clear all corrupted vesting entries
191+
let pallet_prefix = storage_prefix(b"LinearRelease", b"Vesting");
192+
let removed_keys = frame_support::storage::unhashed::clear_prefix(&pallet_prefix, None, None);
193+
log::info!("Cleared {:#?} vesting storage keys", removed_keys.deconstruct());
194+
195+
// Step 4: Create fresh vesting entries for all accounts with holds
196+
let mut success_count = 0u64;
197+
let mut failure_count = 0u64;
198+
199+
for (account, _corrupted_reason, vesting_info) in account_holds {
200+
// Create a BoundedVec with this schedule - reuse original VestingInfo value, but set the k2 to the new pallet_funding HoldReason::Participation.
201+
let mut schedules = BoundedVec::<_, MaxVestingSchedulesGet<Runtime>>::default();
202+
203+
match schedules.try_push(vesting_info) {
204+
Ok(_) => {
205+
pallet_linear_release::Vesting::<Runtime>::insert(
206+
&account,
207+
&RuntimeHoldReason::Funding(pallet_funding::HoldReason::Participation),
208+
schedules,
209+
);
210+
success_count += 1;
211+
total_weight = total_weight.saturating_add(db_weight.writes(1));
212+
},
213+
Err(_) => {
214+
log::error!("Failed to add vesting schedule to BoundedVec for account {:?}", account);
215+
failure_count += 1;
216+
},
217+
}
218+
}
92219

93-
total_weight
94-
}
220+
log::info!(
221+
"Migration complete. Successfully created {} new vesting entries. Failed: {}",
222+
success_count,
223+
failure_count
224+
);
95225

96-
#[cfg(feature = "try-runtime")]
97-
fn post_upgrade(versioned_post_upgrade_data_bytes: alloc::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
98-
let storage = pallet_linear_release::Vesting::<Runtime>::iter().collect_vec();
99-
ensure!(storage.len() == 15, "LinearReleaseVestingMigration: Invalid storage length in post_upgrade");
226+
total_weight
227+
}
100228

101-
match <VersionedPostUpgradeData>::decode_all(&mut &versioned_post_upgrade_data_bytes[..])
102-
.map_err(|_| "VersionedMigration post_upgrade failed to decode PreUpgradeData")?
103-
{
104-
VersionedPostUpgradeData::MigrationExecuted(_inner_bytes) => Ok(()),
105-
VersionedPostUpgradeData::Noop => Ok(()),
229+
#[cfg(feature = "try-runtime")]
230+
fn post_upgrade(
231+
versioned_post_upgrade_data_bytes: alloc::vec::Vec<u8>,
232+
) -> Result<(), sp_runtime::TryRuntimeError> {
233+
let storage = pallet_linear_release::Vesting::<Runtime>::iter().collect_vec();
234+
ensure!(storage.len() == 15, "LinearReleaseVestingMigration: Invalid storage length in post_upgrade");
235+
236+
match <VersionedPostUpgradeData>::decode_all(&mut &versioned_post_upgrade_data_bytes[..])
237+
.map_err(|_| "VersionedMigration post_upgrade failed to decode PreUpgradeData")?
238+
{
239+
VersionedPostUpgradeData::MigrationExecuted(_inner_bytes) => Ok(()),
240+
VersionedPostUpgradeData::Noop => Ok(()),
241+
}
106242
}
107243
}
108244
}

‎runtimes/polimec/src/lib.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,10 +188,11 @@ pub mod migrations {
188188
#[allow(unused_parens)]
189189
pub type Unreleased = (
190190
super::custom_migrations::asset_id_migration::FromOldAssetIdMigration,
191-
super::custom_migrations::linear_release::LinearReleaseVestingMigration,
191+
// super::custom_migrations::linear_release::unversioned::LinearReleaseVestingMigration,
192192
pallet_funding::migrations::storage_migrations::v6::MigrationToV6<Runtime>,
193193
RemovePallet<IdentityPalletName, RuntimeDbWeight>,
194194
pallet_funding::migrations::vesting_info::v7::MigrationToV8<Runtime>,
195+
super::custom_migrations::linear_release::v1::LinearReleaseVestingMigrationV1<Runtime>,
195196
);
196197
}
197198

0 commit comments

Comments
 (0)