refactor: add PokeTokenWeights to Pool interface (#1232)

* updates * updates * updates * updates * updates * updates * updates * updates
osmosis-labs · Apr 12, 2022 · b6d342a · b6d342a
1 parent 204d73d
commit b6d342a
Show file tree

Hide file tree

Showing 22 changed files with 117 additions and 91 deletions.
diff --git a/app/apptesting/test_suite.go b/app/apptesting/test_suite.go
@@ -6,24 +6,21 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
 	"github.com/cosmos/cosmos-sdk/simapp"
-
-	abci "github.com/tendermint/tendermint/abci/types"
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	slashingtypes "github.com/cosmos/cosmos-sdk/x/slashing/types"
 	"github.com/cosmos/cosmos-sdk/x/staking/teststaking"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+	"github.com/stretchr/testify/suite"
+	abci "github.com/tendermint/tendermint/abci/types"
+	"github.com/tendermint/tendermint/crypto/ed25519"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 
 	"github.com/osmosis-labs/osmosis/v7/app"
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/pool-models/balancer"
 	gammtypes "github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 	lockupkeeper "github.com/osmosis-labs/osmosis/v7/x/lockup/keeper"
 	lockuptypes "github.com/osmosis-labs/osmosis/v7/x/lockup/types"
-
-	"github.com/stretchr/testify/suite"
-	"github.com/tendermint/tendermint/crypto/ed25519"
 )
 
 type KeeperTestHelper struct {
@@ -179,7 +176,7 @@ func (keeperTestHelper *KeeperTestHelper) SetupGammPoolsWithBondDenomMultiplier(
 		poolId, err := keeperTestHelper.App.GAMMKeeper.CreatePool(keeperTestHelper.Ctx, msg)
 		keeperTestHelper.Require().NoError(err)
 
-		pool, err := keeperTestHelper.App.GAMMKeeper.GetPool(keeperTestHelper.Ctx, poolId)
+		pool, err := keeperTestHelper.App.GAMMKeeper.GetPoolAndPoke(keeperTestHelper.Ctx, poolId)
 		keeperTestHelper.Require().NoError(err)
 
 		pools = append(pools, pool)

diff --git a/app/upgrades/v5/whitelist_feetokens.go b/app/upgrades/v5/whitelist_feetokens.go
@@ -58,7 +58,7 @@ func InitialWhitelistedFeetokens(ctx sdk.Context, gamm *gammkeeper.Keeper) []typ
 			panic(err)
 		}
 
-		pool, poolExistsErr := gamm.GetPool(ctx, poolId)
+		pool, poolExistsErr := gamm.GetPoolAndPoke(ctx, poolId)
 		_ = pool
 		if poolExistsErr != nil {
 			continue

diff --git a/app/wasm/queries.go b/app/wasm/queries.go
@@ -24,7 +24,7 @@ func NewQueryPlugin(
 }
 
 func (qp QueryPlugin) GetPoolState(ctx sdk.Context, poolID uint64) (*wasmbindings.PoolAssets, error) {
-	poolData, err := qp.gammKeeper.GetPool(ctx, poolID)
+	poolData, err := qp.gammKeeper.GetPoolAndPoke(ctx, poolID)
 	if err != nil {
 		return nil, sdkerrors.Wrap(err, "gamm get pool")
 	}
@@ -50,7 +50,7 @@ func (qp QueryPlugin) GetSpotPrice(ctx sdk.Context, spotPrice *wasmbindings.Spot
 		return nil, sdkerrors.Wrap(err, "gamm get spot price")
 	}
 	if withSwapFee {
-		poolData, err := qp.gammKeeper.GetPool(ctx, poolId)
+		poolData, err := qp.gammKeeper.GetPoolAndPoke(ctx, poolId)
 		if err != nil {
 			return nil, sdkerrors.Wrap(err, "gamm get pool")
 		}

diff --git a/x/gamm/genesis.go b/x/gamm/genesis.go
@@ -3,6 +3,7 @@ package gamm
 import (
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/keeper"
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 )
@@ -36,7 +37,7 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState,
 
 // ExportGenesis returns the capability module's exported genesis.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) *types.GenesisState {
-	pools, err := k.GetPools(ctx)
+	pools, err := k.GetPoolsAndPoke(ctx)
 	if err != nil {
 		panic(err)
 	}

diff --git a/x/gamm/genesis_test.go b/x/gamm/genesis_test.go
@@ -6,14 +6,15 @@ import (
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	osmoapp "github.com/osmosis-labs/osmosis/v7/app"
-	"github.com/osmosis-labs/osmosis/v7/x/gamm"
-	"github.com/osmosis-labs/osmosis/v7/x/gamm/pool-models/balancer"
-	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tendermint/tendermint/crypto/ed25519"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
+
+	osmoapp "github.com/osmosis-labs/osmosis/v7/app"
+	"github.com/osmosis-labs/osmosis/v7/x/gamm"
+	"github.com/osmosis-labs/osmosis/v7/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
 )
 
 func TestGammInitGenesis(t *testing.T) {
@@ -47,7 +48,7 @@ func TestGammInitGenesis(t *testing.T) {
 	}, app.AppCodec())
 
 	require.Equal(t, app.GAMMKeeper.GetNextPoolNumberAndIncrement(ctx), uint64(2))
-	poolStored, err := app.GAMMKeeper.GetPool(ctx, 1)
+	poolStored, err := app.GAMMKeeper.GetPoolAndPoke(ctx, 1)
 	require.NoError(t, err)
 	require.Equal(t, balancerPool.GetId(), poolStored.GetId())
 	require.Equal(t, balancerPool.GetAddress(), poolStored.GetAddress())
@@ -58,7 +59,7 @@ func TestGammInitGenesis(t *testing.T) {
 	// require.Equal(t, balancerPool.GetAllPoolAssets(), poolStored.GetAllPoolAssets())
 	require.Equal(t, balancerPool.String(), poolStored.String())
 
-	_, err = app.GAMMKeeper.GetPool(ctx, 2)
+	_, err = app.GAMMKeeper.GetPoolAndPoke(ctx, 2)
 	require.Error(t, err)
 
 	liquidity := app.GAMMKeeper.GetTotalLiquidity(ctx)

diff --git a/x/gamm/keeper/grpc_query.go b/x/gamm/keeper/grpc_query.go
@@ -53,7 +53,7 @@ func (q Querier) Pool(
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	pool, err := q.Keeper.GetPool(sdkCtx, req.PoolId)
+	pool, err := q.Keeper.GetPoolAndPoke(sdkCtx, req.PoolId)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
@@ -85,8 +85,8 @@ func (q Querier) Pools(
 			return err
 		}
 
-		// Use GetPool function because it runs PokeWeights
-		poolI, err = q.Keeper.GetPool(sdkCtx, poolI.GetId())
+		// Use GetPoolAndPoke function because it runs PokeWeights
+		poolI, err = q.Keeper.GetPoolAndPoke(sdkCtx, poolI.GetId())
 		if err != nil {
 			return err
 		}
@@ -130,7 +130,7 @@ func (q Querier) PoolParams(ctx context.Context, req *types.QueryPoolParamsReque
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	pool, err := q.Keeper.GetPool(sdkCtx, req.PoolId)
+	pool, err := q.Keeper.GetPoolAndPoke(sdkCtx, req.PoolId)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
@@ -159,7 +159,7 @@ func (q Querier) TotalPoolLiquidity(ctx context.Context, req *types.QueryTotalPo
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	pool, err := q.Keeper.GetPool(sdkCtx, req.PoolId)
+	pool, err := q.Keeper.GetPoolAndPoke(sdkCtx, req.PoolId)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
@@ -176,7 +176,7 @@ func (q Querier) TotalShares(ctx context.Context, req *types.QueryTotalSharesReq
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	pool, err := q.Keeper.GetPool(sdkCtx, req.PoolId)
+	pool, err := q.Keeper.GetPoolAndPoke(sdkCtx, req.PoolId)
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
@@ -203,7 +203,7 @@ func (q Querier) SpotPrice(ctx context.Context, req *types.QuerySpotPriceRequest
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
-	pool, err := q.Keeper.GetPool(sdkCtx, req.PoolId)
+	pool, err := q.Keeper.GetPoolAndPoke(sdkCtx, req.PoolId)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed to get pool by ID: %s", err)
 	}

diff --git a/x/gamm/keeper/invariants.go b/x/gamm/keeper/invariants.go
@@ -29,7 +29,7 @@ func AllInvariants(keeper Keeper, bk types.BankKeeper) sdk.Invariant {
 // pool assets
 func PoolAccountInvariant(keeper Keeper, bk types.BankKeeper) sdk.Invariant {
 	return func(ctx sdk.Context) (string, bool) {
-		pools, err := keeper.GetPools(ctx)
+		pools, err := keeper.GetPoolsAndPoke(ctx)
 		if err != nil {
 			return sdk.FormatInvariant(types.ModuleName, poolBalanceInvariantName,
 				"\tgamm pool retrieval failed"), true

diff --git a/x/gamm/keeper/multihop_test.go b/x/gamm/keeper/multihop_test.go
@@ -59,7 +59,7 @@ func (suite *KeeperTestSuite) TestBalancerPoolSimpleMultihopSwapExactAmountIn()
 					if i != 0 {
 						tokenInDenom = test.param.routes[i-1].TokenOutDenom
 					}
-					pool, err := keeper.GetPool(suite.ctx, route.PoolId)
+					pool, err := keeper.GetPoolAndPoke(suite.ctx, route.PoolId)
 					suite.NoError(err, "test: %v", test.name)
 
 					sp, err := pool.SpotPrice(suite.ctx, tokenInDenom, route.TokenOutDenom)
@@ -80,7 +80,8 @@ func (suite *KeeperTestSuite) TestBalancerPoolSimpleMultihopSwapExactAmountIn()
 					if i != 0 {
 						tokenInDenom = test.param.routes[i-1].TokenOutDenom
 					}
-					pool, err := keeper.GetPool(suite.ctx, route.PoolId)
+
+					pool, err := keeper.GetPoolAndPoke(suite.ctx, route.PoolId)
 					suite.NoError(err, "test: %v", test.name)
 
 					sp, err := pool.SpotPrice(suite.ctx, tokenInDenom, route.TokenOutDenom)
@@ -151,7 +152,8 @@ func (suite *KeeperTestSuite) TestBalancerPoolSimpleMultihopSwapExactAmountOut()
 					if i != len(test.param.routes)-1 {
 						tokenOutDenom = test.param.routes[i+1].TokenInDenom
 					}
-					pool, err := keeper.GetPool(suite.ctx, route.PoolId)
+
+					pool, err := keeper.GetPoolAndPoke(suite.ctx, route.PoolId)
 					suite.NoError(err, "test: %v", test.name)
 
 					sp, err := pool.SpotPrice(suite.ctx, route.TokenInDenom, tokenOutDenom)
@@ -172,7 +174,8 @@ func (suite *KeeperTestSuite) TestBalancerPoolSimpleMultihopSwapExactAmountOut()
 					if i != len(test.param.routes)-1 {
 						tokenOutDenom = test.param.routes[i+1].TokenInDenom
 					}
-					pool, err := keeper.GetPool(suite.ctx, route.PoolId)
+
+					pool, err := keeper.GetPoolAndPoke(suite.ctx, route.PoolId)
 					suite.NoError(err, "test: %v", test.name)
 
 					sp, err := pool.SpotPrice(suite.ctx, route.TokenInDenom, tokenOutDenom)

diff --git a/x/gamm/keeper/pool.go b/x/gamm/keeper/pool.go
@@ -20,7 +20,9 @@ func (k Keeper) UnmarshalPool(bz []byte) (types.PoolI, error) {
 	return acc, k.cdc.UnmarshalInterface(bz, &acc)
 }
 
-func (k Keeper) GetPool(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
+// GetPoolAndPoke returns a PoolI based on it's identifier if one exists. Prior
+// to returning the pool, the weights of the pool are updated via PokePool.
+func (k Keeper) GetPoolAndPoke(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
 	store := ctx.KVStore(k.storeKey)
 	poolKey := types.GetKeyPrefixPools(poolId)
 	if !store.Has(poolKey) {
@@ -34,14 +36,14 @@ func (k Keeper) GetPool(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
 		return nil, err
 	}
 
-	// pool.PokeTokenWeights(ctx.BlockTime())
+	pool.PokePool(ctx.BlockTime())
 
 	return pool, nil
 }
 
 // Get pool, and check if the pool is active / allowed to be swapped against
 func (k Keeper) getPoolForSwap(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
-	pool, err := k.GetPool(ctx, poolId)
+	pool, err := k.GetPoolAndPoke(ctx, poolId)
 	if err != nil {
 		return &balancer.Pool{}, err
 	}
@@ -57,7 +59,7 @@ func (k Keeper) iterator(ctx sdk.Context, prefix []byte) sdk.Iterator {
 	return sdk.KVStorePrefixIterator(store, prefix)
 }
 
-func (k Keeper) GetPools(ctx sdk.Context) (res []types.PoolI, err error) {
+func (k Keeper) GetPoolsAndPoke(ctx sdk.Context) (res []types.PoolI, err error) {
 	iter := k.iterator(ctx, types.KeyPrefixPools)
 	defer iter.Close()
 
@@ -69,12 +71,11 @@ func (k Keeper) GetPools(ctx sdk.Context) (res []types.PoolI, err error) {
 			return nil, err
 		}
 
-		// pool.PokeTokenWeights(ctx.BlockTime())
-
+		pool.PokePool(ctx.BlockTime())
 		res = append(res, pool)
 	}
 
-	return
+	return res, nil
 }
 
 func (k Keeper) SetPool(ctx sdk.Context, pool types.PoolI) error {

diff --git a/x/gamm/keeper/pool_service.go b/x/gamm/keeper/pool_service.go
@@ -21,10 +21,11 @@ func (k Keeper) CalculateSpotPrice(
 	baseAssetDenom string,
 	quoteAssetDenom string,
 ) (sdk.Dec, error) {
-	pool, err := k.GetPool(ctx, poolID)
+	pool, err := k.GetPoolAndPoke(ctx, poolID)
 	if err != nil {
 		return sdk.Dec{}, err
 	}
+
 	return pool.SpotPrice(ctx, baseAssetDenom, quoteAssetDenom)
 }
 
@@ -182,7 +183,7 @@ func (k Keeper) JoinPoolNoSwap(
 	shareOutAmount sdk.Int,
 	tokenInMaxs sdk.Coins,
 ) (err error) {
-	pool, err := k.GetPool(ctx, poolId)
+	pool, err := k.GetPoolAndPoke(ctx, poolId)
 	if err != nil {
 		return err
 	}
@@ -330,7 +331,7 @@ func (k Keeper) ExitPool(
 	shareInAmount sdk.Int,
 	tokenOutMins sdk.Coins,
 ) (exitCoins sdk.Coins, err error) {
-	pool, err := k.GetPool(ctx, poolId)
+	pool, err := k.GetPoolAndPoke(ctx, poolId)
 	if err != nil {
 		return sdk.Coins{}, err
 	}

diff --git a/x/gamm/keeper/pool_service_test.go b/x/gamm/keeper/pool_service_test.go
@@ -72,7 +72,7 @@ func (suite *KeeperTestSuite) TestCreateBalancerPool() {
 			poolId, err := keeper.CreatePool(suite.ctx, msg)
 			suite.Require().NoError(err)
 
-			pool, err := keeper.GetPool(suite.ctx, poolId)
+			pool, err := keeper.GetPoolAndPoke(suite.ctx, poolId)
 			suite.Require().NoError(err)
 			suite.Require().Equal(types.InitPoolSharesSupply.String(), pool.GetTotalShares().String(),
 				fmt.Sprintf("share token should be minted as %s initially", types.InitPoolSharesSupply.String()),
@@ -220,7 +220,7 @@ func (suite *KeeperTestSuite) TestCreateBalancerPool() {
 			}, defaultPoolAssets, defaultFutureGovernor)
 			_, err := keeper.CreatePool(suite.ctx, msg)
 			suite.Require().NoError(err)
-			pools, err := keeper.GetPools(suite.ctx)
+			pools, err := keeper.GetPoolsAndPoke(suite.ctx)
 			suite.Require().Len(pools, 1)
 			suite.Require().NoError(err)
 		},
@@ -236,7 +236,7 @@ func (suite *KeeperTestSuite) TestCreateBalancerPool() {
 			}, defaultPoolAssets, defaultFutureGovernor)
 			_, err := keeper.CreatePool(suite.ctx, msg)
 			suite.Require().NoError(err)
-			pools, err := keeper.GetPools(suite.ctx)
+			pools, err := keeper.GetPoolsAndPoke(suite.ctx)
 			suite.Require().Len(pools, 1)
 			suite.Require().NoError(err)
 		},

diff --git a/x/gamm/pool-models/balancer/amm_test.go b/x/gamm/pool-models/balancer/amm_test.go
@@ -2,6 +2,7 @@ package balancer_test
 
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
+
 	"github.com/osmosis-labs/osmosis/v7/x/gamm/pool-models/balancer"
 )
 
@@ -56,7 +57,7 @@ func (suite *KeeperTestSuite) TestBalancerSpotPrice() {
 			tc.quoteDenomPoolInput,
 		)
 
-		pool, err := suite.App.GAMMKeeper.GetPool(suite.Ctx, poolId)
+		pool, err := suite.App.GAMMKeeper.GetPoolAndPoke(suite.Ctx, poolId)
 		suite.Require().NoError(err, "test: %s", tc.name)
 		balancerPool, isPool := pool.(*balancer.Pool)
 		suite.Require().True(isPool, "test: %s", tc.name)