pallet_democracy/migrations/
unlock_and_unreserve_all_funds.rs1use crate::{PropIndex, Voting, DEMOCRACY_ID};
22use alloc::{collections::btree_map::BTreeMap, vec::Vec};
23use core::iter::Sum;
24use frame_support::{
25 pallet_prelude::ValueQuery,
26 storage_alias,
27 traits::{Currency, LockableCurrency, OnRuntimeUpgrade, ReservableCurrency},
28 weights::RuntimeDbWeight,
29 Parameter, Twox64Concat,
30};
31use sp_core::Get;
32use sp_runtime::{traits::Zero, BoundedVec, Saturating};
33
34const LOG_TARGET: &str = "runtime::democracy::migrations::unlock_and_unreserve_all_funds";
35
36type BalanceOf<T> =
37 <<T as UnlockConfig>::Currency as Currency<<T as UnlockConfig>::AccountId>>::Balance;
38
39pub trait UnlockConfig: 'static {
41 type AccountId: Parameter + Ord;
43 type Currency: LockableCurrency<Self::AccountId> + ReservableCurrency<Self::AccountId>;
47 type PalletName: Get<&'static str>;
50 type MaxVotes: Get<u32>;
52 type MaxDeposits: Get<u32>;
54 type DbWeight: Get<RuntimeDbWeight>;
56 type BlockNumber: Parameter + Zero + Copy + Ord;
58}
59
60#[storage_alias(dynamic)]
61type DepositOf<T: UnlockConfig> = StorageMap<
62 <T as UnlockConfig>::PalletName,
63 Twox64Concat,
64 PropIndex,
65 (BoundedVec<<T as UnlockConfig>::AccountId, <T as UnlockConfig>::MaxDeposits>, BalanceOf<T>),
66>;
67
68#[storage_alias(dynamic)]
69type VotingOf<T: UnlockConfig> = StorageMap<
70 <T as UnlockConfig>::PalletName,
71 Twox64Concat,
72 <T as UnlockConfig>::AccountId,
73 Voting<
74 BalanceOf<T>,
75 <T as UnlockConfig>::AccountId,
76 <T as UnlockConfig>::BlockNumber,
77 <T as UnlockConfig>::MaxVotes,
78 >,
79 ValueQuery,
80>;
81
82pub struct UnlockAndUnreserveAllFunds<T: UnlockConfig>(core::marker::PhantomData<T>);
91
92impl<T: UnlockConfig> UnlockAndUnreserveAllFunds<T> {
93 fn get_account_deposits_and_locks() -> (
110 BTreeMap<T::AccountId, BalanceOf<T>>,
111 BTreeMap<T::AccountId, BalanceOf<T>>,
112 frame_support::weights::Weight,
113 ) {
114 let mut deposit_of_len = 0;
115
116 let mut total_voting_vec_entries: u64 = 0;
118 let account_deposits: BTreeMap<T::AccountId, BalanceOf<T>> = DepositOf::<T>::iter()
119 .flat_map(|(_prop_index, (accounts, balance))| {
120 deposit_of_len.saturating_inc();
122
123 total_voting_vec_entries.saturating_accrue(accounts.len() as u64);
125
126 accounts.into_iter().map(|account| (account, balance)).collect::<Vec<_>>()
128 })
129 .fold(BTreeMap::new(), |mut acc, (account, balance)| {
130 acc.entry(account.clone()).or_insert(Zero::zero()).saturating_accrue(balance);
132 acc
133 });
134
135 let account_stakes: BTreeMap<T::AccountId, BalanceOf<T>> = VotingOf::<T>::iter()
137 .map(|(account_id, voting)| (account_id, voting.locked_balance()))
138 .collect();
139 let voting_of_len = account_stakes.len() as u64;
140
141 (
142 account_deposits,
143 account_stakes,
144 T::DbWeight::get().reads(
145 deposit_of_len.saturating_add(voting_of_len).saturating_add(
146 total_voting_vec_entries
148 .saturating_mul(T::MaxVotes::get().saturating_add(5) as u64),
149 ),
150 ),
151 )
152 }
153}
154
155impl<T: UnlockConfig> OnRuntimeUpgrade for UnlockAndUnreserveAllFunds<T>
156where
157 BalanceOf<T>: Sum,
158{
159 #[cfg(feature = "try-runtime")]
172 fn pre_upgrade() -> Result<Vec<u8>, sp_runtime::TryRuntimeError> {
173 use alloc::collections::btree_set::BTreeSet;
174 use codec::Encode;
175
176 let (account_deposits, account_locks, _) = Self::get_account_deposits_and_locks();
178
179 let all_accounts = account_deposits
180 .keys()
181 .chain(account_locks.keys())
182 .cloned()
183 .collect::<BTreeSet<_>>();
184 let account_reserved_before: BTreeMap<T::AccountId, BalanceOf<T>> = account_deposits
185 .keys()
186 .map(|account| (account.clone(), T::Currency::reserved_balance(&account)))
187 .collect();
188
189 let bugged_deposits = all_accounts
192 .iter()
193 .filter(|account| {
194 account_deposits.get(&account).unwrap_or(&Zero::zero()) >
195 account_reserved_before.get(&account).unwrap_or(&Zero::zero())
196 })
197 .count();
198
199 let total_deposits_to_unreserve =
200 account_deposits.clone().into_values().sum::<BalanceOf<T>>();
201 let total_stake_to_unlock = account_locks.clone().into_values().sum::<BalanceOf<T>>();
202
203 log::info!(target: LOG_TARGET, "Total accounts: {:?}", all_accounts.len());
204 log::info!(target: LOG_TARGET, "Total stake to unlock: {:?}", total_stake_to_unlock);
205 log::info!(
206 target: LOG_TARGET,
207 "Total deposit to unreserve: {:?}",
208 total_deposits_to_unreserve
209 );
210 log::info!(
211 target: LOG_TARGET,
212 "Bugged deposits: {}/{}",
213 bugged_deposits,
214 account_deposits.len()
215 );
216
217 Ok(account_reserved_before.encode())
218 }
219
220 fn on_runtime_upgrade() -> frame_support::weights::Weight {
227 let (account_deposits, account_stakes, initial_reads) =
229 Self::get_account_deposits_and_locks();
230
231 for (account, unreserve_amount) in account_deposits.iter() {
233 if unreserve_amount.is_zero() {
234 log::warn!(target: LOG_TARGET, "Unexpected zero amount to unreserve!");
235 continue
236 }
237 T::Currency::unreserve(&account, *unreserve_amount);
238 }
239
240 for account in account_stakes.keys() {
242 T::Currency::remove_lock(DEMOCRACY_ID, account);
243 }
244
245 T::DbWeight::get()
246 .reads_writes(
247 account_stakes.len().saturating_add(account_deposits.len()) as u64,
248 account_stakes.len().saturating_add(account_deposits.len()) as u64,
249 )
250 .saturating_add(initial_reads)
251 }
252
253 #[cfg(feature = "try-runtime")]
258 fn post_upgrade(
259 account_reserved_before_bytes: Vec<u8>,
260 ) -> Result<(), sp_runtime::TryRuntimeError> {
261 use codec::Decode;
262
263 let account_reserved_before =
264 BTreeMap::<T::AccountId, BalanceOf<T>>::decode(&mut &account_reserved_before_bytes[..])
265 .map_err(|_| "Failed to decode account_reserved_before_bytes")?;
266
267 let (account_deposits, _, _) = Self::get_account_deposits_and_locks();
269
270 for (account, actual_reserved_before) in account_reserved_before {
272 let actual_reserved_after = T::Currency::reserved_balance(&account);
273 let expected_amount_deducted = *account_deposits
274 .get(&account)
275 .expect("account deposit must exist to be in pre_migration_data, qed");
276 let expected_reserved_after =
277 actual_reserved_before.saturating_sub(expected_amount_deducted);
278 assert!(
279 actual_reserved_after == expected_reserved_after,
280 "Reserved balance for {:?} is incorrect. actual before: {:?}, actual after, {:?}, expected deducted: {:?}",
281 account,
282 actual_reserved_before,
283 actual_reserved_after,
284 expected_amount_deducted,
285 );
286 }
287
288 Ok(())
289 }
290}
291
292#[cfg(all(feature = "try-runtime", test))]
293mod test {
294 use super::*;
295 use crate::{
296 tests::{new_test_ext, Balances, Test},
297 DepositOf, Voting, VotingOf,
298 };
299 use frame_support::{
300 assert_ok, parameter_types,
301 traits::{Currency, OnRuntimeUpgrade, ReservableCurrency, WithdrawReasons},
302 BoundedVec,
303 };
304 use frame_system::pallet_prelude::BlockNumberFor;
305 use sp_core::ConstU32;
306
307 parameter_types! {
308 const PalletName: &'static str = "Democracy";
309 }
310
311 struct UnlockConfigImpl;
312
313 impl super::UnlockConfig for UnlockConfigImpl {
314 type Currency = Balances;
315 type MaxVotes = ConstU32<100>;
316 type MaxDeposits = ConstU32<1000>;
317 type AccountId = u64;
318 type BlockNumber = BlockNumberFor<Test>;
319 type DbWeight = ();
320 type PalletName = PalletName;
321 }
322
323 #[test]
324 fn unreserve_works_for_depositor() {
325 let depositor_0 = 10;
326 let depositor_1 = 11;
327 let deposit = 25;
328 let depositor_0_initial_reserved = 0;
329 let depositor_1_initial_reserved = 15;
330 let initial_balance = 100_000;
331 new_test_ext().execute_with(|| {
332 <Test as crate::Config>::Currency::make_free_balance_be(&depositor_0, initial_balance);
334 <Test as crate::Config>::Currency::make_free_balance_be(&depositor_1, initial_balance);
335 assert_ok!(<Test as crate::Config>::Currency::reserve(
336 &depositor_0,
337 depositor_0_initial_reserved + deposit
338 ));
339 assert_ok!(<Test as crate::Config>::Currency::reserve(
340 &depositor_1,
341 depositor_1_initial_reserved + deposit
342 ));
343 let depositors =
344 BoundedVec::<_, <Test as crate::Config>::MaxDeposits>::truncate_from(vec![
345 depositor_0,
346 depositor_1,
347 ]);
348 DepositOf::<Test>::insert(0, (depositors, deposit));
349
350 assert_eq!(
352 <Test as crate::Config>::Currency::reserved_balance(&depositor_0),
353 depositor_0_initial_reserved + deposit
354 );
355 assert_eq!(
356 <Test as crate::Config>::Currency::reserved_balance(&depositor_1),
357 depositor_1_initial_reserved + deposit
358 );
359
360 let bytes = UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::pre_upgrade()
362 .unwrap_or_else(|e| panic!("pre_upgrade failed: {:?}", e));
363 UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::on_runtime_upgrade();
364 assert_ok!(UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::post_upgrade(bytes));
365
366 assert_eq!(
368 <Test as crate::Config>::Currency::reserved_balance(&depositor_0),
369 depositor_0_initial_reserved
370 );
371 assert_eq!(
372 <Test as crate::Config>::Currency::reserved_balance(&depositor_1),
373 depositor_1_initial_reserved
374 );
375 });
376 }
377
378 #[test]
379 fn unlock_works_for_voter() {
380 let voter = 10;
381 let stake = 25;
382 let initial_locks = vec![(b"somethin", 10)];
383 let initial_balance = 100_000;
384 new_test_ext().execute_with(|| {
385 <Test as crate::Config>::Currency::make_free_balance_be(&voter, initial_balance);
387 for lock in initial_locks.clone() {
388 <Test as crate::Config>::Currency::set_lock(
389 *lock.0,
390 &voter,
391 lock.1,
392 WithdrawReasons::all(),
393 );
394 }
395 VotingOf::<Test>::insert(voter, Voting::default());
396 <Test as crate::Config>::Currency::set_lock(
397 DEMOCRACY_ID,
398 &voter,
399 stake,
400 WithdrawReasons::all(),
401 );
402
403 let mut voter_all_locks = initial_locks.clone();
405 voter_all_locks.push((&DEMOCRACY_ID, stake));
406 assert_eq!(
407 <Test as crate::Config>::Currency::locks(&voter)
408 .iter()
409 .map(|lock| (&lock.id, lock.amount))
410 .collect::<Vec<_>>(),
411 voter_all_locks
412 );
413
414 let bytes = UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::pre_upgrade()
416 .unwrap_or_else(|e| panic!("pre_upgrade failed: {:?}", e));
417 UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::on_runtime_upgrade();
418 assert_ok!(UnlockAndUnreserveAllFunds::<UnlockConfigImpl>::post_upgrade(bytes));
419
420 assert_eq!(
422 <Test as crate::Config>::Currency::locks(&voter)
423 .iter()
424 .map(|lock| (&lock.id, lock.amount))
425 .collect::<Vec<_>>(),
426 initial_locks
427 );
428 });
429 }
430}