From 33a736eecfb282c54f5652c2573e848b3e7c7727 Mon Sep 17 00:00:00 2001 From: Chralt Date: Thu, 30 Jan 2025 14:29:21 +0100 Subject: [PATCH] Add storage migrations for Combinatorial Tokens Upgrade (#1401) * add todos for storage migrations * remove todo for asset storage migration * add storage migration and try-runtime tests * correct migration and tests * update todos * update migration tests * add migration to runtime, fix clippy * bump storage version * remove corrupted pools * correct clippy * correct CI --- primitives/src/asset.rs | 1 - runtime/common/src/lib.rs | 4 +- zrml/neo-swaps/src/lib.rs | 2 +- zrml/neo-swaps/src/migration.rs | 398 ++++++++++++++++++++++++++++++++ 4 files changed, 402 insertions(+), 3 deletions(-) diff --git a/primitives/src/asset.rs b/primitives/src/asset.rs index dbd22367d..d3d87223a 100644 --- a/primitives/src/asset.rs +++ b/primitives/src/asset.rs @@ -55,7 +55,6 @@ pub enum Asset { ForeignAsset(u32), ParimutuelShare(MarketId, CategoryIndex), } -// TODO Needs storage migration #[cfg(feature = "runtime-benchmarks")] impl ZeitgeistAssetEnumerator for Asset { diff --git a/runtime/common/src/lib.rs b/runtime/common/src/lib.rs index 2b26b806a..8d3426fce 100644 --- a/runtime/common/src/lib.rs +++ b/runtime/common/src/lib.rs @@ -96,11 +96,13 @@ macro_rules! decl_common_types { #[cfg(feature = "runtime-benchmarks")] use zrml_prediction_markets::types::PredictionMarketsCombinatorialTokensBenchmarkHelper; + use zrml_neo_swaps::migration::MigratePoolStorageItems; + pub type Block = generic::Block; type Address = sp_runtime::MultiAddress; - type Migrations = (); + type Migrations = (MigratePoolStorageItems); pub type Executive = frame_executive::Executive< Runtime, diff --git a/zrml/neo-swaps/src/lib.rs b/zrml/neo-swaps/src/lib.rs index 2fc47f05a..88b45bc3a 100644 --- a/zrml/neo-swaps/src/lib.rs +++ b/zrml/neo-swaps/src/lib.rs @@ -93,7 +93,7 @@ mod pallet { }; use zrml_market_commons::MarketCommonsPalletApi; - pub(crate) const STORAGE_VERSION: StorageVersion = StorageVersion::new(2); + pub(crate) const STORAGE_VERSION: StorageVersion = StorageVersion::new(3); // These should not be config parameters to avoid misconfigurations. pub(crate) const EXIT_FEE: u128 = CENT / 10; diff --git a/zrml/neo-swaps/src/migration.rs b/zrml/neo-swaps/src/migration.rs index 2e9ba478f..3ef4559d5 100644 --- a/zrml/neo-swaps/src/migration.rs +++ b/zrml/neo-swaps/src/migration.rs @@ -14,3 +14,401 @@ // // You should have received a copy of the GNU General Public License // along with Zeitgeist. If not, see . + +use crate::{ + traits::LiquiditySharesManager, + types::{MaxAssets, Pool, PoolType}, + AssetOf, BalanceOf, Config, LiquidityTreeOf, MarketIdOf, MarketIdToPoolId, Pallet, PoolCount, + Pools, +}; +use alloc::{fmt::Debug, vec, vec::Vec}; +use core::marker::PhantomData; +use frame_support::{ + migration::storage_key_iter, + pallet_prelude::Twox64Concat, + storage::bounded_btree_map::BoundedBTreeMap, + traits::{Get, OnRuntimeUpgrade, StorageVersion}, + weights::Weight, + CloneNoBound, PartialEqNoBound, RuntimeDebugNoBound, +}; +use log; +use parity_scale_codec::{Decode, Encode, MaxEncodedLen}; +use scale_info::TypeInfo; +use sp_runtime::{SaturatedConversion, Saturating}; +use zeitgeist_primitives::math::checked_ops_res::CheckedAddRes; +use zrml_market_commons::MarketCommonsPalletApi; + +cfg_if::cfg_if! { + if #[cfg(feature = "try-runtime")] { + use alloc::{format, collections::BTreeMap}; + use sp_runtime::DispatchError; + } +} + +const NEO_SWAPS: &[u8] = b"NeoSwaps"; +const POOLS: &[u8] = b"Pools"; + +const NEO_SWAPS_REQUIRED_STORAGE_VERSION: u16 = 2; +const NEO_SWAPS_NEXT_STORAGE_VERSION: u16 = NEO_SWAPS_REQUIRED_STORAGE_VERSION + 1; + +#[derive( + CloneNoBound, Decode, Encode, Eq, MaxEncodedLen, PartialEqNoBound, RuntimeDebugNoBound, TypeInfo, +)] +#[scale_info(skip_type_params(S, T))] +pub struct OldPool +where + T: Config, + LSM: Clone + Debug + LiquiditySharesManager + PartialEq, + S: Get, +{ + pub account_id: T::AccountId, + pub reserves: BoundedBTreeMap, BalanceOf, S>, + pub collateral: AssetOf, + pub liquidity_parameter: BalanceOf, + pub liquidity_shares_manager: LSM, + pub swap_fee: BalanceOf, +} + +type OldPoolOf = OldPool, MaxAssets>; + +// https://substrate.stackexchange.com/questions/10472/pallet-storage-migration-fails-try-runtime-idempotent-test +// idempotent test fails, because of the manual storage version increment +// VersionedMigration is still an experimental feature for the currently used polkadot version +// that's why the idempotent test is ignored for this migration +pub struct MigratePoolStorageItems(PhantomData, RemovableMarketIds); + +impl OnRuntimeUpgrade for MigratePoolStorageItems +where + T: Config, + RemovableMarketIds: Get>, +{ + fn on_runtime_upgrade() -> Weight { + let mut total_weight = T::DbWeight::get().reads(1); + let neo_swaps_version = StorageVersion::get::>(); + if neo_swaps_version != NEO_SWAPS_REQUIRED_STORAGE_VERSION { + log::info!( + "MigratePoolStorageItems: neo-swaps version is {:?}, but {:?} is required", + neo_swaps_version, + NEO_SWAPS_REQUIRED_STORAGE_VERSION, + ); + return total_weight; + } + log::info!("MigratePoolStorageItems: Starting..."); + // NeoSwaps: 7de9893ad4de67f3510fd09678a13412 + // Pools: 4c72016d74b63ae83d79b02efdb5528e + // failed to decode pool with market id 880: 0x7de9893ad4de67f3510fd09678a134124c72016d74b63ae83d79b02efdb5528e00251b42e33e726f70030000000000000000000000000000 + // failed to decode pool with market id 878: 0x7de9893ad4de67f3510fd09678a134124c72016d74b63ae83d79b02efdb5528e0f7a0cea0db6ee406e030000000000000000000000000000 + // failed to decode pool with market id 882: 0x7de9893ad4de67f3510fd09678a134124c72016d74b63ae83d79b02efdb5528eb49736cf4bc6723372030000000000000000000000000000 + // failed to decode pool with market id 879: 0x7de9893ad4de67f3510fd09678a134124c72016d74b63ae83d79b02efdb5528ed857f1051e4281a76f030000000000000000000000000000 + // failed to decode pool with market id 877: 0x7de9893ad4de67f3510fd09678a134124c72016d74b63ae83d79b02efdb5528ee0edd4b43beb361f6d030000000000000000000000000000 + // The decode failure happens, because the old pool used a CampaignAsset as asset, which is not supported anymore, since the asset system overhaul has been reverted. + + let mut max_pool_id: T::PoolId = Default::default(); + for (market_id, _) in + storage_key_iter::, OldPoolOf, Twox64Concat>(NEO_SWAPS, POOLS) + { + total_weight = total_weight.saturating_add(T::DbWeight::get().reads(2)); + if T::MarketCommons::market(&market_id).is_err() { + log::error!("MigratePoolStorageItems: Market {:?} not found", market_id); + return total_weight; + }; + let pool_id = market_id; + max_pool_id = max_pool_id.max(pool_id); + } + let next_pool_count_id = if let Ok(id) = max_pool_id.checked_add_res(&1u8.into()) { + id + } else { + log::error!("MigratePoolStorageItems: Pool id overflow"); + return total_weight; + }; + let mut translated = 0u64; + Pools::::translate::, _>(|market_id, pool| { + translated.saturating_inc(); + let pool_id = market_id; + MarketIdToPoolId::::insert(pool_id, market_id); + let assets = if let Ok(market) = T::MarketCommons::market(&market_id) { + market.outcome_assets().try_into().ok()? + } else { + log::error!( + "MigratePoolStorageItems: Market {:?} not found. This should not happen, \ + because it is checked above.", + market_id + ); + pool.reserves.keys().cloned().collect::>().try_into().ok()? + }; + Some(Pool { + account_id: pool.account_id, + assets, + reserves: pool.reserves, + collateral: pool.collateral, + liquidity_parameter: pool.liquidity_parameter, + liquidity_shares_manager: pool.liquidity_shares_manager, + swap_fee: pool.swap_fee, + pool_type: PoolType::Standard(market_id), + }) + }); + PoolCount::::set(next_pool_count_id); + // Write for the PoolCount storage item + total_weight = total_weight.saturating_add(T::DbWeight::get().writes(1)); + log::info!("MigratePoolStorageItems: Upgraded {} pools.", translated); + // Reads and writes for the Pools storage item + total_weight = + total_weight.saturating_add(T::DbWeight::get().reads_writes(translated, translated)); + // Read for the market and write for the MarketIdToPoolId storage item + total_weight = + total_weight.saturating_add(T::DbWeight::get().reads_writes(translated, translated)); + + // remove pools that contain a corrupted campaign asset from the reverted asset system overhaul + let mut corrupted_pools = vec![]; + for &market_id in RemovableMarketIds::get().iter() { + let market_id = market_id.saturated_into::>(); + total_weight = total_weight.saturating_add(T::DbWeight::get().reads(2)); + let is_corrupted = + || Pools::::contains_key(market_id) && Pools::::get(market_id).is_none(); + if is_corrupted() { + total_weight = total_weight.saturating_add(T::DbWeight::get().writes(1)); + Pools::::remove(market_id); + corrupted_pools.push(market_id); + } else { + log::warn!( + "RemoveMarkets: Pool with market id {:?} was expected to be corrupted, but \ + isn't.", + market_id + ); + } + } + log::info!("RemovePools: Removed pools with market ids: {:?}.", corrupted_pools); + StorageVersion::new(NEO_SWAPS_NEXT_STORAGE_VERSION).put::>(); + total_weight = total_weight.saturating_add(T::DbWeight::get().writes(1)); + log::info!("MigratePoolStorageItems: Done!"); + total_weight + } + + #[cfg(feature = "try-runtime")] + fn pre_upgrade() -> Result, DispatchError> { + let old_pools = + storage_key_iter::, OldPoolOf, Twox64Concat>(NEO_SWAPS, POOLS) + .collect::>(); + Ok(old_pools.encode()) + } + + #[cfg(feature = "try-runtime")] + fn post_upgrade(previous_state: Vec) -> Result<(), DispatchError> { + let old_pools: BTreeMap, OldPoolOf> = + Decode::decode(&mut &previous_state[..]) + .map_err(|_| "Failed to decode state: Invalid state")?; + let new_pool_count = Pools::::iter().count(); + assert_eq!(old_pools.len(), new_pool_count); + let mut max_pool_id: T::PoolId = Default::default(); + for (market_id, new_pool) in Pools::::iter() { + let old_pool = + old_pools.get(&market_id).expect(&format!("Pool {:?} not found", market_id)[..]); + max_pool_id = max_pool_id.max(market_id); + assert_eq!(new_pool.account_id, old_pool.account_id); + let market = T::MarketCommons::market(&market_id)?; + let outcome_assets = market.outcome_assets(); + for asset in &outcome_assets { + assert!(new_pool.assets.contains(asset)); + } + assert_eq!(new_pool.assets.len(), outcome_assets.len()); + assert_eq!(new_pool.reserves, old_pool.reserves); + assert_eq!(new_pool.collateral, old_pool.collateral); + assert_eq!(new_pool.liquidity_parameter, old_pool.liquidity_parameter); + assert_eq!(new_pool.liquidity_shares_manager, old_pool.liquidity_shares_manager); + assert_eq!(new_pool.swap_fee, old_pool.swap_fee); + assert_eq!(new_pool.pool_type, PoolType::Standard(market_id)); + + assert_eq!( + MarketIdToPoolId::::get(market_id).expect("MarketIdToPoolId mapping not found"), + market_id + ); + } + let next_pool_count_id = PoolCount::::get(); + assert_eq!(next_pool_count_id, max_pool_id.checked_add_res(&1u8.into())?); + log::info!( + "MigratePoolStorageItems: Post-upgrade next pool count id is {:?}!", + next_pool_count_id + ); + for &market_id in RemovableMarketIds::get().iter() { + let market_id = market_id.saturated_into::>(); + assert!(!Pools::::contains_key(market_id)); + assert!(Pools::::try_get(market_id).is_err()); + } + log::info!("MigratePoolStorageItems: Post-upgrade pool count is {}!", new_pool_count); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + liquidity_tree::types::LiquidityTree, + mock::{ExtBuilder, MarketCommons, Runtime, ALICE, BOB}, + MarketIdOf, PoolOf, Pools, + }; + use alloc::collections::BTreeMap; + use core::fmt::Debug; + use frame_support::{migration::put_storage_value, StorageHasher, Twox64Concat}; + use parity_scale_codec::Encode; + use sp_io::storage::root as storage_root; + use sp_runtime::{Perbill, StateVersion}; + use zeitgeist_primitives::types::{ + Asset, Market, MarketCreation, MarketPeriod, MarketStatus, MarketType, ScoringRule, + }; + + struct RemovableMarketIds; + impl Get> for RemovableMarketIds { + fn get() -> Vec { + vec![] + } + } + + #[test] + fn on_runtime_upgrade_increments_the_storage_version() { + ExtBuilder::default().build().execute_with(|| { + set_up_version(); + MigratePoolStorageItems::::on_runtime_upgrade(); + assert_eq!(StorageVersion::get::>(), NEO_SWAPS_NEXT_STORAGE_VERSION); + }); + } + + #[test] + fn on_runtime_upgrade_is_noop_if_versions_are_not_correct() { + ExtBuilder::default().build().execute_with(|| { + StorageVersion::new(NEO_SWAPS_NEXT_STORAGE_VERSION).put::>(); + let (_, new_pools) = construct_old_new_tuple(); + populate_test_data::, PoolOf>( + NEO_SWAPS, POOLS, new_pools, + ); + let tmp = storage_root(StateVersion::V1); + MigratePoolStorageItems::::on_runtime_upgrade(); + assert_eq!(tmp, storage_root(StateVersion::V1)); + }); + } + + #[test] + fn on_runtime_upgrade_correctly_updates_pool_storages() { + ExtBuilder::default().build().execute_with(|| { + set_up_version(); + create_markets(3); + let (old_pools, new_pools) = construct_old_new_tuple(); + populate_test_data::, OldPoolOf>( + NEO_SWAPS, POOLS, old_pools, + ); + MigratePoolStorageItems::::on_runtime_upgrade(); + let actual = Pools::get(0u128).unwrap(); + assert_eq!(actual, new_pools[0]); + let next_pool_count_id = PoolCount::::get(); + assert_eq!(next_pool_count_id, 3u128); + assert_eq!(MarketIdToPoolId::::get(0u128).unwrap(), 0u128); + assert_eq!(MarketIdToPoolId::::get(1u128).unwrap(), 1u128); + assert_eq!(MarketIdToPoolId::::get(2u128).unwrap(), 2u128); + assert!(MarketIdToPoolId::::get(3u128).is_none()); + assert!(MarketIdToPoolId::::iter_keys().count() == 3); + }); + } + + fn set_up_version() { + StorageVersion::new(NEO_SWAPS_REQUIRED_STORAGE_VERSION).put::>(); + } + + fn create_markets(count: u8) { + for _ in 0..count { + let base_asset = Asset::Ztg; + let market = Market { + market_id: 0u8.into(), + base_asset, + creation: MarketCreation::Permissionless, + creator_fee: Perbill::zero(), + creator: ALICE, + oracle: BOB, + metadata: vec![0, 50], + market_type: MarketType::Categorical(3), + period: MarketPeriod::Block(0u32.into()..1u32.into()), + deadlines: Default::default(), + scoring_rule: ScoringRule::AmmCdaHybrid, + status: MarketStatus::Active, + report: None, + resolved_outcome: None, + dispute_mechanism: None, + bonds: Default::default(), + early_close: None, + }; + MarketCommons::push_market(market).unwrap(); + } + } + + fn construct_old_new_tuple() -> (Vec>, Vec>) { + let account_id = 1; + let mut reserves = BTreeMap::new(); + let asset_0 = Asset::CategoricalOutcome(0, 0); + let asset_1 = Asset::CategoricalOutcome(0, 1); + let asset_2 = Asset::CategoricalOutcome(0, 2); + reserves.insert(asset_0, 4); + reserves.insert(asset_1, 5); + reserves.insert(asset_2, 6); + let reserves: BoundedBTreeMap, BalanceOf, MaxAssets> = + reserves.clone().try_into().unwrap(); + let collateral = Asset::Ztg; + let liquidity_parameter = 5; + let swap_fee = 6; + let total_shares = 7; + let fees = 8; + + let mut liquidity_shares_manager = LiquidityTree::new(account_id, total_shares).unwrap(); + liquidity_shares_manager.nodes.get_mut(0).unwrap().fees = fees; + + let old_pool = OldPoolOf { + account_id, + reserves: reserves.clone(), + collateral, + liquidity_parameter, + liquidity_shares_manager: liquidity_shares_manager.clone(), + swap_fee, + }; + let new_pool = Pool { + account_id, + assets: vec![asset_0, asset_1, asset_2].try_into().unwrap(), + reserves, + collateral, + liquidity_parameter, + liquidity_shares_manager, + swap_fee, + pool_type: PoolType::Standard(0), + }; + ( + vec![old_pool.clone(), old_pool.clone(), old_pool.clone()], + vec![new_pool.clone(), new_pool.clone(), new_pool.clone()], + ) + } + + #[allow(unused)] + fn populate_test_data(pallet: &[u8], prefix: &[u8], data: Vec) + where + H: StorageHasher, + K: TryFrom + Encode, + V: Encode + Clone, + >::Error: Debug, + { + for (key, value) in data.iter().enumerate() { + let storage_hash = utility::key_to_hash::(K::try_from(key).unwrap()); + put_storage_value::(pallet, prefix, &storage_hash, (*value).clone()); + } + } +} + +mod utility { + use alloc::vec::Vec; + use frame_support::StorageHasher; + use parity_scale_codec::Encode; + + #[allow(unused)] + pub fn key_to_hash(key: K) -> Vec + where + H: StorageHasher, + K: Encode, + { + key.using_encoded(H::hash).as_ref().to_vec() + } +}