pallet_migrations/
migrations.rs

1// Copyright (C) Parity Technologies (UK) Ltd.
2// SPDX-License-Identifier: Apache-2.0
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// 	http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Generic multi block migrations not specific to any pallet.
17
18use crate::{weights::WeightInfo, Config};
19use codec::Encode;
20use core::marker::PhantomData;
21use frame_support::{
22	migrations::{SteppedMigration, SteppedMigrationError, StoreInCodeStorageVersion},
23	traits::{GetStorageVersion, PalletInfoAccess},
24	weights::WeightMeter,
25};
26use sp_core::Get;
27use sp_crypto_hashing::twox_128;
28use sp_io::{storage::clear_prefix, KillStorageResult};
29use sp_runtime::SaturatedConversion;
30
31/// Remove all of a pallet's state and re-initializes it to the current in-code storage version.
32///
33/// It uses the multi block migration frame. Hence it is safe to use even on
34/// pallets that contain a lot of storage.
35///
36/// # Parameters
37///
38/// - T: The runtime. Used to access the weight definition.
39/// - P: The pallet to resetted as defined in construct runtime
40///
41/// # Note
42///
43/// If your pallet does rely of some state in genesis you need to take care of that
44/// separately. This migration only sets the storage version after wiping.
45pub struct ResetPallet<T, P>(PhantomData<(T, P)>);
46
47impl<T, P> ResetPallet<T, P>
48where
49	P: PalletInfoAccess,
50{
51	#[cfg(feature = "try-runtime")]
52	fn num_keys() -> u64 {
53		let prefix = P::name_hash().to_vec();
54		crate::storage::KeyPrefixIterator::new(prefix.clone(), prefix, |_| Ok(())).count() as _
55	}
56}
57
58impl<T, P, V> SteppedMigration for ResetPallet<T, P>
59where
60	T: Config,
61	P: PalletInfoAccess + GetStorageVersion<InCodeStorageVersion = V>,
62	V: StoreInCodeStorageVersion<P>,
63{
64	type Cursor = bool;
65	type Identifier = [u8; 16];
66
67	fn id() -> Self::Identifier {
68		("RemovePallet::", P::name()).using_encoded(twox_128)
69	}
70
71	fn step(
72		cursor: Option<Self::Cursor>,
73		meter: &mut WeightMeter,
74	) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
75		// we write the storage version in a seperate block
76		if cursor.unwrap_or(false) {
77			let required = T::DbWeight::get().writes(1);
78			meter
79				.try_consume(required)
80				.map_err(|_| SteppedMigrationError::InsufficientWeight { required })?;
81			V::store_in_code_storage_version();
82			return Ok(None);
83		}
84
85		let base_weight = T::WeightInfo::reset_pallet_migration(0);
86		let weight_per_key = T::WeightInfo::reset_pallet_migration(1).saturating_sub(base_weight);
87		let key_budget = meter
88			.remaining()
89			.saturating_sub(base_weight)
90			.checked_div_per_component(&weight_per_key)
91			.unwrap_or_default()
92			.saturated_into();
93
94		if key_budget == 0 {
95			return Err(SteppedMigrationError::InsufficientWeight {
96				required: T::WeightInfo::reset_pallet_migration(1),
97			});
98		}
99
100		let (keys_removed, is_done) = match clear_prefix(&P::name_hash(), Some(key_budget)) {
101			KillStorageResult::AllRemoved(value) => (value, true),
102			KillStorageResult::SomeRemaining(value) => (value, false),
103		};
104
105		meter.consume(T::WeightInfo::reset_pallet_migration(keys_removed));
106
107		Ok(Some(is_done))
108	}
109
110	#[cfg(feature = "try-runtime")]
111	fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
112		let num_keys: u64 = Self::num_keys();
113		log::info!("ResetPallet<{}>: Trying to remove {num_keys} keys.", P::name());
114		Ok(num_keys.encode())
115	}
116
117	#[cfg(feature = "try-runtime")]
118	fn post_upgrade(state: alloc::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
119		use codec::Decode;
120		let keys_before = u64::decode(&mut state.as_ref()).expect("We encoded as u64 above; qed");
121		let keys_now = Self::num_keys();
122		log::info!("ResetPallet<{}>: Keys remaining after migration: {keys_now}", P::name());
123
124		if keys_before <= keys_now {
125			log::error!("ResetPallet<{}>: Did not remove any keys.", P::name());
126			Err("ResetPallet failed")?;
127		}
128
129		if keys_now != 1 {
130			log::error!("ResetPallet<{}>: Should have a single key after reset", P::name());
131			Err("ResetPallet failed")?;
132		}
133
134		Ok(())
135	}
136}
pallet_migrations/migrations.rs

pallet_migrations/
migrations.rs
