referrerpolicy=no-referrer-when-downgrade

pallet_democracy/migrations/
unlock_and_unreserve_all_funds.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! A migration that unreserves all deposit and unlocks all stake held in the context of this
19//! pallet.
20
21use 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
39/// The configuration for [`UnlockAndUnreserveAllFunds`].
40pub trait UnlockConfig: 'static {
41	/// The account ID used in the runtime.
42	type AccountId: Parameter + Ord;
43	/// The currency type used in the runtime.
44	///
45	/// Should match the currency type previously used for the pallet, if applicable.
46	type Currency: LockableCurrency<Self::AccountId> + ReservableCurrency<Self::AccountId>;
47	/// The name of the pallet as previously configured in
48	/// [`construct_runtime!`](frame_support::construct_runtime).
49	type PalletName: Get<&'static str>;
50	/// The maximum number of votes as configured previously in the runtime.
51	type MaxVotes: Get<u32>;
52	/// The maximum deposit as configured previously in the runtime.
53	type MaxDeposits: Get<u32>;
54	/// The DB weight as configured in the runtime to calculate the correct weight.
55	type DbWeight: Get<RuntimeDbWeight>;
56	/// The block number as configured in the runtime.
57	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
82/// A migration that unreserves all deposit and unlocks all stake held in the context of this
83/// pallet.
84///
85/// Useful to prevent funds from being locked up when the pallet is being deprecated.
86///
87/// The pallet should be made inoperable before this migration is run.
88///
89/// (See also [`RemovePallet`][frame_support::migrations::RemovePallet])
90pub struct UnlockAndUnreserveAllFunds<T: UnlockConfig>(core::marker::PhantomData<T>);
91
92impl<T: UnlockConfig> UnlockAndUnreserveAllFunds<T> {
93	/// Calculates and returns the total amounts reserved by each account by this pallet, and all
94	/// accounts with locks in the context of this pallet.
95	///
96	/// There is no need to return the amount locked, because the entire lock is removed (always
97	/// should be zero post-migration). We need to return the amounts reserved to check that the
98	/// reserved amount is deducted correctly.
99	///
100	/// # Returns
101	///
102	/// This function returns a tuple of two `BTreeMap` collections and the weight of the reads:
103	///
104	/// * `BTreeMap<T::AccountId, BalanceOf<T>>`: Map of account IDs to their respective total
105	///   reserved balance by this pallet
106	/// * `BTreeMap<T::AccountId, BalanceOf<T>>`: Map of account IDs to their respective total
107	///   locked balance by this pallet
108	/// * `frame_support::weights::Weight`: the weight consumed by this call.
109	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		// Get all deposits (reserved).
117		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				// Count the number of deposits
121				deposit_of_len.saturating_inc();
122
123				// Track the total number of vec entries to calculate the weight of the reads.
124				total_voting_vec_entries.saturating_accrue(accounts.len() as u64);
125
126				// Create a vec of tuples where each account is associated with the given balance
127				accounts.into_iter().map(|account| (account, balance)).collect::<Vec<_>>()
128			})
129			.fold(BTreeMap::new(), |mut acc, (account, balance)| {
130				// Add the balance to the account's existing balance in the accumulator
131				acc.entry(account.clone()).or_insert(Zero::zero()).saturating_accrue(balance);
132				acc
133			});
134
135		// Voter accounts have amounts locked.
136		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					// Max items in a Voting enum is MaxVotes + 5
147					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	/// Collects pre-migration data useful for validating the migration was successful, and also
160	/// checks the integrity of deposited and reserved balances.
161	///
162	/// Steps:
163	/// 1. Gets the deposited balances for each account stored in this pallet.
164	/// 2. Collects actual pre-migration reserved balances for each account.
165	/// 3. Checks the integrity of the deposited balances.
166	/// 4. Prints summary statistics about the state to be migrated.
167	/// 5. Encodes and returns pre-migration data to be used in post_upgrade.
168	///
169	/// Fails with a `TryRuntimeError` if somehow the amount reserved by this pallet is greater than
170	/// the actual total reserved amount for any accounts.
171	#[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		// Get staked and deposited balances as reported by this pallet.
177		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		// Total deposited for each account *should* be less than or equal to the total reserved,
190		// however this does not hold for all cases due to bugs in the reserve logic of this pallet.
191		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	/// Executes the migration.
221	///
222	/// Steps:
223	/// 1. Retrieves the deposit and accounts with locks for the pallet.
224	/// 2. Unreserves the deposited funds for each account.
225	/// 3. Unlocks the staked funds for each account.
226	fn on_runtime_upgrade() -> frame_support::weights::Weight {
227		// Get staked and deposited balances as reported by this pallet.
228		let (account_deposits, account_stakes, initial_reads) =
229			Self::get_account_deposits_and_locks();
230
231		// Deposited funds need to be unreserved.
232		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		// Staked funds need to be unlocked.
241		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	/// Performs post-upgrade sanity checks:
254	///
255	/// 1. No locks remain for this pallet in Balances.
256	/// 2. The reserved balance for each account has been reduced by the expected amount.
257	#[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		// Get staked and deposited balances as reported by this pallet.
268		let (account_deposits, _, _) = Self::get_account_deposits_and_locks();
269
270		// Check that the reserved balance is reduced by the expected deposited amount.
271		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			// Set up initial state.
333			<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			// Sanity check: ensure initial reserved balance was set correctly.
351			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			// Run the migration.
361			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 the reserved balance was reduced by the expected amount.
367			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			// Set up initial state.
386			<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			// Sanity check: ensure initial Balance state was set up correctly.
404			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			// Run the migration.
415			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 the voter lock was removed
421			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}