From b122292da79c7ff0d89fa7939c91fe0c93c84fca Mon Sep 17 00:00:00 2001 From: Malte Kliemann Date: Wed, 7 Feb 2024 14:59:42 +0100 Subject: [PATCH] Migrate to new pool struct --- runtime/common/src/lib.rs | 3 +- zrml/neo-swaps/src/consts.rs | 2 +- zrml/neo-swaps/src/lib.rs | 2 +- zrml/neo-swaps/src/migration.rs | 233 ++++++++++++++++++++++++ zrml/neo-swaps/src/mock.rs | 2 +- zrml/neo-swaps/src/tests/deploy_pool.rs | 2 +- zrml/neo-swaps/src/types/max_assets.rs | 3 +- zrml/neo-swaps/src/types/pool.rs | 17 +- 8 files changed, 252 insertions(+), 12 deletions(-) diff --git a/runtime/common/src/lib.rs b/runtime/common/src/lib.rs index 12e5a57cb..37abc22b4 100644 --- a/runtime/common/src/lib.rs +++ b/runtime/common/src/lib.rs @@ -55,12 +55,13 @@ macro_rules! decl_common_types { use orml_traits::MultiCurrency; use sp_runtime::{generic, DispatchError, DispatchResult, SaturatedConversion}; use zeitgeist_primitives::traits::{DeployPoolApi, DistributeFees, MarketCommonsPalletApi}; + use zrml_neo_swaps::migration::MigratePoolReservesToBoundedBTreeMap; pub type Block = generic::Block; type Address = sp_runtime::MultiAddress; - type Migrations = (); + type Migrations = (MigratePoolReservesToBoundedBTreeMap); pub type Executive = frame_executive::Executive< Runtime, diff --git a/zrml/neo-swaps/src/consts.rs b/zrml/neo-swaps/src/consts.rs index 19d0af5f7..270c9fc1a 100644 --- a/zrml/neo-swaps/src/consts.rs +++ b/zrml/neo-swaps/src/consts.rs @@ -22,7 +22,7 @@ pub(crate) const EXP_NUMERICAL_LIMIT: u128 = 10; /// Numerical lower limit for ln arguments (fixed point number). pub(crate) const LN_NUMERICAL_LIMIT: u128 = BASE / 10; /// The maximum number of assets allowed in a pool. -pub(crate) const MAX_ASSETS: u16 = 128; +pub(crate) const MAX_ASSETS: u32 = 128; pub(crate) const _1: u128 = BASE; pub(crate) const _2: u128 = 2 * _1; diff --git a/zrml/neo-swaps/src/lib.rs b/zrml/neo-swaps/src/lib.rs index 974dc031f..698c61e26 100644 --- a/zrml/neo-swaps/src/lib.rs +++ b/zrml/neo-swaps/src/lib.rs @@ -74,7 +74,7 @@ mod pallet { }; use zrml_market_commons::MarketCommonsPalletApi; - pub(crate) const STORAGE_VERSION: StorageVersion = StorageVersion::new(1); + pub(crate) const STORAGE_VERSION: StorageVersion = StorageVersion::new(2); // 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 fb9ed7fb7..2aa4996f9 100644 --- a/zrml/neo-swaps/src/migration.rs +++ b/zrml/neo-swaps/src/migration.rs @@ -15,6 +15,239 @@ // You should have received a copy of the GNU General Public License // along with Zeitgeist. If not, see . +use crate::{ + traits::LiquiditySharesManager, types::Pool, AssetOf, BalanceOf, Config, LiquidityTreeOf, + Pallet, Pools, +}; +use alloc::collections::BTreeMap; +use core::marker::PhantomData; +use frame_support::{ + dispatch::Weight, + log, + traits::{Get, OnRuntimeUpgrade, StorageVersion}, + RuntimeDebug, +}; +use parity_scale_codec::{Decode, Encode}; +use scale_info::TypeInfo; +use sp_runtime::Saturating; + +cfg_if::cfg_if! { + if #[cfg(feature = "try-runtime")] { + use crate::{MarketIdOf}; + use alloc::{format, vec::Vec}; + use frame_support::{migration::storage_key_iter, pallet_prelude::Twox64Concat}; + } +} + +cfg_if::cfg_if! { + if #[cfg(any(feature = "try-runtime", test))] { + const NEO_SWAPS: &[u8] = b"NeoSwaps"; + const POOLS: &[u8] = b"Pools"; + } +} + +const NEO_SWAPS_REQUIRED_STORAGE_VERSION: u16 = 1; +const NEO_SWAPS_NEXT_STORAGE_VERSION: u16 = NEO_SWAPS_REQUIRED_STORAGE_VERSION + 1; + +#[derive(Clone, Decode, Encode, Eq, PartialEq, RuntimeDebug, TypeInfo)] +#[scale_info(skip_type_params(T))] +pub struct OldPool +where + T: Config, + LSM: LiquiditySharesManager, +{ + pub account_id: T::AccountId, + pub reserves: BTreeMap, BalanceOf>, + pub collateral: AssetOf, + pub liquidity_parameter: BalanceOf, + pub liquidity_shares_manager: LSM, + pub swap_fee: BalanceOf, +} + +type OldPoolOf = OldPool>; + +pub struct MigratePoolReservesToBoundedBTreeMap(PhantomData); + +impl OnRuntimeUpgrade for MigratePoolReservesToBoundedBTreeMap +where + T: Config, +{ + 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!( + "MigratePoolReservesToBoundedBTreeMap: neo-swaps version is {:?}, but {:?} is \ + required", + neo_swaps_version, + NEO_SWAPS_REQUIRED_STORAGE_VERSION, + ); + return total_weight; + } + log::info!("MigratePoolReservesToBoundedBTreeMap: Starting..."); + let mut translated = 0u64; + Pools::::translate::, _>(|_, pool| { + // Can't fail unless `MaxAssets` is misconfigured. If it fails after all, we delete the + // pool. This may seem drastic, but is actually cleaner than trying some half-baked + // recovery and allows us to do a manual recovery of funds. + let reserves = pool.reserves.try_into().ok()?; + translated.saturating_inc(); + Some(Pool { + account_id: pool.account_id, + reserves, + collateral: pool.collateral, + liquidity_parameter: pool.liquidity_parameter, + liquidity_shares_manager: pool.liquidity_shares_manager, + swap_fee: pool.swap_fee, + }) + }); + log::info!("MigratePoolReservesToBoundedBTreeMap: Upgraded {} pools.", translated); + total_weight = + total_weight.saturating_add(T::DbWeight::get().reads_writes(translated, translated)); + StorageVersion::new(NEO_SWAPS_NEXT_STORAGE_VERSION).put::>(); + total_weight = total_weight.saturating_add(T::DbWeight::get().writes(1)); + log::info!("MigratePoolReservesToBoundedBTreeMap: Done!"); + total_weight + } + + #[cfg(feature = "try-runtime")] + fn pre_upgrade() -> Result, &'static str> { + 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<(), &'static str> { + 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); + for (market_id, new_pool) in Pools::::iter() { + let old_pool = + old_pools.get(&market_id).expect(&format!("Pool {:?} not found", market_id)[..]); + assert_eq!(new_pool.account_id, old_pool.account_id); + assert_eq!(new_pool.reserves.into_inner(), 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); + } + log::info!( + "MigratePoolReservesToBoundedBTreeMap: Post-upgrade pool count is {}!", + new_pool_count + ); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + liquidity_tree::types::LiquidityTree, + mock::{ExtBuilder, Runtime}, + MarketIdOf, PoolOf, Pools, + }; + use alloc::collections::BTreeMap; + use frame_support::{ + dispatch::fmt::Debug, migration::put_storage_value, storage_root, StateVersion, + StorageHasher, Twox64Concat, + }; + use parity_scale_codec::Encode; + use zeitgeist_primitives::types::Asset; + + #[test] + fn on_runtime_upgrade_increments_the_storage_version() { + ExtBuilder::default().build().execute_with(|| { + set_up_version(); + MigratePoolReservesToBoundedBTreeMap::::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); + MigratePoolReservesToBoundedBTreeMap::::on_runtime_upgrade(); + assert_eq!(tmp, storage_root(StateVersion::V1)); + }); + } + + #[test] + fn on_runtime_upgrade_correctly_updates_markets() { + ExtBuilder::default().build().execute_with(|| { + set_up_version(); + let (old_pools, new_pools) = construct_old_new_tuple(); + populate_test_data::, OldPoolOf>( + NEO_SWAPS, POOLS, old_pools, + ); + MigratePoolReservesToBoundedBTreeMap::::on_runtime_upgrade(); + let actual = Pools::get(0u128).unwrap(); + assert_eq!(actual, new_pools[0]); + }); + } + + fn set_up_version() { + StorageVersion::new(NEO_SWAPS_REQUIRED_STORAGE_VERSION).put::>(); + } + + fn construct_old_new_tuple() -> (Vec>, Vec>) { + let account_id = 1; + let mut old_reserves = BTreeMap::new(); + old_reserves.insert(Asset::CategoricalOutcome(2, 3), 4); + let new_reserves = old_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: old_reserves, + collateral, + liquidity_parameter, + liquidity_shares_manager: liquidity_shares_manager.clone(), + swap_fee, + }; + let new_pool = Pool { + account_id, + reserves: new_reserves, + collateral, + liquidity_parameter, + liquidity_shares_manager, + swap_fee, + }; + (vec![old_pool], vec![new_pool]) + } + + #[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; diff --git a/zrml/neo-swaps/src/mock.rs b/zrml/neo-swaps/src/mock.rs index a6d519fd8..8466e155e 100644 --- a/zrml/neo-swaps/src/mock.rs +++ b/zrml/neo-swaps/src/mock.rs @@ -98,7 +98,7 @@ parameter_types! { pub const OracleBond: Balance = 0; pub const ValidityBond: Balance = 0; pub const DisputeBond: Balance = 0; - pub const MaxCategories: u16 = MAX_ASSETS + 1; + pub const MaxCategories: u16 = MAX_ASSETS as u16 + 1; } pub struct DeployPoolNoop; diff --git a/zrml/neo-swaps/src/tests/deploy_pool.rs b/zrml/neo-swaps/src/tests/deploy_pool.rs index 689dd76c0..b6d2ae4a1 100644 --- a/zrml/neo-swaps/src/tests/deploy_pool.rs +++ b/zrml/neo-swaps/src/tests/deploy_pool.rs @@ -262,7 +262,7 @@ fn deploy_pool_fails_on_asset_count_above_max() { let market_id = create_market( ALICE, BASE_ASSET, - MarketType::Categorical(category_count), + MarketType::Categorical(category_count as u16), ScoringRule::Lmsr, ); let liquidity = _10; diff --git a/zrml/neo-swaps/src/types/max_assets.rs b/zrml/neo-swaps/src/types/max_assets.rs index 1a5e96b7c..b939bc3d4 100644 --- a/zrml/neo-swaps/src/types/max_assets.rs +++ b/zrml/neo-swaps/src/types/max_assets.rs @@ -16,11 +16,12 @@ // along with Zeitgeist. If not, see . use sp_runtime::traits::Get; +use crate::consts::MAX_ASSETS; pub(crate) struct MaxAssets; impl Get for MaxAssets { fn get() -> u32 { - 128 + MAX_ASSETS } } diff --git a/zrml/neo-swaps/src/types/pool.rs b/zrml/neo-swaps/src/types/pool.rs index bf355919a..683088840 100644 --- a/zrml/neo-swaps/src/types/pool.rs +++ b/zrml/neo-swaps/src/types/pool.rs @@ -22,21 +22,26 @@ use crate::{ traits::{LiquiditySharesManager, PoolOperations}, Error, }; -use alloc::vec::Vec; -use frame_support::{storage::bounded_btree_map::BoundedBTreeMap, CloneNoBound}; +use alloc::{fmt::Debug, vec::Vec}; +use frame_support::{ + storage::bounded_btree_map::BoundedBTreeMap, CloneNoBound, PartialEqNoBound, + RuntimeDebugNoBound, +}; use parity_scale_codec::{Decode, Encode, MaxEncodedLen}; use scale_info::TypeInfo; use sp_runtime::{ traits::{CheckedAdd, CheckedSub, Get}, - DispatchError, DispatchResult, RuntimeDebug, SaturatedConversion, Saturating, + DispatchError, DispatchResult, SaturatedConversion, Saturating, }; -#[derive(CloneNoBound, Decode, Encode, Eq, MaxEncodedLen, PartialEq, RuntimeDebug, TypeInfo)] +#[derive( + CloneNoBound, Decode, Encode, Eq, MaxEncodedLen, PartialEqNoBound, RuntimeDebugNoBound, TypeInfo, +)] #[scale_info(skip_type_params(S, T))] pub struct Pool where T: Config, - LSM: Clone + LiquiditySharesManager, + LSM: Clone + Debug + LiquiditySharesManager + PartialEq, S: Get, { pub account_id: T::AccountId, @@ -51,7 +56,7 @@ impl PoolOperations for Pool where T: Config, BalanceOf: SaturatedConversion, - LSM: Clone + LiquiditySharesManager + TypeInfo, + LSM: Clone + Debug + LiquiditySharesManager + TypeInfo + PartialEq, S: Get, { fn assets(&self) -> Vec> {