Problem: pre-estimation don't run in parallel (#523)

* Problem: pre-estimation don't run in parallel * fix build * fix race * cleanup chunking * keep unchanged --------- Co-authored-by: mmsqe <[email protected]>
crypto-org-chain · Sep 12, 2024 · 56f8a5b · 56f8a5b
1 parent 79bb39e
commit 56f8a5b
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 44 deletions.
diff --git a/app/executor.go b/app/executor.go
@@ -3,6 +3,7 @@ package app
 import (
 	"context"
 	"io"
+	"sync"
 	"sync/atomic"
 
 	"cosmossdk.io/collections"
@@ -21,6 +22,8 @@ import (
 	blockstm "github.com/crypto-org-chain/go-block-stm"
 )
 
+const MinimalParallelPreEstimate = 16
+
 func DefaultTxExecutor(_ context.Context,
 	txs [][]byte,
 	ms storetypes.MultiStore,
@@ -73,17 +76,14 @@ func STMTxExecutor(
 			incarnationCache[i].Store(&m)
 		}
 
-		var estimates map[int]blockstm.MultiLocations
-		memTxs := make([]sdk.Tx, len(txs))
+		var (
+			estimates []blockstm.MultiLocations
+			memTxs    []sdk.Tx
+		)
 		if estimate {
-			for i, rawTx := range txs {
-				if memTx, err := txDecoder(rawTx); err == nil {
-					memTxs[i] = memTx
-				}
-			}
 			// pre-estimation
 			evmDenom := evmKeeper.GetParams(sdk.NewContext(ms, cmtproto.Header{}, false, log.NewNopLogger())).EvmDenom
-			estimates = preEstimates(memTxs, authStore, bankStore, evmDenom)
+			memTxs, estimates = preEstimates(txs, workers, authStore, bankStore, evmDenom, txDecoder)
 		}
 
 		if err := blockstm.ExecuteBlockWithEstimates(
@@ -103,7 +103,11 @@ func STMTxExecutor(
 					cache = *v
 				}
 
-				results[txn] = deliverTxWithMultiStore(int(txn), memTxs[txn], msWrapper{ms}, cache)
+				var memTx sdk.Tx
+				if memTxs != nil {
+					memTx = memTxs[txn]
+				}
+				results[txn] = deliverTxWithMultiStore(int(txn), memTx, msWrapper{ms}, cache)
 
 				if v != nil {
 					incarnationCache[txn].Store(v)
@@ -188,40 +192,65 @@ func (ms stmMultiStoreWrapper) GetObjKVStore(key storetypes.StoreKey) storetypes
 
 // preEstimates returns a static estimation of the written keys for each transaction.
 // NOTE: make sure it sync with the latest sdk logic when sdk upgrade.
-func preEstimates(txs []sdk.Tx, authStore, bankStore int, evmDenom string) map[int]blockstm.MultiLocations {
-	estimates := make(map[int]blockstm.MultiLocations, len(txs))
-	for i, tx := range txs {
-		feeTx, ok := tx.(sdk.FeeTx)
-		if !ok {
-			continue
-		}
-		feePayer := sdk.AccAddress(feeTx.FeePayer())
+func preEstimates(txs [][]byte, workers, authStore, bankStore int, evmDenom string, txDecoder sdk.TxDecoder) ([]sdk.Tx, []blockstm.MultiLocations) {
+	memTxs := make([]sdk.Tx, len(txs))
+	estimates := make([]blockstm.MultiLocations, len(txs))
+
+	job := func(start, end int) {
+		for i := start; i < end; i++ {
+			rawTx := txs[i]
+			tx, err := txDecoder(rawTx)
+			if err != nil {
+				continue
+			}
+			memTxs[i] = tx
 
-		// account key
-		accKey, err := collections.EncodeKeyWithPrefix(
-			authtypes.AddressStoreKeyPrefix,
-			sdk.AccAddressKey,
-			feePayer,
-		)
-		if err != nil {
-			continue
-		}
+			feeTx, ok := tx.(sdk.FeeTx)
+			if !ok {
+				continue
+			}
+			feePayer := sdk.AccAddress(feeTx.FeePayer())
+
+			// account key
+			accKey, err := collections.EncodeKeyWithPrefix(
+				authtypes.AddressStoreKeyPrefix,
+				sdk.AccAddressKey,
+				feePayer,
+			)
+			if err != nil {
+				continue
+			}
 
-		// balance key
-		balanceKey, err := collections.EncodeKeyWithPrefix(
-			banktypes.BalancesPrefix,
-			collections.PairKeyCodec(sdk.AccAddressKey, collections.StringKey),
-			collections.Join(feePayer, evmDenom),
-		)
-		if err != nil {
-			continue
-		}
+			// balance key
+			balanceKey, err := collections.EncodeKeyWithPrefix(
+				banktypes.BalancesPrefix,
+				collections.PairKeyCodec(sdk.AccAddressKey, collections.StringKey),
+				collections.Join(feePayer, evmDenom),
+			)
+			if err != nil {
+				continue
+			}
 
-		estimates[i] = blockstm.MultiLocations{
-			authStore: {accKey},
-			bankStore: {balanceKey},
+			estimates[i] = blockstm.MultiLocations{
+				authStore: {accKey},
+				bankStore: {balanceKey},
+			}
 		}
 	}
 
-	return estimates
+	blockSize := len(txs)
+	chunk := (blockSize + workers - 1) / workers
+	var wg sync.WaitGroup
+	for i := 0; i < blockSize; i += chunk {
+		start := i
+		end := min(i+chunk, blockSize)
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			job(start, end)
+		}()
+	}
+	wg.Wait()
+
+	return memTxs, estimates
 }
diff --git a/go.mod b/go.mod
@@ -29,7 +29,7 @@ require (
 	github.com/cosmos/ibc-go/modules/capability v1.0.0
 	github.com/cosmos/ibc-go/v8 v8.1.0
 	github.com/cosmos/rosetta v0.50.3-1
-	github.com/crypto-org-chain/go-block-stm v0.0.0-20240911081142-92839e79a3ae
+	github.com/crypto-org-chain/go-block-stm v0.0.0-20240912024944-1cd89976aa5e
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc
 	github.com/ethereum/go-ethereum v1.10.26
 	github.com/gogo/protobuf v1.3.2

diff --git a/go.sum b/go.sum
@@ -419,8 +419,8 @@ github.com/crypto-org-chain/cosmos-sdk/store v0.0.0-20240911084450-6870ba130be2
 github.com/crypto-org-chain/cosmos-sdk/store v0.0.0-20240911084450-6870ba130be2/go.mod h1:gjE3DZe4t/+VeIk6CmrouyqiuDbZ7QOVDDq3nLqBTpg=
 github.com/crypto-org-chain/cosmos-sdk/x/tx v0.0.0-20240911084450-6870ba130be2 h1:mxlOSCru7YgmX055rrlkCSUu0D8lAqJ8Dnhp0yXCBuM=
 github.com/crypto-org-chain/cosmos-sdk/x/tx v0.0.0-20240911084450-6870ba130be2/go.mod h1:RTiTs4hkXG6IvYGknvB8p79YgjYJdcbzLUOGJChsPnY=
-github.com/crypto-org-chain/go-block-stm v0.0.0-20240911081142-92839e79a3ae h1:gakWYsVubWX8P9NpxaPnvg0UJYfIZigfko5WN57t7OA=
-github.com/crypto-org-chain/go-block-stm v0.0.0-20240911081142-92839e79a3ae/go.mod h1:iwQTX9xMX8NV9k3o2BiWXA0SswpsZrDk5q3gA7nWYiE=
+github.com/crypto-org-chain/go-block-stm v0.0.0-20240912024944-1cd89976aa5e h1:FFpE6+Y4o5GxkeGwUcETM6amgohh7msWvWf1MDqueVc=
+github.com/crypto-org-chain/go-block-stm v0.0.0-20240912024944-1cd89976aa5e/go.mod h1:iwQTX9xMX8NV9k3o2BiWXA0SswpsZrDk5q3gA7nWYiE=
 github.com/crypto-org-chain/go-ethereum v1.10.20-0.20240425065928-ebb09502e7a7 h1:V43F3JFcqG4MUThf9W/DytnPblpR6CcaLBw2Wx6zTgE=
 github.com/crypto-org-chain/go-ethereum v1.10.20-0.20240425065928-ebb09502e7a7/go.mod h1:+a8pUj1tOyJ2RinsNQD4326YS+leSoKGiG/uVVb0x6Y=
 github.com/danieljoos/wincred v1.2.0 h1:ozqKHaLK0W/ii4KVbbvluM91W2H3Sh0BncbUNPS7jLE=

diff --git a/gomod2nix.toml b/gomod2nix.toml
@@ -203,8 +203,8 @@ schema = 3
     version = "v0.0.24"
     hash = "sha256-4vUukHONOjNn0qfQr4esK6TWfPWsIp+rbdz65og84lw="
   [mod."github.com/crypto-org-chain/go-block-stm"]
-    version = "v0.0.0-20240911081142-92839e79a3ae"
-    hash = "sha256-8MhSeC5BB5BwOet3k3Rfua0TBeBATtLSwaW6s5WpYCM="
+    version = "v0.0.0-20240912024944-1cd89976aa5e"
+    hash = "sha256-rY8W4dSciOXT29MCySbH5sw0Fp15IQVgBK9QlMX0JeU="
   [mod."github.com/danieljoos/wincred"]
     version = "v1.2.0"
     hash = "sha256-LHcvTJCc8++bFndbd8ZgMSTe4L5h2C4rN+cSWHCz54Y="