Improve endblock performance and fix tracer (#654)

Co-authored-by: Yiming Zang <[email protected]>

app/abci.go

@@ -40,11 +40,8 @@ func (app *App) CheckTx(ctx context.Context, req *abci.RequestCheckTx) (*abci.Re
 
 func (app *App) DeliverTx(ctx sdk.Context, req abci.RequestDeliverTx) abci.ResponseDeliverTx {
 	defer metrics.MeasureDeliverTxDuration(time.Now())
-	tracectx, span := app.tracingInfo.Start("DeliverTx")
-	oldCtx := app.tracingInfo.GetContext()
-	app.tracingInfo.SetContext(tracectx)
+	_, span := app.tracingInfo.Start("DeliverTx")
 	defer span.End()
-	defer func() { app.tracingInfo.SetContext(oldCtx) }()
 	return app.BaseApp.DeliverTx(ctx, req)
 }
 

utils/tracing/tracer.go

@@ -59,7 +59,7 @@ func (i *Info) Start(name string) (context.Context, otrace.Span) {
 	if i.tracerContext == nil {
 		i.tracerContext = context.Background()
 	}
-	return (*i.Tracer).Start(i.tracerContext, "DeliverTx")
+	return (*i.Tracer).Start(i.tracerContext, name)
 }
 
 func (i *Info) GetContext() context.Context {

x/dex/contract/abci.go

@@ -48,7 +48,7 @@ func EndBlockerAtomic(ctx sdk.Context, keeper *keeper.Keeper, validContractsInfo
 	env := newEnv(ctx, validContractsInfo, keeper)
 	cachedCtx, msCached := cacheContext(ctx, env)
 	memStateCopy := dexutils.GetMemState(cachedCtx.Context()).DeepCopy()
-	handleDeposits(cachedCtx, env, keeper, tracer)
+	handleDeposits(spanCtx, cachedCtx, env, keeper, tracer)
 
 	runner := NewParallelRunner(func(contract types.ContractInfoV2) {
 		orderMatchingRunnable(spanCtx, cachedCtx, env, keeper, contract, tracer)
@@ -103,20 +103,17 @@ func newEnv(ctx sdk.Context, validContractsInfo []types.ContractInfoV2, keeper *
 	settlementsByContract := datastructures.NewTypedSyncMap[string, []*types.SettlementEntry]()
 	executionTerminationSignals := datastructures.NewTypedSyncMap[string, chan struct{}]()
 	registeredPairs := datastructures.NewTypedSyncMap[string, []types.Pair]()
-	orderBooks := datastructures.NewTypedNestedSyncMap[string, dextypesutils.PairString, *types.OrderBook]()
+	allContractAndPairs := map[string][]types.Pair{}
 	for _, contract := range validContractsInfo {
 		finalizeBlockMessages.Store(contract.ContractAddr, dextypeswasm.NewSudoFinalizeBlockMsg())
 		settlementsByContract.Store(contract.ContractAddr, []*types.SettlementEntry{})
 		executionTerminationSignals.Store(contract.ContractAddr, make(chan struct{}, 1))
 		contractPairs := keeper.GetAllRegisteredPairs(ctx, contract.ContractAddr)
 		registeredPairs.Store(contract.ContractAddr, contractPairs)
-		for _, pair := range contractPairs {
-			pair := pair
-			orderBooks.StoreNested(contract.ContractAddr, dextypesutils.GetPairString(&pair), dexkeeperutils.PopulateOrderbook(
-				ctx, keeper, dextypesutils.ContractAddress(contract.ContractAddr), pair,
-			))
-		}
+		allContractAndPairs[contract.ContractAddr] = contractPairs
 	}
+	// Parallelize populating orderbooks for performance improvements
+	orderBooks := dexkeeperutils.PopulateAllOrderbooks(ctx, keeper, allContractAndPairs)
 	return &environment{
 		validContractsInfo:          validContractsInfo,
 		failedContractAddresses:     datastructures.NewSyncSet([]string{}),
@@ -146,14 +143,16 @@ func decorateContextForContract(ctx sdk.Context, contractInfo types.ContractInfo
 	)
 }
 
-func handleDeposits(ctx sdk.Context, env *environment, keeper *keeper.Keeper, tracer *otrace.Tracer) {
+func handleDeposits(spanCtx context.Context, ctx sdk.Context, env *environment, keeper *keeper.Keeper, tracer *otrace.Tracer) {
 	// Handle deposit sequentially since they mutate `bank` state which is shared by all contracts
+	_, span := (*tracer).Start(spanCtx, "handleDeposits")
+	defer span.End()
 	keeperWrapper := dexkeeperabci.KeeperWrapper{Keeper: keeper}
 	for _, contract := range env.validContractsInfo {
 		if !contract.NeedOrderMatching {
 			continue
 		}
-		if err := keeperWrapper.HandleEBDeposit(ctx.Context(), ctx, tracer, contract.ContractAddr); err != nil {
+		if err := keeperWrapper.HandleEBDeposit(spanCtx, ctx, tracer, contract.ContractAddr); err != nil {
 			env.failedContractAddresses.Add(contract.ContractAddr)
 		}
 	}
@@ -221,6 +220,8 @@ func handleFinalizedBlocks(ctx context.Context, sdkCtx sdk.Context, env *environ
 }
 
 func orderMatchingRunnable(ctx context.Context, sdkContext sdk.Context, env *environment, keeper *keeper.Keeper, contractInfo types.ContractInfoV2, tracer *otrace.Tracer) {
+	_, span := (*tracer).Start(ctx, "orderMatchingRunnable")
+	defer span.End()
 	defer func() {
 		if channel, ok := env.executionTerminationSignals.Load(contractInfo.ContractAddr); ok {
 			_, err := logging.LogIfNotDoneAfter(sdkContext.Logger(), func() (struct{}, error) {

x/dex/keeper/utils/order_book.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"sort"
+	"sync"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/sei-protocol/sei-chain/utils/datastructures"
@@ -32,6 +33,27 @@ func PopulateOrderbook(
 	}
 }
 
+func PopulateAllOrderbooks(
+	ctx sdk.Context,
+	keeper *keeper.Keeper,
+	contractsAndPairs map[string][]types.Pair,
+) *datastructures.TypedNestedSyncMap[string, dextypesutils.PairString, *types.OrderBook] {
+	var orderBooks = datastructures.NewTypedNestedSyncMap[string, dextypesutils.PairString, *types.OrderBook]()
+	wg := sync.WaitGroup{}
+	for contractAddr, pairs := range contractsAndPairs {
+		for _, pair := range pairs {
+			wg.Add(1)
+			go func(contractAddr string, pair types.Pair) {
+				defer wg.Done()
+				orderBook := PopulateOrderbook(ctx, keeper, dextypesutils.ContractAddress(contractAddr), pair)
+				orderBooks.StoreNested(contractAddr, dextypesutils.GetPairString(&pair), orderBook)
+			}(contractAddr, pair)
+		}
+	}
+	wg.Wait()
+	return orderBooks
+}
+
 func sortOrderBookEntries(entries []types.OrderBookEntry) {
 	sort.SliceStable(entries, func(i, j int) bool {
 		return entries[i].GetPrice().LT(entries[j].GetPrice())