Speedup protorev by having route calculation phase not unmarshal pools (

#8157) (#8164) * Another expected ~1% state machine speedup * Add changelog * minor other improvement * Reduce addr decoding overhead (cherry picked from commit 64a9078) Co-authored-by: Dev Ojha <[email protected]>
osmosis-labs · Apr 29, 2024 · a969e10 · a969e10
1 parent 24c4128
commit a969e10
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 12 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -54,6 +54,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * [#8136](https://github.com/osmosis-labs/osmosis/pull/8136) Don't allow gauge creation/addition with rewards that have no protorev route (i.e. no way to determine if rewards meet minimum epoch value distribution requirements)
 * [#8144](https://github.com/osmosis-labs/osmosis/pull/8144) IBC wasm clients can now make stargate queries and support abort.
 * [#8147](https://github.com/osmosis-labs/osmosis/pull/8147) Process unbonding locks once per minute, rather than every single block.
+* [#8157](https://github.com/osmosis-labs/osmosis/pull/8157) Speedup protorev by not unmarshalling pools in cost-estimation phase.
+
 
 ### State Compatible
 

diff --git a/x/concentrated-liquidity/pool.go b/x/concentrated-liquidity/pool.go
@@ -215,8 +215,9 @@ func (k Keeper) GetTotalPoolLiquidity(ctx sdk.Context, poolId uint64) (sdk.Coins
 		return nil, err
 	}
 
-	token0Bal := k.bankKeeper.GetBalance(ctx, pool.GetAddress(), pool.GetToken0())
-	token1Bal := k.bankKeeper.GetBalance(ctx, pool.GetAddress(), pool.GetToken1())
+	addr := pool.GetAddress()
+	token0Bal := k.bankKeeper.GetBalance(ctx, addr, pool.GetToken0())
+	token1Bal := k.bankKeeper.GetBalance(ctx, addr, pool.GetToken1())
 
 	return sdk.NewCoins(token0Bal, token1Bal), nil
 }

diff --git a/x/poolmanager/create_pool.go b/x/poolmanager/create_pool.go
@@ -157,12 +157,30 @@ func (k *Keeper) SetPoolRoute(ctx sdk.Context, poolId uint64, poolType types.Poo
 }
 
 type poolModuleCacheValue struct {
+	pooltype types.PoolType
 	module   types.PoolModuleI
 	gasFlat  uint64
 	gasKey   uint64
 	gasValue uint64
 }
 
+func (k *Keeper) GetPoolType(ctx sdk.Context, poolId uint64) (types.PoolType, error) {
+	poolModuleCandidate, cacheHit := k.cachedPoolModules.Load(poolId)
+	if !cacheHit {
+		_, err := k.GetPoolModule(ctx, poolId)
+		if err != nil {
+			return 0, err
+		}
+		poolModuleCandidate, _ = k.cachedPoolModules.Load(poolId)
+	}
+	v, _ := poolModuleCandidate.(poolModuleCacheValue)
+	if cacheHit {
+		osmoutils.ChargeMockReadGas(ctx, v.gasFlat, v.gasKey, v.gasValue)
+	}
+
+	return v.pooltype, nil
+}
+
 // GetPoolModule returns the swap module for the given pool ID.
 // Returns error if:
 // - any database error occurs.
@@ -195,6 +213,7 @@ func (k *Keeper) GetPoolModule(ctx sdk.Context, poolId uint64) (types.PoolModule
 	}
 
 	k.cachedPoolModules.Store(poolId, poolModuleCacheValue{
+		pooltype: moduleRoute.PoolType,
 		module:   swapModule,
 		gasFlat:  gasFlat,
 		gasKey:   gasKey,

diff --git a/x/poolmanager/router_test.go b/x/poolmanager/router_test.go
@@ -249,6 +249,30 @@ func (s *KeeperTestSuite) TestGetPoolModule() {
 	}
 }
 
+// TestGetPoolTypeGas tests that the result for GetPoolType charges the
+// same gas whether its a cache hit or cache fail.
+func (s *KeeperTestSuite) TestGetPoolTypeGas() {
+	s.SetupTest()
+	poolmanagerKeeper := s.App.PoolManagerKeeper
+
+	createdPoolId := s.CreatePoolFromType(types.Balancer)
+
+	// cache miss
+	s.App.PoolManagerKeeper.ResetCaches()
+	startGas := s.Ctx.GasMeter().GasConsumed()
+	_, err := poolmanagerKeeper.GetPoolType(s.Ctx, createdPoolId)
+	s.Require().NoError(err)
+	endGas := s.Ctx.GasMeter().GasConsumed()
+	cacheMissGas := endGas - startGas
+
+	startGas = s.Ctx.GasMeter().GasConsumed()
+	_, err = poolmanagerKeeper.GetPoolType(s.Ctx, createdPoolId)
+	s.Require().NoError(err)
+	endGas = s.Ctx.GasMeter().GasConsumed()
+	cacheHitGas := endGas - startGas
+	s.Require().Equal(cacheMissGas, cacheHitGas)
+}
+
 func (s *KeeperTestSuite) TestRouteGetPoolDenoms() {
 	tests := map[string]struct {
 		poolId            uint64

diff --git a/x/protorev/keeper/routes.go b/x/protorev/keeper/routes.go
@@ -226,17 +226,12 @@ func (k Keeper) CalculateRoutePoolPoints(ctx sdk.Context, route poolmanagertypes
 	totalWeight := uint64(0)
 
 	for _, poolId := range route.PoolIds() {
-		pool, err := k.poolmanagerKeeper.GetPool(ctx, poolId)
+		poolType, err := k.poolmanagerKeeper.GetPoolType(ctx, poolId)
 		if err != nil {
 			return 0, err
 		}
 
-		// Ensure that all of the pools in the route exist and are active
-		if err := k.IsValidPool(ctx, pool); err != nil {
-			return 0, err
-		}
-
-		switch pool.GetType() {
+		switch poolType {
 		case poolmanagertypes.Balancer:
 			totalWeight += infoByPoolType.Balancer.Weight
 		case poolmanagertypes.Stableswap:
@@ -245,8 +240,13 @@ func (k Keeper) CalculateRoutePoolPoints(ctx sdk.Context, route poolmanagertypes
 			totalWeight += infoByPoolType.Concentrated.Weight
 		case poolmanagertypes.CosmWasm:
 			weight, ok := uint64(0), false
+			pool, err := k.poolmanagerKeeper.GetPool(ctx, poolId)
+			if err != nil {
+				return 0, err
+			}
+			poolAddrString := pool.GetAddress().String()
 			for _, weightMap := range infoByPoolType.Cosmwasm.WeightMaps {
-				if weightMap.ContractAddress == pool.GetAddress().String() {
+				if weightMap.ContractAddress == poolAddrString {
 					weight = weightMap.Weight
 					ok = true
 					break

diff --git a/x/protorev/keeper/statistics.go b/x/protorev/keeper/statistics.go
@@ -39,7 +39,7 @@ func (k Keeper) IncrementNumberOfTrades(ctx sdk.Context) error {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixNumberOfTrades)
 
 	numberOfTrades, _ := k.GetNumberOfTrades(ctx)
-	numberOfTrades = numberOfTrades.Add(osmomath.OneInt())
+	numberOfTrades = numberOfTrades.Add(oneInt)
 
 	bz, err := numberOfTrades.Marshal()
 	if err != nil {
@@ -189,7 +189,7 @@ func (k Keeper) IncrementTradesByRoute(ctx sdk.Context, route []uint64) error {
 	key := types.GetKeyPrefixTradesByRoute(route)
 
 	trades, _ := k.GetTradesByRoute(ctx, route)
-	trades = trades.Add(osmomath.OneInt())
+	trades = trades.Add(oneInt)
 	bz, err := trades.Marshal()
 	if err != nil {
 		return err

diff --git a/x/protorev/types/expected_keepers.go b/x/protorev/types/expected_keepers.go
@@ -59,6 +59,7 @@ type PoolManagerKeeper interface {
 		ctx sdk.Context,
 		poolId uint64,
 	) (poolmanagertypes.PoolI, error)
+	GetPoolType(ctx sdk.Context, poolId uint64) (poolmanagertypes.PoolType, error)
 	GetPoolModule(ctx sdk.Context, poolId uint64) (poolmanagertypes.PoolModuleI, error)
 	GetTotalPoolLiquidity(ctx sdk.Context, poolId uint64) (sdk.Coins, error)
 	RouteGetPoolDenoms(ctx sdk.Context, poolId uint64) ([]string, error)