Rework on runtime upgrade logic for pallet-balanced-currency-swap-bri…

…dges-initializer

crates/pallet-balanced-currency-swap-bridges-initializer/src/lib.rs

@@ -14,13 +14,12 @@ use frame_support::{
 };
 pub use pallet::*;
 use sp_std::cmp::Ordering;
-#[cfg(feature = "try-runtime")]
-use sp_std::vec::Vec;
-pub use weights::*;
 
-pub mod weights;
+mod weights;
+pub use weights::*;
 
 mod upgrade_init;
+pub use upgrade_init::UpgradeInit;
 
 #[cfg(test)]
 mod mock;
@@ -147,23 +146,6 @@ pub mod pallet {
         NotBalanced,
     }
 
-    #[pallet::hooks]
-    impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-        fn on_runtime_upgrade() -> Weight {
-            upgrade_init::on_runtime_upgrade::<T>()
-        }
-
-        #[cfg(feature = "try-runtime")]
-        fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
-            upgrade_init::pre_upgrade()
-        }
-
-        #[cfg(feature = "try-runtime")]
-        fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
-            upgrade_init::post_upgrade::<T>(state)
-        }
-    }
-
     #[pallet::call(weight(T::WeightInfo))]
     impl<T: Config> Pallet<T> {
         /// Verify if currencies are balanced.

crates/pallet-balanced-currency-swap-bridges-initializer/src/tests.rs

@@ -5,7 +5,7 @@ use frame_support::{
 
 use crate::{
     mock::{new_test_ext_with, v0, v1, v2, with_runtime_lock, *},
-    swappable_balance, LastForceRebalanceAskCounter, LastInitializerVersion,
+    swappable_balance, LastForceRebalanceAskCounter, LastInitializerVersion, UpgradeInit,
 };
 
 /// This test verifies that balanced bridges initialization works in case bridge pot accounts
@@ -411,7 +411,7 @@ fn runtime_upgrade() {
             assert_eq!(<LastForceRebalanceAskCounter<v1::Test>>::get(), 0);
 
             // Do runtime upgrade hook.
-            v1::AllPalletsWithoutSystem::on_runtime_upgrade();
+            UpgradeInit::<v1::Test>::on_runtime_upgrade();
 
             // Verify bridges initialization result.
             assert_eq!(
@@ -457,7 +457,7 @@ fn runtime_upgrade() {
                 v2::EvmBalances::total_balance(&v2::SwapBridgeEvmToNativePot::account_id());
 
             // Do runtime upgrade hook.
-            v2::AllPalletsWithoutSystem::on_runtime_upgrade();
+            UpgradeInit::<v2::Test>::on_runtime_upgrade();
 
             // Verify result.
             assert_eq!(

crates/pallet-balanced-currency-swap-bridges-initializer/src/upgrade_init.rs

@@ -1,6 +1,10 @@
 //! Initialization of the bridge pot accounts on runtime upgrade.
 
-use frame_support::{log, pallet_prelude::*};
+use frame_support::{
+    log::{error, info},
+    pallet_prelude::*,
+    traits::OnRuntimeUpgrade,
+};
 #[cfg(feature = "try-runtime")]
 use sp_std::vec::Vec;
 
@@ -9,54 +13,54 @@ use crate::{
     CURRENT_BRIDGES_INITIALIZER_VERSION,
 };
 
-/// Initialize the bridges pot accounts if required.
-pub fn on_runtime_upgrade<T: Config>() -> Weight {
-    let last_initializer_version = <LastInitializerVersion<T>>::get();
-    let last_force_rebalance_ask_counter = <LastForceRebalanceAskCounter<T>>::get();
-    let current_force_rebalance_ask_counter = T::ForceRebalanceAskCounter::get();
+/// Execute upgrade init.
+pub struct UpgradeInit<T>(sp_std::marker::PhantomData<T>);
+
+impl<T: Config> OnRuntimeUpgrade for UpgradeInit<T> {
+    fn on_runtime_upgrade() -> Weight {
+        let last_initializer_version = <LastInitializerVersion<T>>::get();
+        let last_force_rebalance_ask_counter = <LastForceRebalanceAskCounter<T>>::get();
+        let current_force_rebalance_ask_counter = T::ForceRebalanceAskCounter::get();
+
+        let mut weight = T::DbWeight::get().reads(2);
 
-    let mut weight = T::DbWeight::get().reads(2);
+        let is_version_mismatch = last_initializer_version != CURRENT_BRIDGES_INITIALIZER_VERSION;
+        let is_forced = last_force_rebalance_ask_counter < current_force_rebalance_ask_counter;
 
-    let is_version_mismatch = last_initializer_version != CURRENT_BRIDGES_INITIALIZER_VERSION;
-    let is_forced = last_force_rebalance_ask_counter < current_force_rebalance_ask_counter;
+        if is_version_mismatch || is_forced {
+            match Pallet::<T>::initialize() {
+                Ok(w) => weight.saturating_accrue(w),
+                Err(err) => error!("error during bridges initialization: {err:?}"),
+            }
 
-    if is_version_mismatch || is_forced {
-        match Pallet::<T>::initialize() {
-            Ok(w) => weight.saturating_accrue(w),
-            Err(err) => log::error!("error during bridges initialization: {err:?}"),
+            <LastInitializerVersion<T>>::put(CURRENT_BRIDGES_INITIALIZER_VERSION);
+            <LastForceRebalanceAskCounter<T>>::put(current_force_rebalance_ask_counter);
+            weight.saturating_accrue(T::DbWeight::get().writes(2));
+        } else {
+            info!("Nothing to do. This runtime upgrade probably should be removed");
         }
 
-        <LastInitializerVersion<T>>::put(CURRENT_BRIDGES_INITIALIZER_VERSION);
-        <LastForceRebalanceAskCounter<T>>::put(current_force_rebalance_ask_counter);
-        weight.saturating_accrue(T::DbWeight::get().writes(2));
+        weight
     }
 
-    weight
-}
-
-/// Check the state before the bridges initialization.
-///
-/// Panics if anything goes wrong.
-#[cfg(feature = "try-runtime")]
-pub fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
-    // Do nothing.
-    Ok(Vec::new())
-}
+    #[cfg(feature = "try-runtime")]
+    fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+        // Do nothing.
+        Ok(Vec::new())
+    }
 
-/// Check the state after the bridges initialization.
-///
-/// Panics if anything goes wrong.
-#[cfg(feature = "try-runtime")]
-pub fn post_upgrade<T: Config>(_state: Vec<u8>) -> Result<(), &'static str> {
-    use frame_support::{storage_root, StateVersion};
+    #[cfg(feature = "try-runtime")]
+    fn post_upgrade(_state: Vec<u8>) -> Result<(), &'static str> {
+        use frame_support::{storage_root, StateVersion};
 
-    let storage_root_before = storage_root(StateVersion::V1);
+        let storage_root_before = storage_root(StateVersion::V1);
 
-    if !Pallet::<T>::is_balanced()? {
-        return Err("currencies are not balanced");
-    }
+        if !Pallet::<T>::is_balanced()? {
+            return Err("currencies are not balanced");
+        }
 
-    assert_eq!(storage_root_before, storage_root(StateVersion::V1));
+        assert_eq!(storage_root_before, storage_root(StateVersion::V1));
 
-    Ok(())
+        Ok(())
+    }
 }