cleanup

Author: mmsqe

blockstm/bench_test.go

@@ -2,32 +2,16 @@ package blockstm
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"runtime"
 	"strconv"
 	"sync/atomic"
 	"testing"
 
 	"github.com/test-go/testify/require"
-	"golang.org/x/sync/errgroup"
 
 	storetypes "cosmossdk.io/store/types"
 )
 
-func executeBlock(stores map[storetypes.StoreKey]int, storage MultiStore, worker int, block *MockBlock) error {
-	incarnationCache := make([]atomic.Pointer[map[string]any], block.Size())
-	for i := 0; i < block.Size(); i++ {
-		m := make(map[string]any)
-		incarnationCache[i].Store(&m)
-	}
-	return ExecuteBlock(context.Background(), block.Size(), stores, storage, worker, func(txn TxnIndex, store MultiStore) {
-		cache := incarnationCache[txn].Swap(nil)
-		block.ExecuteTx(txn, store, *cache)
-		incarnationCache[txn].Store(cache)
-	})
-}
-
 func BenchmarkBlockSTM(b *testing.B) {
 	stores := map[storetypes.StoreKey]int{StoreKeyAuth: 0, StoreKeyBank: 1}
 	for i := 0; i < 26; i++ {
@@ -92,17 +76,19 @@ func BenchmarkBlockSTM(b *testing.B) {
 				b.ResetTimer()
 				var executedTotal, validatedTotal uint64
 				for i := 0; i < b.N; i++ {
-					executed, validated, err := executeBlockForBench(
+					executed, validated, err := executeBlockWithEstimatesImpl(
 						context.Background(),
 						tc.block.Size(),
 						stores,
 						storage,
 						worker,
+						nil,
 						func(txn TxnIndex, store MultiStore) {
 							cache := incarnationCache[txn].Swap(nil)
 							tc.block.ExecuteTx(txn, store, *cache)
 							incarnationCache[txn].Store(cache)
 						},
+						false,
 					)
 					require.NoError(b, err)
 					executedTotal += executed
@@ -148,47 +134,6 @@ func iterateNewKeysBlock(size int) *MockBlock {
 	return NewMockBlock(txs)
 }
 
-// executeBlockForBench is a benchmark-only variant of ExecuteBlockWithEstimates that
-// exposes scheduler counters (executed / validated) without emitting telemetry.
-func executeBlockForBench(
-	ctx context.Context,
-	blockSize int,
-	stores map[storetypes.StoreKey]int,
-	storage MultiStore,
-	executors int,
-	txExecutor TxExecutor,
-) (executed, validated uint64, err error) {
-	if executors < 0 {
-		return 0, 0, fmt.Errorf("invalid number of executors: %d", executors)
-	}
-	if executors == 0 {
-		executors = min(runtime.GOMAXPROCS(0), runtime.NumCPU())
-	}
-
-	scheduler := NewScheduler(blockSize)
-	mvMemory := NewMVMemory(blockSize, stores, storage, scheduler)
-
-	var wg errgroup.Group
-	wg.SetLimit(executors)
-	for i := 0; i < executors; i++ {
-		e := NewExecutor(ctx, scheduler, txExecutor, mvMemory, i)
-		wg.Go(e.Run)
-	}
-	if err := wg.Wait(); err != nil {
-		return 0, 0, err
-	}
-
-	if !scheduler.Done() {
-		if ctx.Err() != nil {
-			return 0, 0, ctx.Err()
-		}
-		return 0, 0, errors.New("scheduler did not complete")
-	}
-
-	mvMemory.WriteSnapshot(storage)
-	return uint64(scheduler.executedTxns.Load()), uint64(scheduler.validatedTxns.Load()), nil
-}
-
 func runSequential(storage MultiStore, block *MockBlock) {
 	for i, tx := range block.Txs {
 		block.Results[i] = tx(storage, nil)

blockstm/stm.go

@@ -36,8 +36,31 @@ func ExecuteBlockWithEstimates(
 	estimates []MultiLocations, // txn -> multi-locations
 	txExecutor TxExecutor,
 ) error {
+	_, _, err := executeBlockWithEstimatesImpl(
+		ctx,
+		blockSize,
+		stores,
+		storage,
+		executors,
+		estimates,
+		txExecutor,
+		true,
+	)
+	return err
+}
+
+func executeBlockWithEstimatesImpl(
+	ctx context.Context,
+	blockSize int,
+	stores map[storetypes.StoreKey]int,
+	storage MultiStore,
+	executors int,
+	estimates []MultiLocations, // txn -> multi-locations
+	txExecutor TxExecutor,
+	emitTelemetry bool,
+) (executed, validated uint64, err error) {
 	if executors < 0 {
-		return fmt.Errorf("invalid number of executors: %d", executors)
+		return 0, 0, fmt.Errorf("invalid number of executors: %d", executors)
 	}
 	if executors == 0 {
 		executors = maxParallelism()
@@ -68,25 +91,30 @@ func ExecuteBlockWithEstimates(
 		}
 	}()
 
-	err := wg.Wait()
+	err = wg.Wait()
 	close(cancelDone)
 	if err != nil {
-		return err
+		return 0, 0, err
 	}
 
 	if !scheduler.Done() {
 		if ctx.Err() != nil {
-			return ctx.Err()
+			return 0, 0, ctx.Err()
 		}
-		return errors.New("scheduler did not complete")
+		return 0, 0, errors.New("scheduler did not complete")
 	}
 
-	telemetry.SetGauge(float32(scheduler.executedTxns.Load()), TelemetrySubsystem, KeyExecutedTxs)   //nolint:staticcheck // TODO: switch to OpenTelemetry
-	telemetry.SetGauge(float32(scheduler.validatedTxns.Load()), TelemetrySubsystem, KeyValidatedTxs) //nolint:staticcheck // TODO: switch to OpenTelemetry
+	executed = uint64(scheduler.executedTxns.Load())
+	validated = uint64(scheduler.validatedTxns.Load())
+
+	if emitTelemetry {
+		telemetry.SetGauge(float32(executed), TelemetrySubsystem, KeyExecutedTxs)   //nolint:staticcheck // TODO: switch to OpenTelemetry
+		telemetry.SetGauge(float32(validated), TelemetrySubsystem, KeyValidatedTxs) //nolint:staticcheck // TODO: switch to OpenTelemetry
+	}
 
 	// Write the snapshot into the storage
 	mvMemory.WriteSnapshot(storage)
-	return nil
+	return executed, validated, nil
 }
 
 func maxParallelism() int {