Parallelize begin block for dex module

x/dex/keeper/price.go

@@ -16,12 +16,12 @@ func (k Keeper) SetPriceState(ctx sdk.Context, price types.Price, contractAddr s
 
 func (k Keeper) DeletePriceStateBefore(ctx sdk.Context, contractAddr string, timestamp uint64, pair types.Pair) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.PricePrefix(contractAddr, pair.PriceDenom, pair.AssetDenom))
-	for _, key := range k.getPriceKeysToDelete(store, timestamp) {
+	for _, key := range k.GetPriceKeysToDelete(store, timestamp) {
 		store.Delete(key)
 	}
 }
 
-func (k Keeper) getPriceKeysToDelete(store sdk.KVStore, timestamp uint64) [][]byte {
+func (k Keeper) GetPriceKeysToDelete(store sdk.KVStore, timestamp uint64) [][]byte {
 	keys := [][]byte{}
 	iterator := sdk.KVStorePrefixIterator(store, []byte{})
 	defer iterator.Close()

x/dex/module.go

@@ -4,23 +4,22 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	"sync"
 	"time"
 
 	"github.com/CosmWasm/wasmd/x/wasm"
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	"github.com/gorilla/mux"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cobra"
-	"go.opentelemetry.io/otel/attribute"
-
 	abci "github.com/tendermint/tendermint/abci/types"
 
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/codec"
 	cdctypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
-	seisync "github.com/sei-protocol/sei-chain/sync"
 	"github.com/sei-protocol/sei-chain/utils"
 	"github.com/sei-protocol/sei-chain/utils/datastructures"
 	"github.com/sei-protocol/sei-chain/utils/tracing"
@@ -251,30 +250,54 @@ func (am AppModule) BeginBlock(ctx sdk.Context, _ abci.RequestBeginBlock) {
 		am.keeper.SetEpoch(ctx, currentEpoch)
 	}
 	cachedCtx, cachedStore := store.GetCachedContext(ctx)
-	gasLimit := am.keeper.GetParams(ctx).BeginBlockGasLimit
-	for _, contract := range am.getAllContractInfo(ctx) {
-		am.beginBlockForContract(cachedCtx, contract, gasLimit)
+	priceRetention := am.keeper.GetParams(ctx).PriceSnapshotRetention
+	cutOffTime := uint64(ctx.BlockTime().Unix()) - priceRetention
+	wg := sync.WaitGroup{}
+	allContracts := am.getAllContractInfo(ctx)
+	allPricesToDelete := make(map[string][]*types.PriceStore, len(allContracts))
+
+	// Parallelize the logic to find all prices to delete
+	for _, contract := range allContracts {
+		wg.Add(1)
+		go func(contract types.ContractInfoV2) {
+			priceKeysToDelete := am.getPriceToDelete(cachedCtx, contract, cutOffTime)
+			allPricesToDelete[contract.ContractAddr] = priceKeysToDelete
+			wg.Done()
+		}(contract)
+	}
+	wg.Wait()
+
+	// Execute the deletion in order
+	for _, contract := range allContracts {
+		if priceStores, found := allPricesToDelete[contract.ContractAddr]; found {
+			for _, priceStore := range priceStores {
+				for _, key := range priceStore.PriceKeys {
+					priceStore.Store.Delete(key)
+				}
+			}
+		}
 	}
 	// only write if all contracts have been processed
 	cachedStore.Write()
 }
 
-func (am AppModule) beginBlockForContract(ctx sdk.Context, contract types.ContractInfoV2, gasLimit uint64) {
-	_, span := am.tracingInfo.Start("DexBeginBlock")
-	contractAddr := contract.ContractAddr
-	span.SetAttributes(attribute.String("contract", contractAddr))
-	defer span.End()
-
-	ctx = ctx.WithGasMeter(seisync.NewGasWrapper(dexutils.GetGasMeterForLimit(gasLimit)))
-
+func (am AppModule) getPriceToDelete(
+	ctx sdk.Context,
+	contract types.ContractInfoV2,
+	timestamp uint64,
+) []*types.PriceStore {
+	var result []*types.PriceStore
 	if contract.NeedOrderMatching {
-		currentTimestamp := uint64(ctx.BlockTime().Unix())
-		ctx.Logger().Debug(fmt.Sprintf("Removing stale prices for ts %d", currentTimestamp))
-		priceRetention := am.keeper.GetParams(ctx).PriceSnapshotRetention
-		for _, pair := range am.keeper.GetAllRegisteredPairs(ctx, contractAddr) {
-			am.keeper.DeletePriceStateBefore(ctx, contractAddr, currentTimestamp-priceRetention, pair)
+		for _, pair := range am.keeper.GetAllRegisteredPairs(ctx, contract.ContractAddr) {
+			store := prefix.NewStore(ctx.KVStore(am.keeper.GetStoreKey()), types.PricePrefix(contract.ContractAddr, pair.PriceDenom, pair.AssetDenom))
+			keysToDelete := am.keeper.GetPriceKeysToDelete(store, timestamp)
+			result = append(result, &types.PriceStore{
+				Store:     store,
+				PriceKeys: keysToDelete,
+			})
 		}
 	}
+	return result
 }
 
 // EndBlock executes all ABCI EndBlock logic respective to the capability module. It

x/dex/types/orders.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/sei-protocol/sei-chain/utils"
 	"github.com/sei-protocol/sei-chain/utils/datastructures"
@@ -41,6 +42,11 @@ type OrderBookEntry interface {
 	DeepCopy() OrderBookEntry
 }
 
+type PriceStore struct {
+	Store     prefix.Store
+	PriceKeys [][]byte
+}
+
 func (m *LongBook) GetPrice() sdk.Dec {
 	if m != nil {
 		return m.Price