1
// Copyright 2019-2025 PureStake Inc.
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
//! # Common Moonbeam Migrations
18

            
19
use core::marker::PhantomData;
20
use frame_support::migrations::SteppedMigrationError;
21
use frame_support::traits::PalletInfoAccess;
22
use frame_support::weights::WeightMeter;
23
use frame_support::{migrations::SteppedMigration, parameter_types};
24
use pallet_migrations::WeightInfo;
25
use parity_scale_codec::Encode;
26
use sp_core::{twox_128, Get};
27
use sp_io::{storage::clear_prefix, KillStorageResult};
28
use 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 name to be removed 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.
44
pub struct RemovePallet<T, P>(PhantomData<(T, P)>);
45

            
46
impl<T, P> RemovePallet<T, P>
47
where
48
	P: Get<&'static str>,
49
{
50
	#[cfg(feature = "try-runtime")]
51
	fn num_keys() -> u64 {
52
		let prefix = twox_128(P::get().as_bytes()).to_vec();
53
		frame_support::storage::KeyPrefixIterator::new(prefix.clone(), prefix, |_| Ok(())).count()
54
			as _
55
	}
56
}
57

            
58
impl<T, P> SteppedMigration for RemovePallet<T, P>
59
where
60
	T: pallet_migrations::Config,
61
	P: Get<&'static str>,
62
{
63
	type Cursor = bool;
64
	type Identifier = [u8; 16];
65

            
66
6
	fn id() -> Self::Identifier {
67
6
		("RemovePallet::", P::get()).using_encoded(twox_128)
68
6
	}
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
			return Ok(None);
81
		}
82

            
83
		let base_weight = T::WeightInfo::reset_pallet_migration(0);
84
		let weight_per_key = T::WeightInfo::reset_pallet_migration(1).saturating_sub(base_weight);
85
		let key_budget = meter
86
			.remaining()
87
			.saturating_sub(base_weight)
88
			.checked_div_per_component(&weight_per_key)
89
			.unwrap_or_default()
90
			.saturated_into();
91

            
92
		if key_budget == 0 {
93
			return Err(SteppedMigrationError::InsufficientWeight {
94
				required: T::WeightInfo::reset_pallet_migration(1),
95
			});
96
		}
97

            
98
		let (keys_removed, is_done) =
99
			match clear_prefix(&twox_128(P::get().as_bytes()), 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<sp_std::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
111
		let num_keys: u64 = Self::num_keys();
112
		log::info!(
113
			"RemovePallet<{}>: Trying to remove {num_keys} keys.",
114
			P::get()
115
		);
116
		Ok(num_keys.encode())
117
	}
118

            
119
	#[cfg(feature = "try-runtime")]
120
	fn post_upgrade(state: sp_std::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
121
		use parity_scale_codec::Decode;
122
		let keys_before = u64::decode(&mut state.as_ref()).expect("We encoded as u64 above; qed");
123
		let keys_now = Self::num_keys();
124
		log::info!(
125
			"RemovePallet<{}>: Keys remaining after migration: {keys_now}",
126
			P::get()
127
		);
128

            
129
		if keys_before <= keys_now {
130
			log::error!("RemovePallet<{}>: Did not remove any keys.", P::get());
131
			Err("RemovePallet failed")?;
132
		}
133

            
134
		if keys_now != 1 {
135
			log::error!("RemovePallet<{}>: Should have a single key after", P::get());
136
			Err("RemovePallet failed")?;
137
		}
138

            
139
		Ok(())
140
	}
141
}
142

            
143
/// Reset the pallet's storage.
144
///
145
/// It uses the multi block migration frame. Hence it is safe to use even on
146
/// pallets that contain a lot of storage.
147
///
148
/// # Parameters
149
///
150
/// - T: The runtime. Used to access the weight definition.
151
/// - Pallet: The pallet to be resetted
152
/// - Storage: The pallet storage to be resetted
153
pub struct ResetStorage<T, Pallet, Storage>(PhantomData<(T, Pallet, Storage)>);
154

            
155
impl<T, Pallet, Storage> ResetStorage<T, Pallet, Storage>
156
where
157
	Pallet: PalletInfoAccess,
158
	Storage: Get<&'static str>,
159
{
160
	#[cfg(feature = "try-runtime")]
161
	fn num_keys() -> u64 {
162
		let storage_prefix = frame_support::storage::storage_prefix(
163
			Pallet::name().as_bytes(),
164
			Storage::get().as_bytes(),
165
		)
166
		.to_vec();
167
		frame_support::storage::KeyPrefixIterator::new(
168
			storage_prefix.clone(),
169
			storage_prefix,
170
			|_| Ok(()),
171
		)
172
		.count() as _
173
	}
174
}
175

            
176
impl<T, Pallet, Storage> SteppedMigration for ResetStorage<T, Pallet, Storage>
177
where
178
	T: pallet_migrations::Config,
179
	Pallet: PalletInfoAccess,
180
	Storage: Get<&'static str>,
181
{
182
	type Cursor = bool;
183
	type Identifier = [u8; 16];
184

            
185
6
	fn id() -> Self::Identifier {
186
6
		("ResetStorage", Pallet::name(), Storage::get()).using_encoded(twox_128)
187
6
	}
188

            
189
	fn step(
190
		cursor: Option<Self::Cursor>,
191
		meter: &mut WeightMeter,
192
	) -> Result<Option<Self::Cursor>, SteppedMigrationError> {
193
		// we write the storage version in a separate block
194
		if cursor.unwrap_or(false) {
195
			let required = T::DbWeight::get().writes(1);
196
			meter
197
				.try_consume(required)
198
				.map_err(|_| SteppedMigrationError::InsufficientWeight { required })?;
199
			return Ok(None);
200
		}
201

            
202
		let base_weight = T::WeightInfo::reset_pallet_migration(0);
203
		let weight_per_key = T::WeightInfo::reset_pallet_migration(1).saturating_sub(base_weight);
204
		let key_budget = meter
205
			.remaining()
206
			.saturating_sub(base_weight)
207
			.checked_div_per_component(&weight_per_key)
208
			.unwrap_or_default()
209
			.saturated_into();
210

            
211
		if key_budget == 0 {
212
			return Err(SteppedMigrationError::InsufficientWeight {
213
				required: T::WeightInfo::reset_pallet_migration(1),
214
			});
215
		}
216

            
217
		let storage_prefix = frame_support::storage::storage_prefix(
218
			Pallet::name().as_bytes(),
219
			Storage::get().as_bytes(),
220
		);
221
		let (keys_removed, is_done) = match clear_prefix(&storage_prefix, Some(key_budget)) {
222
			KillStorageResult::AllRemoved(value) => (value, true),
223
			KillStorageResult::SomeRemaining(value) => (value, false),
224
		};
225

            
226
		meter.consume(T::WeightInfo::reset_pallet_migration(keys_removed));
227

            
228
		Ok(Some(is_done))
229
	}
230

            
231
	#[cfg(feature = "try-runtime")]
232
	fn pre_upgrade() -> Result<sp_std::vec::Vec<u8>, sp_runtime::TryRuntimeError> {
233
		let num_keys: u64 = Self::num_keys();
234
		log::info!(
235
			"ResetStorage<{}, {}>: Trying to remove {num_keys} keys.",
236
			Pallet::name(),
237
			Storage::get()
238
		);
239
		Ok(num_keys.encode())
240
	}
241

            
242
	#[cfg(feature = "try-runtime")]
243
	fn post_upgrade(state: sp_std::vec::Vec<u8>) -> Result<(), sp_runtime::TryRuntimeError> {
244
		use parity_scale_codec::Decode;
245
		let keys_before = u64::decode(&mut state.as_ref()).expect("We encoded as u64 above; qed");
246
		let keys_now = Self::num_keys();
247
		log::info!(
248
			"ResetStorage<{}, {}>: Keys remaining after migration: {keys_now}",
249
			Pallet::name(),
250
			Storage::get()
251
		);
252

            
253
		if keys_before <= keys_now {
254
			log::error!(
255
				"ResetStorage<{}, {}>: Did not remove any keys.",
256
				Pallet::name(),
257
				Storage::get()
258
			);
259
			Err("ResetStorage failed")?;
260
		}
261

            
262
		if keys_now != 1 {
263
			log::error!(
264
				"ResetStorage<{}, {}>: Should have a single key after reset",
265
				Pallet::name(),
266
				Storage::get()
267
			);
268
			Err("ResetStorage failed")?;
269
		}
270

            
271
		Ok(())
272
	}
273
}
274

            
275
/// Unreleased migrations. Add new ones here:
276
pub type UnreleasedSingleBlockMigrations = ();
277

            
278
/// Migrations/checks that do not need to be versioned and can run on every update.
279
pub type PermanentSingleBlockMigrations<Runtime> =
280
	(pallet_xcm::migration::MigrateToLatestXcmVersion<Runtime>,);
281

            
282
/// All migrations that will run on the next runtime upgrade.
283
pub type SingleBlockMigrations<Runtime> = (
284
	UnreleasedSingleBlockMigrations,
285
	PermanentSingleBlockMigrations<Runtime>,
286
);
287

            
288
parameter_types! {
289
	pub const AssetManagerPalletName: &'static str = "AssetManager";
290
	pub const AssetsPalletName: &'static str = "Assets";
291

            
292
	pub const MigratedCandidatesStorageName: &'static str = "MigratedCandidates";
293
	pub const MigratedDelegatorsStorageName: &'static str = "MigratedDelegators";
294
}
295

            
296
/// List of common multiblock migrations to be executed by the pallet_multiblock_migrations.
297
/// The migrations listed here are common to every moonbeam runtime.
298
pub type MultiBlockMigrations<Runtime> = (
299
	RemovePallet<Runtime, AssetManagerPalletName>,
300
	RemovePallet<Runtime, AssetsPalletName>,
301
	ResetStorage<Runtime, pallet_parachain_staking::Pallet<Runtime>, MigratedCandidatesStorageName>,
302
	ResetStorage<Runtime, pallet_parachain_staking::Pallet<Runtime>, MigratedDelegatorsStorageName>,
303
);