|
17 | 17 | //! A module that is responsible for migration of storage for Collator Selection. |
18 | 18 |
|
19 | 19 | use super::*; |
20 | | -use frame_support::traits::OnRuntimeUpgrade; |
| 20 | +use frame_support::traits::{OnRuntimeUpgrade, UncheckedOnRuntimeUpgrade}; |
21 | 21 | use log; |
22 | 22 |
|
| 23 | +/// Migrate to v2. Should have been part of <https://github.com/paritytech/polkadot-sdk/pull/1340>. |
| 24 | +pub mod v2 { |
| 25 | + use super::*; |
| 26 | + use frame_support::{ |
| 27 | + pallet_prelude::*, |
| 28 | + storage_alias, |
| 29 | + traits::{Currency, ReservableCurrency}, |
| 30 | + }; |
| 31 | + use sp_runtime::traits::{Saturating, Zero}; |
| 32 | + #[cfg(feature = "try-runtime")] |
| 33 | + use sp_std::vec::Vec; |
| 34 | + |
| 35 | + /// [`UncheckedMigrationToV2`] wrapped in a |
| 36 | + /// [`VersionedMigration`](frame_support::migrations::VersionedMigration), ensuring the |
| 37 | + /// migration is only performed when on-chain version is 1. |
| 38 | + pub type MigrationToV2<T> = frame_support::migrations::VersionedMigration< |
| 39 | + 1, |
| 40 | + 2, |
| 41 | + UncheckedMigrationToV2<T>, |
| 42 | + Pallet<T>, |
| 43 | + <T as frame_system::Config>::DbWeight, |
| 44 | + >; |
| 45 | + |
| 46 | + #[storage_alias] |
| 47 | + pub type Candidates<T: Config> = StorageValue< |
| 48 | + Pallet<T>, |
| 49 | + BoundedVec<CandidateInfo<<T as frame_system::Config>::AccountId, <<T as Config>::Currency as Currency<<T as frame_system::Config>::AccountId>>::Balance>, <T as Config>::MaxCandidates>, |
| 50 | + ValueQuery, |
| 51 | + >; |
| 52 | + |
| 53 | + /// Migrate to V2. |
| 54 | + pub struct UncheckedMigrationToV2<T>(sp_std::marker::PhantomData<T>); |
| 55 | + impl<T: Config + pallet_balances::Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV2<T> { |
| 56 | + fn on_runtime_upgrade() -> Weight { |
| 57 | + let mut weight = Weight::zero(); |
| 58 | + let mut count: u64 = 0; |
| 59 | + // candidates who exist under the old `Candidates` key |
| 60 | + let candidates = Candidates::<T>::take(); |
| 61 | + |
| 62 | + // New candidates who have registered since the upgrade. Under normal circumstances, |
| 63 | + // this should not exist because the migration should be applied when the upgrade |
| 64 | + // happens. But in Polkadot/Kusama we messed this up, and people registered under |
| 65 | + // `CandidateList` while their funds were locked in `Candidates`. |
| 66 | + let new_candidate_list = CandidateList::<T>::get(); |
| 67 | + if new_candidate_list.len().is_zero() { |
| 68 | + // The new list is empty, so this is essentially being applied correctly. We just |
| 69 | + // put the candidates into the new storage item. |
| 70 | + CandidateList::<T>::put(&candidates); |
| 71 | + // 1 write for the new list |
| 72 | + weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1)); |
| 73 | + } else { |
| 74 | + // Oops, the runtime upgraded without the migration. There are new candidates in |
| 75 | + // `CandidateList`. So, let's just refund the old ones and assume they have already |
| 76 | + // started participating in the new system. |
| 77 | + for candidate in candidates { |
| 78 | + let err = T::Currency::unreserve(&candidate.who, candidate.deposit); |
| 79 | + if err > Zero::zero() { |
| 80 | + log::error!( |
| 81 | + target: LOG_TARGET, |
| 82 | + "{:?} balance was unable to be unreserved from {:?}", |
| 83 | + err, &candidate.who, |
| 84 | + ); |
| 85 | + } |
| 86 | + count.saturating_inc(); |
| 87 | + } |
| 88 | + weight.saturating_accrue( |
| 89 | + <<T as pallet_balances::Config>::WeightInfo as pallet_balances::WeightInfo>::force_unreserve().saturating_mul(count.into()), |
| 90 | + ); |
| 91 | + } |
| 92 | + |
| 93 | + log::info!( |
| 94 | + target: LOG_TARGET, |
| 95 | + "Unreserved locked bond of {} candidates, upgraded storage to version 2", |
| 96 | + count, |
| 97 | + ); |
| 98 | + |
| 99 | + weight.saturating_accrue(T::DbWeight::get().reads_writes(3, 2)); |
| 100 | + weight |
| 101 | + } |
| 102 | + |
| 103 | + #[cfg(feature = "try-runtime")] |
| 104 | + fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::DispatchError> { |
| 105 | + let number_of_candidates = Candidates::<T>::get().to_vec().len(); |
| 106 | + Ok((number_of_candidates as u32).encode()) |
| 107 | + } |
| 108 | + |
| 109 | + #[cfg(feature = "try-runtime")] |
| 110 | + fn post_upgrade(_number_of_candidates: Vec<u8>) -> Result<(), sp_runtime::DispatchError> { |
| 111 | + let new_number_of_candidates = Candidates::<T>::get().to_vec().len(); |
| 112 | + assert_eq!( |
| 113 | + new_number_of_candidates, 0 as usize, |
| 114 | + "after migration, the candidates map should be empty" |
| 115 | + ); |
| 116 | + Ok(()) |
| 117 | + } |
| 118 | + } |
| 119 | +} |
| 120 | + |
23 | 121 | /// Version 1 Migration |
24 | 122 | /// This migration ensures that any existing `Invulnerables` storage lists are sorted. |
25 | 123 | pub mod v1 { |
@@ -90,3 +188,136 @@ pub mod v1 { |
90 | 188 | } |
91 | 189 | } |
92 | 190 | } |
| 191 | + |
| 192 | +#[cfg(all(feature = "try-runtime", test))] |
| 193 | +mod tests { |
| 194 | + use super::*; |
| 195 | + use crate::{ |
| 196 | + migration::v2::Candidates, |
| 197 | + mock::{new_test_ext, Balances, Test}, |
| 198 | + }; |
| 199 | + use frame_support::{ |
| 200 | + traits::{Currency, ReservableCurrency, StorageVersion}, |
| 201 | + BoundedVec, |
| 202 | + }; |
| 203 | + use sp_runtime::traits::ConstU32; |
| 204 | + |
| 205 | + #[test] |
| 206 | + fn migrate_to_v2_with_new_candidates() { |
| 207 | + new_test_ext().execute_with(|| { |
| 208 | + let storage_version = StorageVersion::new(1); |
| 209 | + storage_version.put::<Pallet<Test>>(); |
| 210 | + |
| 211 | + let one = 1u64; |
| 212 | + let two = 2u64; |
| 213 | + let three = 3u64; |
| 214 | + let deposit = 10u64; |
| 215 | + |
| 216 | + // Set balance to 100 |
| 217 | + Balances::make_free_balance_be(&one, 100u64); |
| 218 | + Balances::make_free_balance_be(&two, 100u64); |
| 219 | + Balances::make_free_balance_be(&three, 100u64); |
| 220 | + |
| 221 | + // Reservations: 10 for the "old" candidacy and 10 for the "new" |
| 222 | + Balances::reserve(&one, 10u64).unwrap(); // old |
| 223 | + Balances::reserve(&two, 20u64).unwrap(); // old + new |
| 224 | + Balances::reserve(&three, 10u64).unwrap(); // new |
| 225 | + |
| 226 | + // Candidate info |
| 227 | + let candidate_one = CandidateInfo { who: one, deposit }; |
| 228 | + let candidate_two = CandidateInfo { who: two, deposit }; |
| 229 | + let candidate_three = CandidateInfo { who: three, deposit }; |
| 230 | + |
| 231 | + // Storage lists |
| 232 | + let bounded_candidates = |
| 233 | + BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![ |
| 234 | + candidate_one.clone(), |
| 235 | + candidate_two.clone(), |
| 236 | + ]) |
| 237 | + .expect("it works"); |
| 238 | + let bounded_candidate_list = |
| 239 | + BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![ |
| 240 | + candidate_two.clone(), |
| 241 | + candidate_three.clone(), |
| 242 | + ]) |
| 243 | + .expect("it works"); |
| 244 | + |
| 245 | + // Set storage |
| 246 | + Candidates::<Test>::put(bounded_candidates); |
| 247 | + CandidateList::<Test>::put(bounded_candidate_list.clone()); |
| 248 | + |
| 249 | + // Sanity check |
| 250 | + assert_eq!(Balances::free_balance(one), 90); |
| 251 | + assert_eq!(Balances::free_balance(two), 80); |
| 252 | + assert_eq!(Balances::free_balance(three), 90); |
| 253 | + |
| 254 | + // Run migration |
| 255 | + v2::MigrationToV2::<Test>::on_runtime_upgrade(); |
| 256 | + |
| 257 | + let new_storage_version = StorageVersion::get::<Pallet<Test>>(); |
| 258 | + assert_eq!(new_storage_version, 2); |
| 259 | + |
| 260 | + // 10 should have been unreserved from the old candidacy |
| 261 | + assert_eq!(Balances::free_balance(one), 100); |
| 262 | + assert_eq!(Balances::free_balance(two), 90); |
| 263 | + assert_eq!(Balances::free_balance(three), 90); |
| 264 | + // The storage item should be gone |
| 265 | + assert!(Candidates::<Test>::get().is_empty()); |
| 266 | + // The new storage item should be preserved |
| 267 | + assert_eq!(CandidateList::<Test>::get(), bounded_candidate_list); |
| 268 | + }); |
| 269 | + } |
| 270 | + |
| 271 | + #[test] |
| 272 | + fn migrate_to_v2_without_new_candidates() { |
| 273 | + new_test_ext().execute_with(|| { |
| 274 | + let storage_version = StorageVersion::new(1); |
| 275 | + storage_version.put::<Pallet<Test>>(); |
| 276 | + |
| 277 | + let one = 1u64; |
| 278 | + let two = 2u64; |
| 279 | + let deposit = 10u64; |
| 280 | + |
| 281 | + // Set balance to 100 |
| 282 | + Balances::make_free_balance_be(&one, 100u64); |
| 283 | + Balances::make_free_balance_be(&two, 100u64); |
| 284 | + |
| 285 | + // Reservations |
| 286 | + Balances::reserve(&one, 10u64).unwrap(); // old |
| 287 | + Balances::reserve(&two, 10u64).unwrap(); // old |
| 288 | + |
| 289 | + // Candidate info |
| 290 | + let candidate_one = CandidateInfo { who: one, deposit }; |
| 291 | + let candidate_two = CandidateInfo { who: two, deposit }; |
| 292 | + |
| 293 | + // Storage lists |
| 294 | + let bounded_candidates = |
| 295 | + BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![ |
| 296 | + candidate_one.clone(), |
| 297 | + candidate_two.clone(), |
| 298 | + ]) |
| 299 | + .expect("it works"); |
| 300 | + |
| 301 | + // Set storage |
| 302 | + Candidates::<Test>::put(bounded_candidates.clone()); |
| 303 | + |
| 304 | + // Sanity check |
| 305 | + assert_eq!(Balances::free_balance(one), 90); |
| 306 | + assert_eq!(Balances::free_balance(two), 90); |
| 307 | + |
| 308 | + // Run migration |
| 309 | + v2::MigrationToV2::<Test>::on_runtime_upgrade(); |
| 310 | + |
| 311 | + let new_storage_version = StorageVersion::get::<Pallet<Test>>(); |
| 312 | + assert_eq!(new_storage_version, 2); |
| 313 | + |
| 314 | + // Nothing changes deposit-wise |
| 315 | + assert_eq!(Balances::free_balance(one), 90); |
| 316 | + assert_eq!(Balances::free_balance(two), 90); |
| 317 | + // The storage item should be gone |
| 318 | + assert!(Candidates::<Test>::get().is_empty()); |
| 319 | + // The new storage item should have the info now |
| 320 | + assert_eq!(CandidateList::<Test>::get(), bounded_candidates); |
| 321 | + }); |
| 322 | + } |
| 323 | +} |
0 commit comments