Refactor stochastic replay.

cmd/aida-sdb/trace/replay.go

@@ -53,11 +53,11 @@ type operationProcessor struct {
 }
 
 func (p operationProcessor) Process(state executor.State[[]operation.Operation], ctx *executor.Context) error {
-	p.runTransaction(uint64(state.Block), state.Data, ctx.State)
+	p.runTransaction(state.Block, state.Data, ctx.State)
 	return nil
 }
 
-func (p operationProcessor) runTransaction(block uint64, operations []operation.Operation, stateDb state.StateDB) {
+func (p operationProcessor) runTransaction(block int, operations []operation.Operation, stateDb state.StateDB) {
 	for _, op := range operations {
 		operation.Execute(op, stateDb, p.rCtx)
 		if p.cfg.Debug && block >= p.cfg.DebugFrom {

cmd/aida-sdb/trace/replay_substate.go

@@ -58,7 +58,7 @@ type substateProcessor struct {
 
 func (p substateProcessor) Process(state executor.State[*substate.Substate], ctx *executor.Context) error {
 	return p.operationProvider.Run(state.Block, state.Block, func(t executor.TransactionInfo[[]operation.Operation]) error {
-		p.runTransaction(uint64(state.Block), t.Data, ctx.State)
+		p.runTransaction(state.Block, t.Data, ctx.State)
 		return nil
 	})
 }

cmd/aida-stochastic-sdb/main.go

@@ -19,7 +19,7 @@ func initStochasticApp() *cli.App {
 			&stochastic.StochasticEstimateCommand,
 			&stochastic.StochasticGenerateCommand,
 			&stochastic.StochasticRecordCommand,
-			&stochastic.StochasticReplayCommand,
+			&StochasticReplayCommand,
 			&stochastic.StochasticVisualizeCommand,
 		},
 	}

cmd/aida-stochastic-sdb/replay.go

@@ -0,0 +1,152 @@
+package main
+
+import (
+	"fmt"
+	"math/rand"
+	"time"
+
+	"github.com/Fantom-foundation/Aida/executor"
+	"github.com/Fantom-foundation/Aida/executor/extension/aidadb"
+	"github.com/Fantom-foundation/Aida/executor/extension/primer"
+	"github.com/Fantom-foundation/Aida/executor/extension/profiler"
+	"github.com/Fantom-foundation/Aida/executor/extension/statedb"
+	"github.com/Fantom-foundation/Aida/executor/extension/tracker"
+	"github.com/Fantom-foundation/Aida/executor/extension/validator"
+	"github.com/Fantom-foundation/Aida/logger"
+	"github.com/Fantom-foundation/Aida/state"
+	"github.com/Fantom-foundation/Aida/stochastic"
+	"github.com/Fantom-foundation/Aida/utils"
+	"github.com/urfave/cli/v2"
+)
+
+// StochasticReplayCommand data structure for the replay app.
+var StochasticReplayCommand = cli.Command{
+	Action:    RunStochasticReplay,
+	Name:      "replay",
+	Usage:     "Simulates StateDB operations using a random generator with realistic distributions",
+	ArgsUsage: "<simulation-length> <simulation-file>",
+	Flags: []cli.Flag{
+		&utils.BalanceRangeFlag,
+		&utils.CarmenSchemaFlag,
+		&utils.ContinueOnFailureFlag,
+		&utils.CpuProfileFlag,
+		&utils.DebugFromFlag,
+		&utils.MemoryBreakdownFlag,
+		&utils.NonceRangeFlag,
+		&utils.RandomSeedFlag,
+		&utils.StateDbImplementationFlag,
+		&utils.StateDbVariantFlag,
+		&utils.DbTmpFlag,
+		&utils.StateDbLoggingFlag,
+		&utils.TraceFileFlag,
+		&utils.TraceDebugFlag,
+		&utils.TraceFlag,
+		&utils.ShadowDbImplementationFlag,
+		&utils.ShadowDbVariantFlag,
+		&logger.LogLevelFlag,
+	},
+	Description: `
+The stochastic replay command requires two argument:
+<simulation-length> <simulation.json> 
+
+<simulation-length> determines the number of blocks
+<simulation.json> contains the simulation parameters produced by the stochastic estimator.`,
+}
+
+func RunStochasticReplay(ctx *cli.Context) error {
+	cfg, err := utils.NewConfig(ctx, utils.BlockRangeArgs)
+	if err != nil {
+		return err
+	}
+
+	if cfg.StochasticSimulationFile == "" {
+		return fmt.Errorf("you must define path to simulation file (--%v)", utils.StochasticSimulationFileFlag.Name)
+	}
+
+	simulation, err := stochastic.ReadSimulation(cfg.StochasticSimulationFile)
+	if err != nil {
+		return fmt.Errorf("cannot read simulation; %v", err)
+	}
+
+	rg := rand.New(rand.NewSource(cfg.RandomSeed))
+
+	simulations, err := executor.OpenSimulations(simulation, ctx, rg)
+	if err != nil {
+		return err
+	}
+	defer simulations.Close()
+
+	return runStochasticReplay(cfg, simulations, nil, makeStochasticProcessor(cfg, simulation, rg), nil)
+
+}
+
+func makeStochasticProcessor(cfg *utils.Config, e *stochastic.EstimationModelJSON, rg *rand.Rand) executor.Processor[stochastic.Data] {
+	return stochasticProcessor{
+		stochastic.CreateState(e, rg, logger.NewLogger(cfg.LogLevel, "Stochastic Processor")), cfg,
+	}
+}
+
+type stochasticProcessor struct {
+	*stochastic.State
+	cfg *utils.Config
+}
+
+func (p stochasticProcessor) Process(state executor.State[stochastic.Data], ctx *executor.Context) error {
+	if p.cfg.Debug && state.Block >= p.cfg.DebugFrom {
+		p.EnableDebug()
+	}
+
+	p.Execute(state.Block, state.Transaction, state.Data, ctx.State)
+	return nil
+}
+
+func runStochasticReplay(
+	cfg *utils.Config,
+	provider executor.Provider[stochastic.Data],
+	stateDb state.StateDB,
+	processor executor.Processor[stochastic.Data],
+	extra []executor.Extension[stochastic.Data],
+) error {
+	// order of extensionList has to be maintained
+	var extensionList = []executor.Extension[stochastic.Data]{
+		profiler.MakeCpuProfiler[stochastic.Data](cfg),
+		profiler.MakeDiagnosticServer[stochastic.Data](cfg),
+	}
+
+	if stateDb == nil {
+		extensionList = append(
+			extensionList,
+			statedb.MakeStateDbManager[stochastic.Data](cfg),
+			tracker.MakeDbLogger[stochastic.Data](cfg),
+		)
+	}
+
+	extensionList = append(extensionList, extra...)
+
+	extensionList = append(extensionList, []executor.Extension[stochastic.Data]{
+		profiler.MakeThreadLocker[stochastic.Data](),
+		aidadb.MakeAidaDbManager[stochastic.Data](cfg),
+		profiler.MakeVirtualMachineStatisticsPrinter[stochastic.Data](cfg),
+		tracker.MakeProgressLogger[stochastic.Data](cfg, 15*time.Second),
+		tracker.MakeErrorLogger[stochastic.Data](cfg),
+		primer.MakeStateDbPrimer[stochastic.Data](cfg),
+		profiler.MakeMemoryUsagePrinter[stochastic.Data](cfg),
+		profiler.MakeMemoryProfiler[stochastic.Data](cfg),
+		validator.MakeStateHashValidator[stochastic.Data](cfg),
+
+		profiler.MakeOperationProfiler[stochastic.Data](cfg),
+	}...,
+	)
+
+	return executor.NewExecutor(provider, cfg.LogLevel).Run(
+		executor.Params{
+			From:                   int(cfg.First),
+			To:                     int(cfg.Last) + 1,
+			NumWorkers:             1, // stochastic can run only with one worker
+			State:                  stateDb,
+			ParallelismGranularity: executor.BlockLevel,
+		},
+		processor,
+		extensionList,
+	)
+}

cmd/aida-stochastic-sdb/replay_test.go

@@ -0,0 +1,185 @@
+package main
+
+import (
+	"errors"
+	"math/rand"
+	"strings"
+	"testing"
+
+	"github.com/Fantom-foundation/Aida/executor"
+	"github.com/Fantom-foundation/Aida/state"
+	"github.com/Fantom-foundation/Aida/stochastic"
+	"github.com/Fantom-foundation/Aida/stochastic/generator"
+	"github.com/Fantom-foundation/Aida/utils"
+	"github.com/ethereum/go-ethereum/common"
+	"go.uber.org/mock/gomock"
+)
+
+var simulation = &stochastic.EstimationModelJSON{
+	FileId:           "1",
+	Operations:       []string{},
+	StochasticMatrix: [][]float64{{1.0}, {2.0}},
+	Contracts: stochastic.EstimationStatsJSON{
+		NumKeys:           generator.MinRandomAccessSize,
+		Lambda:            1.1,
+		QueueDistribution: []float64{1.0, 2.0},
+	},
+	Keys: stochastic.EstimationStatsJSON{
+		NumKeys:           generator.MinRandomAccessSize,
+		Lambda:            1.1,
+		QueueDistribution: []float64{1.0, 2.0},
+	},
+	Values: stochastic.EstimationStatsJSON{
+		NumKeys:           generator.MinRandomAccessSize,
+		Lambda:            1.1,
+		QueueDistribution: []float64{1.0, 2.0},
+	},
+	SnapshotLambda: 1,
+}
+
+var rg = rand.New(rand.NewSource(1))
+
+func TestVmSdb_Substate_AllDbEventsAreIssuedInOrder(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	provider := executor.NewMockProvider[stochastic.Data](ctrl)
+	db := state.NewMockStateDB(ctrl)
+	cfg := &utils.Config{
+		First:             2,
+		Last:              4,
+		ChainID:           utils.MainnetChainID,
+		SkipPriming:       true,
+		ContinueOnFailure: true,
+		LogLevel:          "Critical",
+	}
+
+	// Simulate the execution of three transactions in two blocks.
+	provider.EXPECT().
+		Run(2, 5, gomock.Any()).
+		DoAndReturn(func(_ int, _ int, consumer executor.Consumer[stochastic.Data]) error {
+			// Block 2
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 2, Transaction: 1, Data: existData})
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 2, Transaction: 2, Data: beginTransactionData})
+			// Block 3
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 3, Transaction: 1, Data: beginBlockData})
+			// Block 4
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 4, Transaction: utils.PseudoTx, Data: addBalanceData})
+			return nil
+		})
+
+	// The expectation is that all of those blocks and transactions
+	// are properly opened, prepared, executed, and closed.
+	gomock.InOrder(
+		db.EXPECT().Exist(common.Address{byte(0)}),
+		db.EXPECT().BeginTransaction(uint32(2)),
+		db.EXPECT().BeginBlock(uint64(3)),
+		db.EXPECT().AddBalance(common.Address{byte(0)}, executor.WithBigIntOfAnySize()),
+	)
+
+	// since we are working with mock transactions, run logically fails on 'intrinsic gas too low'
+	// since this is a test that tests orded of the db events, we can ignore this error
+	err := runStochasticReplay(cfg, provider, db, makeStochasticProcessor(cfg, simulation, rg), nil)
+	if err != nil {
+		errors.Unwrap(err)
+		if strings.Contains(err.Error(), "intrinsic gas too low") {
+			return
+		}
+		t.Fatal("run failed")
+	}
+}
+
+func TestVmSdb_Substate_AllTransactionsAreProcessedInOrder(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	provider := executor.NewMockProvider[stochastic.Data](ctrl)
+	db := state.NewMockStateDB(ctrl)
+	ext := executor.NewMockExtension[stochastic.Data](ctrl)
+	processor := executor.NewMockProcessor[stochastic.Data](ctrl)
+	cfg := &utils.Config{
+		First:       2,
+		Last:        4,
+		ChainID:     utils.MainnetChainID,
+		LogLevel:    "Critical",
+		SkipPriming: true,
+	}
+
+	// Simulate the execution of three transactions in two blocks.
+	provider.EXPECT().
+		Run(2, 5, gomock.Any()).
+		DoAndReturn(func(_ int, _ int, consumer executor.Consumer[stochastic.Data]) error {
+			// Block 2
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 2, Transaction: 1, Data: stochastic.Data{}})
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 2, Transaction: 2, Data: stochastic.Data{}})
+			// Block 3
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 3, Transaction: 1, Data: stochastic.Data{}})
+			// Block 4
+			consumer(executor.TransactionInfo[stochastic.Data]{Block: 4, Transaction: utils.PseudoTx, Data: stochastic.Data{}})
+			return nil
+		})
+
+	// The expectation is that all of those blocks and transactions
+	// are properly opened, prepared, executed, and closed.
+	// Since we are running sequential mode with 1 worker,
+	// all block and transactions need to be in order.
+	gomock.InOrder(
+		ext.EXPECT().PreRun(executor.AtBlock[stochastic.Data](2), gomock.Any()),
+
+		// Block 2
+		// Tx 1
+		ext.EXPECT().PreBlock(executor.AtBlock[stochastic.Data](2), gomock.Any()),
+		ext.EXPECT().PreTransaction(executor.AtTransaction[stochastic.Data](2, 1), gomock.Any()),
+		processor.EXPECT().Process(executor.AtTransaction[stochastic.Data](2, 1), gomock.Any()),
+		ext.EXPECT().PostTransaction(executor.AtTransaction[stochastic.Data](2, 1), gomock.Any()),
+		ext.EXPECT().PreTransaction(executor.AtTransaction[stochastic.Data](2, 2), gomock.Any()),
+		// Tx 2
+		processor.EXPECT().Process(executor.AtTransaction[stochastic.Data](2, 2), gomock.Any()),
+		ext.EXPECT().PostTransaction(executor.AtTransaction[stochastic.Data](2, 2), gomock.Any()),
+		ext.EXPECT().PostBlock(executor.AtTransaction[stochastic.Data](2, 2), gomock.Any()),
+
+		// Block 3
+		ext.EXPECT().PreBlock(executor.AtBlock[stochastic.Data](3), gomock.Any()),
+		ext.EXPECT().PreTransaction(executor.AtTransaction[stochastic.Data](3, 1), gomock.Any()),
+		processor.EXPECT().Process(executor.AtTransaction[stochastic.Data](3, 1), gomock.Any()),
+		ext.EXPECT().PostTransaction(executor.AtTransaction[stochastic.Data](3, 1), gomock.Any()),
+		ext.EXPECT().PostBlock(executor.AtTransaction[stochastic.Data](3, 1), gomock.Any()),
+
+		// Block 4
+		ext.EXPECT().PreBlock(executor.AtBlock[stochastic.Data](4), gomock.Any()),
+		ext.EXPECT().PreTransaction(executor.AtTransaction[stochastic.Data](4, utils.PseudoTx), gomock.Any()),
+		processor.EXPECT().Process(executor.AtTransaction[stochastic.Data](4, utils.PseudoTx), gomock.Any()),
+		ext.EXPECT().PostTransaction(executor.AtTransaction[stochastic.Data](4, utils.PseudoTx), gomock.Any()),
+		ext.EXPECT().PostBlock(executor.AtTransaction[stochastic.Data](4, utils.PseudoTx), gomock.Any()),
+
+		ext.EXPECT().PostRun(executor.AtBlock[stochastic.Data](5), gomock.Any(), nil),
+	)
+
+	if err := runStochasticReplay(cfg, provider, db, processor, []executor.Extension[stochastic.Data]{ext}); err != nil {
+		t.Errorf("run failed: %v", err)
+	}
+}
+
+var beginBlockData = stochastic.Data{
+	Operation: stochastic.BeginBlockID,
+	Address:   0,
+	Key:       0,
+	Value:     0,
+}
+
+var beginTransactionData = stochastic.Data{
+	Operation: stochastic.BeginTransactionID,
+	Address:   0,
+	Key:       0,
+	Value:     0,
+}
+
+var existData = stochastic.Data{
+	Operation: stochastic.ExistID,
+	Address:   0,
+	Key:       0,
+	Value:     0,
+}
+
+var addBalanceData = stochastic.Data{
+	Operation: stochastic.AddBalanceID,
+	Address:   0,
+	Key:       0,
+	Value:     1,
+}