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

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

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

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

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

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

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

            
140
		Ok(())
141
	}
142
}
143

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

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

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

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

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

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

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

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

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

            
229
		Ok(Some(is_done))
230
	}
231

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

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

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

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

            
272
		Ok(())
273
	}
274
}
275

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

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

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

            
289
parameter_types! {
290
	pub const MigratedCandidatesStorageName: &'static str = "MigratedCandidates";
291
	pub const MigratedDelegatorsStorageName: &'static str = "MigratedDelegators";
292
}
293

            
294
/// List of common multiblock migrations to be executed by the pallet-migrations pallet.
295
/// The migrations listed here are common to every moonbeam runtime.
296
pub type MultiBlockMigrations<Runtime> = (
297
	ResetStorage<Runtime, pallet_parachain_staking::Pallet<Runtime>, MigratedCandidatesStorageName>,
298
	ResetStorage<Runtime, pallet_parachain_staking::Pallet<Runtime>, MigratedDelegatorsStorageName>,
299
	pallet_parachain_staking::migrations::MigrateDelegationScheduledRequestsToDoubleMap<Runtime>,
300
);