Repository: paritytech/substrate
File: frame/elections-phragmen/src/migrations/unlock_and_unreserve_all_funds.rs (Rust)

Code before:
		let (account_deposited_sums, account_staked_sums, initial_reads_weight) =
			Self::get_account_deposited_and_staked_sums();

		// Deposited funds need to be unreserved.
		for (account, unreserve_amount) in account_deposited_sums {
			if unreserve_amount.is_zero() {
				continue
			}
			T::Currency::unreserve(&account, unreserve_amount);
		}

		// Staked funds need to be unlocked.
		for (account, amount) in account_staked_sums {
			if amount.is_zero() {
				continue
			}
			// Question for reviewers:
			// remove_lock triggers a few
			// `ERROR main runtime::system: Logic error: Unexpected underflow in reducing consumer`
			// errors. are they safe to ignore?
			//
			// I have verified that a non-zero lock always exists before remove_lock is called.
			T::Currency::remove_lock(T::PalletId::get(), &account);
		}

		// Question for reviewers: how do I know the weight of the unreserve & remove_lock calls?
		RocksDbWeight::get().reads_writes(1, 1) + initial_reads_weight
	}

	/// Performs post-upgrade sanity checks:
	///
	/// 1. All expected locks were removed after the migration.
	/// 2. There are no 'hanging' locks for this pallet in Balances.
	/// 3. The reserved balance for each account has been reduced by the expected amount.
	#[cfg(feature = "try-runtime")]
	fn post_upgrade(pre_migration_data_bytes: Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
		use sp_runtime::Saturating;

		let pre_migration_data = PreMigrationData::<T>::decode(&mut &pre_migration_data_bytes[..])
			.map_err(|_| "Failed to decode pre-migration data")?;

		// Get staked and deposited balances as reported by this pallet.
		let (account_deposited_sums, account_staked_sums, _) =
			Self::get_account_deposited_and_staked_sums();

		// Check that all expected locks were removed.
		account_staked_sums.iter().for_each(|(acc, _)| {
			assert_eq!(Self::get_actual_locked_amount(acc.clone()), None, "Lock still exists!");
		});

		// Extra sanity check to ensure there're no 'hanging' locks.

Review comment (suggestion):
> unreserve

There is a [benchmark](https://github.com/paritytech/substrate/blob/master/frame/balances/src/weights.rs#L57) for the `force_unreserve` extrinsic in the balances pallet which is basically does an account lookup and calls `unreserve`. Maybe the corresponding weight can be used here as an aprox?

> remove_lock

We may need to check the implementation in the balances pallet and estimate the expected r/w for removing a lock

Code after:
				.unwrap_or(&Zero::zero())
				// .min here to handle bugged deposits where actual_reserved_before is less than the
				// amount the pallet reports is reserved
				.min(&actual_reserved_before);
			let expected_reserved_after =
				actual_reserved_before.saturating_sub(expected_amount_deducted);
			assert!(
				actual_reserved_after == expected_reserved_after,
				"Reserved balance for {:?} is incorrect. actual before: {:?}, actual after, {:?}, expected deducted: {:?}",
				account,
				actual_reserved_before,
				actual_reserved_after,
				expected_amount_deducted,
			);
		}

		Ok(())
	}
}

#[cfg(all(feature = "try-runtime", test))]
mod test {
	use super::*;
	use crate::{
		tests::{ExtBuilder, Test},
		Candidates, Members, RunnersUp, SeatHolder, Voter, Voting,
	};
	use frame_support::{
		assert_ok,
		traits::{Currency, OnRuntimeUpgrade, ReservableCurrency, WithdrawReasons},
	};

	#[test]
	fn unreserve_works_for_candidate() {
		let candidate = 10;
		let deposit = 100;
		let initial_reserved = 15;
		let initial_balance = 100_000;
		ExtBuilder::default().build_and_execute(|| {
			// Set up initial state.
			<Test as crate::Config>::Currency::make_free_balance_be(&candidate, initial_balance);
			assert_ok!(<Test as crate::Config>::Currency::reserve(&candidate, initial_reserved));
			Candidates::<Test>::set(vec![(candidate, deposit)]);
			assert_ok!(<Test as crate::Config>::Currency::reserve(&candidate, deposit));

			// Sanity check: ensure initial Balance state was set up correctly.
			assert_eq!(
				<Test as crate::Config>::Currency::reserved_balance(&candidate),
				deposit + initial_reserved
			);