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::{twox_128, Get};
27use sp_io::{storage::clear_prefix, KillStorageResult};
28use sp_runtime::SaturatedConversion;
29
30/// Remove all of a pallet's state and re-initializes it to the current in-code storage version.
31///
32/// It uses the multi block migration frame. Hence it is safe to use even on
33/// pallets that contain a lot of storage.
34///
35/// # Parameters
36///
37/// - T: The runtime. Used to access the weight definition.
38/// - P: The pallet to resetted as defined in construct runtime
39///
40/// # Note
41///
42/// If your pallet does rely of some state in genesis you need to take care of that
43/// separately. This migration only sets the storage version after wiping.
44pub struct ResetPallet<T, P>(PhantomData<(T, P)>);
45
46impl<T, P> ResetPallet<T, P>
47where
48	P: PalletInfoAccess,
49{
50	#[cfg(feature = "try-runtime")]
51	fn num_keys() -> u64 {
52		let prefix = P::name_hash().to_vec();
53		crate::storage::KeyPrefixIterator::new(prefix.clone(), prefix, |_| Ok(())).count() as _
54	}
55}
56
57impl<T, P, V> SteppedMigration for ResetPallet<T, P>
58where
59	T: Config,
60	P: PalletInfoAccess + GetStorageVersion<InCodeStorageVersion = V>,
61	V: StoreInCodeStorageVersion<P>,
62{
63	type Cursor = bool;
64	type Identifier = [u8; 16];
65
66	fn id() -> Self::Identifier {
67		("RemovePallet::", P::name()).using_encoded(twox_128)
68	}
69
70	fn step(
71		cursor: Option<Self::Cursor>,
72		meter: &mut WeightMeter,
73	) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
74		// we write the storage version in a seperate block
75		if cursor.unwrap_or(false) {
76			let required = T::DbWeight::get().writes(1);
77			meter
78				.try_consume(required)
79				.map_err(|_| SteppedMigrationError::InsufficientWeight { required })?;
80			V::store_in_code_storage_version();
81			return Ok(None);
82		}
83
84		let base_weight = T::WeightInfo::reset_pallet_migration(0);
85		let weight_per_key = T::WeightInfo::reset_pallet_migration(1).saturating_sub(base_weight);
86		let key_budget = meter
87			.remaining()
88			.saturating_sub(base_weight)
89			.checked_div_per_component(&weight_per_key)
90			.unwrap_or_default()
91			.saturated_into();
92
93		if key_budget == 0 {
94			return Err(SteppedMigrationError::InsufficientWeight {
95				required: T::WeightInfo::reset_pallet_migration(1),
96			})
97		}
98
99		let (keys_removed, is_done) = match clear_prefix(&P::name_hash(), Some(key_budget)) {
100			KillStorageResult::AllRemoved(value) => (value, true),
101			KillStorageResult::SomeRemaining(value) => (value, false),
102		};
103
104		meter.consume(T::WeightInfo::reset_pallet_migration(keys_removed));
105
106		Ok(Some(is_done))
107	}
108
109	#[cfg(feature = "try-runtime")]
110	fn pre_upgrade() -> Result<alloc::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
111		let num_keys: u64 = Self::num_keys();
112		log::info!("ResetPallet<{}>: Trying to remove {num_keys} keys.", P::name());
113		Ok(num_keys.encode())
114	}
115
116	#[cfg(feature = "try-runtime")]
117	fn post_upgrade(state: alloc::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
118		use codec::Decode;
119		let keys_before = u64::decode(&mut state.as_ref()).expect("We encoded as u64 above; qed");
120		let keys_now = Self::num_keys();
121		log::info!("ResetPallet<{}>: Keys remaining after migration: {keys_now}", P::name());
122
123		if keys_before <= keys_now {
124			log::error!("ResetPallet<{}>: Did not remove any keys.", P::name());
125			Err("ResetPallet failed")?;
126		}
127
128		if keys_now != 1 {
129			log::error!("ResetPallet<{}>: Should have a single key after reset", P::name());
130			Err("ResetPallet failed")?;
131		}
132
133		Ok(())
134	}
135}
pallet_migrations/migrations.rs

pallet_migrations/
migrations.rs
