From 9a07e9d564636198afef9260de390ebcd40ac1fc Mon Sep 17 00:00:00 2001 From: Liam Aharon Date: Fri, 7 Apr 2023 00:27:25 +1000 Subject: [PATCH] `RemovePallet` migration utility struct (#13835) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * remove pallet utility struct * Update frame/support/src/remove_pallet.rs Co-authored-by: Bastian Köcher * Update frame/support/src/remove_pallet.rs Co-authored-by: Bastian Köcher * make removepallet a module of migrations * fix rust doc lint --------- Co-authored-by: Bastian Köcher --- frame/support/src/migrations.rs | 116 ++++++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/frame/support/src/migrations.rs b/frame/support/src/migrations.rs index 645316042b38c..f54b73aad8889 100644 --- a/frame/support/src/migrations.rs +++ b/frame/support/src/migrations.rs @@ -15,11 +15,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(feature = "try-runtime")] +use crate::storage::unhashed::contains_prefixed_key; use crate::{ traits::{GetStorageVersion, PalletInfoAccess}, weights::{RuntimeDbWeight, Weight}, }; use impl_trait_for_tuples::impl_for_tuples; +use sp_core::Get; +use sp_io::{hashing::twox_128, storage::clear_prefix, KillStorageResult}; +use sp_std::marker::PhantomData; +#[cfg(feature = "try-runtime")] +use sp_std::vec::Vec; /// Trait used by [`migrate_from_pallet_version_to_storage_version`] to do the actual migration. pub trait PalletVersionToStorageVersionHelper { @@ -67,3 +74,112 @@ pub fn migrate_from_pallet_version_to_storage_version< ) -> Weight { Pallets::migrate(db_weight) } + +/// `RemovePallet` is a utility struct used to remove all storage items associated with a specific +/// pallet. +/// +/// This struct is generic over two parameters: +/// - `P` is a type that implements the `Get` trait for a static string, representing the pallet's +/// name. +/// - `DbWeight` is a type that implements the `Get` trait for `RuntimeDbWeight`, providing the +/// weight for database operations. +/// +/// On runtime upgrade, the `on_runtime_upgrade` function will clear all storage items associated +/// with the specified pallet, logging the number of keys removed. If the `try-runtime` feature is +/// enabled, the `pre_upgrade` and `post_upgrade` functions can be used to verify the storage +/// removal before and after the upgrade. +/// +/// # Examples: +/// ```ignore +/// construct_runtime! { +/// pub enum Runtime where +/// Block = Block, +/// NodeBlock = primitives::Block, +/// UncheckedExtrinsic = UncheckedExtrinsic +/// { +/// System: frame_system::{Pallet, Call, Storage, Config, Event} = 0, +/// +/// SomePalletToRemove: pallet_something::{Pallet, Call, Storage, Event} = 1, +/// AnotherPalletToRemove: pallet_something_else::{Pallet, Call, Storage, Event} = 2, +/// +/// YourOtherPallets... +/// } +/// }; +/// +/// parameter_types! { +/// pub const SomePalletToRemoveStr: &'static str = "SomePalletToRemove"; +/// pub const AnotherPalletToRemoveStr: &'static str = "AnotherPalletToRemove"; +/// } +/// +/// pub type Migrations = ( +/// RemovePallet, +/// RemovePallet, +/// AnyOtherMigrations... +/// ); +/// +/// pub type Executive = frame_executive::Executive< +/// Runtime, +/// Block, +/// frame_system::ChainContext, +/// Runtime, +/// Migrations +/// >; +/// ``` +/// +/// WARNING: `RemovePallet` has no guard rails preventing it from bricking the chain if the +/// operation of removing storage for the given pallet would exceed the block weight limit. +/// +/// If your pallet has too many keys to be removed in a single block, it is advised to wait for +/// a multi-block scheduler currently under development which will allow for removal of storage +/// items (and performing other heavy migrations) over multiple blocks +/// (see ). +pub struct RemovePallet, DbWeight: Get>( + PhantomData<(P, DbWeight)>, +); +impl, DbWeight: Get> frame_support::traits::OnRuntimeUpgrade + for RemovePallet +{ + fn on_runtime_upgrade() -> frame_support::weights::Weight { + let hashed_prefix = twox_128(P::get().as_bytes()); + let keys_removed = match clear_prefix(&hashed_prefix, None) { + KillStorageResult::AllRemoved(value) => value, + KillStorageResult::SomeRemaining(value) => { + log::error!( + "`clear_prefix` failed to remove all keys for {}. THIS SHOULD NEVER HAPPEN! 🚨", + P::get() + ); + value + }, + } as u64; + + log::info!("Removed {} {} keys 🧹", keys_removed, P::get()); + + DbWeight::get().reads_writes(keys_removed + 1, keys_removed) + } + + #[cfg(feature = "try-runtime")] + fn pre_upgrade() -> Result, &'static str> { + let hashed_prefix = twox_128(P::get().as_bytes()); + match contains_prefixed_key(&hashed_prefix) { + true => log::info!("Found {} keys pre-removal 👀", P::get()), + false => log::warn!( + "Migration RemovePallet<{}> can be removed (no keys found pre-removal).", + P::get() + ), + }; + Ok(Vec::new()) + } + + #[cfg(feature = "try-runtime")] + fn post_upgrade(_state: Vec) -> Result<(), &'static str> { + let hashed_prefix = twox_128(P::get().as_bytes()); + match contains_prefixed_key(&hashed_prefix) { + true => { + log::error!("{} has keys remaining post-removal ❗", P::get()); + return Err("Keys remaining post-removal, this should never happen 🚨") + }, + false => log::info!("No {} keys found post-removal 🎉", P::get()), + }; + Ok(()) + } +}