perf(protorev): base denom as param

CHANGELOG.md

@@ -45,6 +45,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### State Breaking
 
 * [#7250](https://github.com/osmosis-labs/osmosis/pull/7250) Further filter spam gauges from epoch distribution.
+* [#7508](https://github.com/osmosis-labs/osmosis/pull/7508) Improve protorev performance by removing iterator and storing base denoms as a parameter.
 
 ## v23.0.0
 

app/upgrades/v24/upgrades.go

@@ -23,6 +23,17 @@ func CreateUpgradeHandler(
 			return nil, err
 		}
 
+		// We no longer use the KVStore base denoms and instead use the param store for performance reasons.
+		// We retrieve the base denoms from the KVStore, set them in the param store, then delete them here.
+		baseDenoms, err := keepers.ProtoRevKeeper.DeprecatedGetAllBaseDenoms(ctx)
+		if err != nil {
+			return nil, err
+		}
+		protorevParams := keepers.ProtoRevKeeper.GetParams(ctx)
+		protorevParams.BaseDenoms = baseDenoms
+		keepers.ProtoRevKeeper.SetParams(ctx, protorevParams)
+		keepers.ProtoRevKeeper.DeprecatedDeleteBaseDenoms(ctx)
+
 		return migrations, nil
 	}
 }

app/upgrades/v24/upgrades_test.go

@@ -0,0 +1,74 @@
+package v24_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+
+	abci "github.com/cometbft/cometbft/abci/types"
+
+	"github.com/osmosis-labs/osmosis/osmomath"
+	"github.com/osmosis-labs/osmosis/v23/app/apptesting"
+	protorevtypes "github.com/osmosis-labs/osmosis/v23/x/protorev/types"
+)
+
+const (
+	v24UpgradeHeight              = int64(10)
+	HistoricalTWAPTimeIndexPrefix = "historical_time_index"
+	KeySeparator                  = "|"
+)
+
+type UpgradeTestSuite struct {
+	apptesting.KeeperTestHelper
+}
+
+func TestUpgradeTestSuite(t *testing.T) {
+	suite.Run(t, new(UpgradeTestSuite))
+}
+
+func (s *UpgradeTestSuite) TestUpgrade() {
+	s.Setup()
+
+	// Set the old KVStore base denoms
+	s.App.ProtoRevKeeper.DeprecatedSetBaseDenoms(s.Ctx, []protorevtypes.BaseDenom{
+		{Denom: protorevtypes.OsmosisDenomination, StepSize: osmomath.NewInt(1_000_000)},
+		{Denom: "atom", StepSize: osmomath.NewInt(1_000_000)},
+		{Denom: "weth", StepSize: osmomath.NewInt(1_000_000)}})
+	oldBaseDenoms, err := s.App.ProtoRevKeeper.DeprecatedGetAllBaseDenoms(s.Ctx)
+	s.Require().NoError(err)
+	s.Require().Equal(3, len(oldBaseDenoms))
+	s.Require().Equal(oldBaseDenoms[0].Denom, protorevtypes.OsmosisDenomination)
+	s.Require().Equal(oldBaseDenoms[1].Denom, "atom")
+	s.Require().Equal(oldBaseDenoms[2].Denom, "weth")
+
+	// The new param store should return the default value
+	newBaseDenoms := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
+	s.Require().Equal(protorevtypes.DefaultBaseDenoms, newBaseDenoms)
+
+	dummyUpgrade(s)
+	s.Require().NotPanics(func() {
+		s.App.BeginBlocker(s.Ctx, abci.RequestBeginBlock{})
+	})
+
+	// The new param store should return the old KVStore values
+	newBaseDenoms = s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
+	s.Require().Equal(oldBaseDenoms, newBaseDenoms)
+
+	// The old KVStore base denoms should be deleted
+	oldBaseDenoms, err = s.App.ProtoRevKeeper.DeprecatedGetAllBaseDenoms(s.Ctx)
+	s.Require().NoError(err)
+	s.Require().Empty(oldBaseDenoms)
+}
+
+func dummyUpgrade(s *UpgradeTestSuite) {
+	s.Ctx = s.Ctx.WithBlockHeight(v24UpgradeHeight - 1)
+	plan := upgradetypes.Plan{Name: "v24", Height: v24UpgradeHeight}
+	err := s.App.UpgradeKeeper.ScheduleUpgrade(s.Ctx, plan)
+	s.Require().NoError(err)
+	_, exists := s.App.UpgradeKeeper.GetUpgradePlan(s.Ctx)
+	s.Require().True(exists)
+
+	s.Ctx = s.Ctx.WithBlockHeight(v24UpgradeHeight)
+}

proto/osmosis/protorev/v1beta1/params.proto

@@ -12,4 +12,25 @@ message Params {
   bool enabled = 1 [ (gogoproto.moretags) = "yaml:\"enabled\"" ];
   // The admin account (settings manager) of the protorev module.
   string admin = 2 [ (gogoproto.moretags) = "yaml:\"admin\"" ];
+  // All enabled base denominations for the protorev module.
+  repeated BaseDenom base_denoms = 3 [
+    (gogoproto.nullable) = false,
+    (gogoproto.moretags) = "yaml:\"base_denoms\""
+  ];
+}
+
+// BaseDenom represents a single base denom that the module uses for its
+// arbitrage trades. It contains the denom name alongside the step size of the
+// binary search that is used to find the optimal swap amount
+message BaseDenom {
+  // The denom i.e. name of the base denom (ex. uosmo)
+  string denom = 1 [ (gogoproto.moretags) = "yaml:\"denom\"" ];
+  // The step size of the binary search that is used to find the optimal swap
+  // amount
+  string step_size = 2 [
+
+    (gogoproto.customtype) = "cosmossdk.io/math.Int",
+    (gogoproto.nullable) = false,
+    (gogoproto.moretags) = "yaml:\"step_size\""
+  ];
 }

proto/osmosis/protorev/v1beta1/protorev.proto

@@ -168,22 +168,6 @@ message WeightMap {
       [ (gogoproto.moretags) = "yaml:\"contract_address\"" ];
 }
 
-// BaseDenom represents a single base denom that the module uses for its
-// arbitrage trades. It contains the denom name alongside the step size of the
-// binary search that is used to find the optimal swap amount
-message BaseDenom {
-  // The denom i.e. name of the base denom (ex. uosmo)
-  string denom = 1 [ (gogoproto.moretags) = "yaml:\"denom\"" ];
-  // The step size of the binary search that is used to find the optimal swap
-  // amount
-  string step_size = 2 [
-
-    (gogoproto.customtype) = "cosmossdk.io/math.Int",
-    (gogoproto.nullable) = false,
-    (gogoproto.moretags) = "yaml:\"step_size\""
-  ];
-}
-
 message AllProtocolRevenue {
   osmosis.poolmanager.v1beta1.TakerFeesTracker taker_fees_tracker = 1 [
     (gogoproto.moretags) = "yaml:\"taker_fees_tracker\"",

proto/osmosis/protorev/v1beta1/tx.proto

@@ -6,6 +6,7 @@ import "amino/amino.proto";
 import "google/api/annotations.proto";
 import "osmosis/protorev/v1beta1/protorev.proto";
 import "cosmos_proto/cosmos.proto";
+import "osmosis/protorev/v1beta1/params.proto";
 
 option go_package = "github.com/osmosis-labs/osmosis/v23/x/protorev/types";
 

x/protorev/keeper/epoch_hook.go

@@ -53,10 +53,7 @@ func (k Keeper) UpdatePools(ctx sdk.Context) error {
 	// baseDenomPools maps each base denom to a map of the highest liquidity pools paired with that base denom
 	// ex. {osmo -> {atom : 100, weth : 200}}
 	baseDenomPools := make(map[string]map[string]LiquidityPoolStruct)
-	baseDenoms, err := k.GetAllBaseDenoms(ctx)
-	if err != nil {
-		return err
-	}
+	baseDenoms := k.GetAllBaseDenoms(ctx)
 
 	// Delete any pools that currently exist in the store + initialize baseDenomPools
 	for _, baseDenom := range baseDenoms {

x/protorev/keeper/epoch_hook_test.go

@@ -47,8 +47,7 @@ func (s *KeeperTestSuite) TestEpochHook() {
 
 	totalNumberExpected := 0
 	expectedToSee := make(map[string]Pool)
-	baseDenoms, err := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
+	baseDenoms := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 	for _, pool := range s.pools {
 		// Module currently limited to two asset pools
 		// Instantiate asset and amounts for the pool

x/protorev/keeper/genesis.go

@@ -28,9 +28,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, genState types.GenesisState) {
 	// Configure the initial base denoms used for cyclic route building. The order of the list of base
 	// denoms is the order in which routes will be prioritized i.e. routes will be built and simulated in a
 	// first come first serve basis that is based on the order of the base denoms.
-	if err := k.SetBaseDenoms(ctx, genState.BaseDenoms); err != nil {
-		panic(err)
-	}
+	k.SetBaseDenoms(ctx, genState.BaseDenoms)
 
 	// Update the pools on genesis.
 	if err := k.UpdatePools(ctx); err != nil {
@@ -122,10 +120,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 	genesis.TokenPairArbRoutes = routes
 
 	// Export the base denoms used for cyclic route building.
-	baseDenoms, err := k.GetAllBaseDenoms(ctx)
-	if err != nil {
-		panic(err)
-	}
+	baseDenoms := k.GetAllBaseDenoms(ctx)
 	genesis.BaseDenoms = baseDenoms
 
 	// Export the developer fees that have been collected.

x/protorev/keeper/genesis_test.go

@@ -15,8 +15,7 @@ func (s *KeeperTestSuite) TestInitGenesis() {
 		s.Require().Contains(tokenPairArbRoutes, route)
 	}
 
-	baseDenoms, err := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
+	baseDenoms := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 	s.Require().Equal(len(baseDenoms), len(exportedGenesis.BaseDenoms))
 	for _, baseDenom := range exportedGenesis.BaseDenoms {
 		s.Require().Contains(baseDenoms, baseDenom)

x/protorev/keeper/grpc_query.go

@@ -214,10 +214,7 @@ func (q Querier) GetProtoRevBaseDenoms(c context.Context, req *types.QueryGetPro
 		return nil, status.Error(codes.InvalidArgument, "invalid request")
 	}
 	ctx := sdk.UnwrapSDKContext(c)
-	baseDenoms, err := q.Keeper.GetAllBaseDenoms(ctx)
-	if err != nil {
-		return nil, status.Error(codes.Internal, err.Error())
-	}
+	baseDenoms := q.Keeper.GetAllBaseDenoms(ctx)
 
 	return &types.QueryGetProtoRevBaseDenomsResponse{BaseDenoms: baseDenoms}, nil
 }

x/protorev/keeper/grpc_query_test.go

@@ -330,8 +330,7 @@ func (s *KeeperTestSuite) TestGetProtoRevMaxPoolPointsPerBlock() {
 // TestGetProtoRevBaseDenoms tests the query to retrieve the base denoms
 func (s *KeeperTestSuite) TestGetProtoRevBaseDenoms() {
 	// base denoms already set in setup
-	baseDenoms, err := s.App.AppKeepers.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
+	baseDenoms := s.App.AppKeepers.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 
 	req := &types.QueryGetProtoRevBaseDenomsRequest{}
 	res, err := s.queryClient.GetProtoRevBaseDenoms(sdk.WrapSDKContext(s.Ctx), req)

x/protorev/keeper/hooks.go

@@ -154,11 +154,7 @@ func (k Keeper) StoreJoinExitPoolSwaps(ctx sdk.Context, sender sdk.AccAddress, p
 // AfterPoolCreatedWithCoins checks if the new pool should be stored as the highest liquidity pool
 // for any of the base denoms, and stores it if so.
 func (k Keeper) AfterPoolCreatedWithCoins(ctx sdk.Context, poolId uint64) {
-	baseDenoms, err := k.GetAllBaseDenoms(ctx)
-	if err != nil {
-		ctx.Logger().Error("Protorev error getting base denoms in AfterCFMMPoolCreated hook: " + err.Error())
-		return
-	}
+	baseDenoms := k.GetAllBaseDenoms(ctx)
 
 	baseDenomMap := make(map[string]bool)
 	for _, baseDenom := range baseDenoms {

x/protorev/keeper/keeper_test.go

@@ -94,8 +94,7 @@ func (s *KeeperTestSuite) SetupTest() {
 			StepSize: osmomath.NewInt(1_000_000),
 		},
 	}
-	err := s.App.ProtoRevKeeper.SetBaseDenoms(s.Ctx, baseDenomPriorities)
-	s.Require().NoError(err)
+	s.App.ProtoRevKeeper.SetBaseDenoms(s.Ctx, baseDenomPriorities)
 
 	encodingConfig := osmosisapp.MakeEncodingConfig()
 	s.clientCtx = client.Context{}.
@@ -151,7 +150,7 @@ func (s *KeeperTestSuite) SetupTest() {
 
 	// Set the Admin Account
 	s.adminAccount = apptesting.CreateRandomAccounts(1)[0]
-	err = protorev.HandleSetProtoRevAdminAccount(s.Ctx, *s.App.ProtoRevKeeper, &types.SetProtoRevAdminAccountProposal{Account: s.adminAccount.String()})
+	err := protorev.HandleSetProtoRevAdminAccount(s.Ctx, *s.App.ProtoRevKeeper, &types.SetProtoRevAdminAccountProposal{Account: s.adminAccount.String()})
 	s.Require().NoError(err)
 
 	queryHelper := baseapp.NewQueryServerTestHelper(s.Ctx, s.App.InterfaceRegistry())

x/protorev/keeper/msg_server.go

@@ -139,22 +139,17 @@ func (m MsgServer) SetBaseDenoms(c context.Context, msg *types.MsgSetBaseDenoms)
 	}
 
 	// Get the old base denoms
-	baseDenoms, err := m.k.GetAllBaseDenoms(ctx)
-	if err != nil {
-		return nil, err
-	}
+	oldBaseDenoms := m.k.GetAllBaseDenoms(ctx)
 
 	// Delete all pools associated with the base denoms
-	for _, baseDenom := range baseDenoms {
+	for _, baseDenom := range oldBaseDenoms {
 		m.k.DeleteAllPoolsForBaseDenom(ctx, baseDenom.Denom)
 	}
 
-	// Delete the old base denoms
-	m.k.DeleteBaseDenoms(ctx)
+	// // Delete the old base denoms
+	// m.k.DeleteBaseDenoms(ctx)
 
-	if err := m.k.SetBaseDenoms(ctx, msg.BaseDenoms); err != nil {
-		return nil, err
-	}
+	m.k.SetBaseDenoms(ctx, msg.BaseDenoms)
 
 	// Update all of the pools
 	if err := m.k.UpdatePools(ctx); err != nil {

x/protorev/keeper/msg_server_test.go

@@ -655,8 +655,7 @@ func (s *KeeperTestSuite) TestMsgSetBaseDenoms() {
 				s.Require().NoError(err)
 				s.Require().Equal(response, &types.MsgSetBaseDenomsResponse{})
 
-				baseDenoms, err := s.App.AppKeepers.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-				s.Require().NoError(err)
+				baseDenoms := s.App.AppKeepers.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 				s.Require().Equal(testCase.baseDenoms, baseDenoms)
 			} else {
 				s.Require().Error(err)

x/protorev/keeper/protorev.go

@@ -76,8 +76,9 @@ func (k Keeper) DeleteAllTokenPairArbRoutes(ctx sdk.Context) {
 	k.DeleteAllEntriesForKeyPrefix(ctx, types.KeyPrefixTokenPairRoutes)
 }
 
-// GetAllBaseDenoms returns all of the base denoms (sorted by priority in descending order) used to build cyclic arbitrage routes
-func (k Keeper) GetAllBaseDenoms(ctx sdk.Context) ([]types.BaseDenom, error) {
+// DeprecatedGetAllBaseDenoms returns all of the base denoms (sorted by priority in descending order) used to build cyclic arbitrage routes
+// After v24 upgrade, this method should be deleted. We now use the param store.
+func (k Keeper) DeprecatedGetAllBaseDenoms(ctx sdk.Context) ([]types.BaseDenom, error) {
 	baseDenoms := make([]types.BaseDenom, 0)
 
 	store := ctx.KVStore(k.storeKey)
@@ -97,9 +98,10 @@ func (k Keeper) GetAllBaseDenoms(ctx sdk.Context) ([]types.BaseDenom, error) {
 	return baseDenoms, nil
 }
 
-// SetBaseDenoms sets all of the base denoms used to build cyclic arbitrage routes. The base denoms priority
+// DeprecatedSetBaseDenoms sets all of the base denoms used to build cyclic arbitrage routes. The base denoms priority
 // order is going to match the order of the base denoms in the slice.
-func (k Keeper) SetBaseDenoms(ctx sdk.Context, baseDenoms []types.BaseDenom) error {
+// After v24 upgrade, this method should be deleted. We now use the param store.
+func (k Keeper) DeprecatedSetBaseDenoms(ctx sdk.Context, baseDenoms []types.BaseDenom) error {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixBaseDenoms)
 
 	for i, baseDenom := range baseDenoms {
@@ -115,8 +117,19 @@ func (k Keeper) SetBaseDenoms(ctx sdk.Context, baseDenoms []types.BaseDenom) err
 	return nil
 }
 
-// DeleteBaseDenoms deletes all of the base denoms
-func (k Keeper) DeleteBaseDenoms(ctx sdk.Context) {
+func (k Keeper) GetAllBaseDenoms(ctx sdk.Context) []types.BaseDenom {
+	return k.GetParams(ctx).BaseDenoms
+}
+
+func (k Keeper) SetBaseDenoms(ctx sdk.Context, newBaseDenoms []types.BaseDenom) {
+	params := k.GetParams(ctx)
+	params.BaseDenoms = newBaseDenoms
+	k.SetParams(ctx, params)
+}
+
+// DeprecatedDeleteBaseDenoms deletes all of the base denoms.
+// After v24 upgrade, this method should be deleted. We now use the param store.
+func (k Keeper) DeprecatedDeleteBaseDenoms(ctx sdk.Context) {
 	k.DeleteAllEntriesForKeyPrefix(ctx, types.KeyPrefixBaseDenoms)
 }
 

x/protorev/keeper/protorev_test.go

@@ -60,26 +60,17 @@ func (s *KeeperTestSuite) TestDeleteAllTokenPairArbRoutes() {
 // TestGetAllBaseDenoms tests the GetAllBaseDenoms, SetBaseDenoms, and DeleteBaseDenoms functions.
 func (s *KeeperTestSuite) TestGetAllBaseDenoms() {
 	// Should be initialized on genesis
-	baseDenoms, err := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
+	baseDenoms := s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 	s.Require().Equal(3, len(baseDenoms))
 	s.Require().Equal(baseDenoms[0].Denom, types.OsmosisDenomination)
 	s.Require().Equal(baseDenoms[1].Denom, "Atom")
 	s.Require().Equal(baseDenoms[2].Denom, "ibc/0CD3A0285E1341859B5E86B6AB7682F023D03E97607CCC1DC95706411D866DF7")
 
-	// Should be able to delete all base denoms
-	s.App.ProtoRevKeeper.DeleteBaseDenoms(s.Ctx)
-	baseDenoms, err = s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
-	s.Require().Equal(0, len(baseDenoms))
-
 	// Should be able to set the base denoms
-	err = s.App.ProtoRevKeeper.SetBaseDenoms(s.Ctx, []types.BaseDenom{{Denom: "osmo"}, {Denom: "atom"}, {Denom: "weth"}})
-	s.Require().NoError(err)
-	baseDenoms, err = s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
-	s.Require().NoError(err)
+	s.App.ProtoRevKeeper.SetBaseDenoms(s.Ctx, []types.BaseDenom{{Denom: types.OsmosisDenomination}, {Denom: "atom"}, {Denom: "weth"}})
+	baseDenoms = s.App.ProtoRevKeeper.GetAllBaseDenoms(s.Ctx)
 	s.Require().Equal(3, len(baseDenoms))
-	s.Require().Equal(baseDenoms[0].Denom, "osmo")
+	s.Require().Equal(baseDenoms[0].Denom, types.OsmosisDenomination)
 	s.Require().Equal(baseDenoms[1].Denom, "atom")
 	s.Require().Equal(baseDenoms[2].Denom, "weth")
 }

x/protorev/keeper/routes.go

@@ -91,10 +91,7 @@ func (k Keeper) BuildHotRoute(ctx sdk.Context, route types.Route, poolId uint64)
 // and routes are built in a greedy manner.
 func (k Keeper) BuildHighestLiquidityRoutes(ctx sdk.Context, tokenIn, tokenOut string, poolId uint64) ([]RouteMetaData, error) {
 	routes := make([]RouteMetaData, 0)
-	baseDenoms, err := k.GetAllBaseDenoms(ctx)
-	if err != nil {
-		return routes, err
-	}
+	baseDenoms := k.GetAllBaseDenoms(ctx)
 
 	// Iterate through all denoms greedily. When simulating and executing trades, routes that are closer to the beginning of the list
 	// have priority over those that are later in the list. This way we can build routes that are more likely to succeed and bring in