feat: allow specifying test data dirs (#132)

Author: meyer9

runner/benchmark/definition.go

@@ -91,9 +91,15 @@ type BenchmarkConfig struct {
 	TransactionPayloads []payload.Definition `yaml:"payloads"`
 }
 
+type DatadirConfig struct {
+	Sequencer *string `yaml:"sequencer"`
+	Validator *string `yaml:"validator"`
+}
+
 // TestDefinition is the user-facing YAML configuration for specifying a
 // matrix of benchmark runs.
 type TestDefinition struct {
+	Datadir      *DatadirConfig       `yaml:"datadirs"`
 	Snapshot     *SnapshotDefinition  `yaml:"snapshot"`
 	Metrics      *ThresholdConfig     `yaml:"metrics"`
 	Tags         *map[string]string   `yaml:"tags"`

runner/benchmark/matrix.go

@@ -13,6 +13,7 @@ type ThresholdConfig struct {
 // TestPlan represents a list of test runs to be executed.
 type TestPlan struct {
 	Runs         []TestRun
+	Datadir      *DatadirConfig
 	Snapshot     *SnapshotDefinition
 	ProofProgram *ProofProgramOptions
 	Thresholds   *ThresholdConfig
@@ -37,6 +38,7 @@ func NewTestPlanFromConfig(c TestDefinition, testFileName string, config *Benchm
 
 	return &TestPlan{
 		Runs:         testRuns,
+		Datadir:      c.Datadir,
 		Snapshot:     c.Snapshot,
 		ProofProgram: proofProgram,
 		Thresholds:   c.Metrics,

runner/benchmark/snapshots.go

@@ -12,7 +12,7 @@ type SnapshotManager interface {
 	// EnsureSnapshot ensures that a snapshot exists for the given node type and
 	// role. If it does not exist, it will create it using the given snapshot
 	// definition. It returns the path to the snapshot.
-	EnsureSnapshot(definition SnapshotDefinition, nodeType string, role string) (string, error)
+	EnsureSnapshot(datadirsConfig *DatadirConfig, definition SnapshotDefinition, nodeType string, role string) (string, error)
 }
 
 type snapshotStoragePath struct {
@@ -58,7 +58,7 @@ func NewSnapshotManager(snapshotsDir string) SnapshotManager {
 	}
 }
 
-func (b *benchmarkDatadirState) EnsureSnapshot(definition SnapshotDefinition, nodeType string, role string) (string, error) {
+func (b *benchmarkDatadirState) EnsureSnapshot(datadirsConfig *DatadirConfig, definition SnapshotDefinition, nodeType string, role string) (string, error) {
 	snapshotDatadir := snapshotStoragePath{
 		nodeType: nodeType,
 		role:     role,
@@ -69,10 +69,15 @@ func (b *benchmarkDatadirState) EnsureSnapshot(definition SnapshotDefinition, no
 		return datadir, nil
 	}
 
-	hashCommand := sha256.New().Sum([]byte(definition.Command))
-
-	snapshotPath := filepath.Join(b.snapshotsDir, fmt.Sprintf("%s_%s_%x", nodeType, role, hashCommand[:12]))
-
+	var snapshotPath string
+	if datadirsConfig != nil && role == "validator" && datadirsConfig.Validator != nil {
+		snapshotPath = *datadirsConfig.Validator
+	} else if datadirsConfig != nil && role == "sequencer" && datadirsConfig.Sequencer != nil {
+		snapshotPath = *datadirsConfig.Sequencer
+	} else {
+		hashCommand := sha256.New().Sum([]byte(definition.Command))
+		snapshotPath = filepath.Join(b.snapshotsDir, fmt.Sprintf("%s_%s_%x", nodeType, role, hashCommand[:12]))
+	}
 	// Create a new datadir for this snapshot.
 	err := definition.CreateSnapshot(nodeType, snapshotPath)
 	if err != nil {

runner/service.go

@@ -73,7 +73,7 @@ func readBenchmarkConfig(path string) (*benchmark.BenchmarkConfig, error) {
 	return config, err
 }
 
-func (s *service) setupInternalDirectories(testDir string, params types.RunParams, genesis *core.Genesis, snapshot *benchmark.SnapshotDefinition, role string) (*config.InternalClientOptions, error) {
+func (s *service) setupInternalDirectories(testDir string, params types.RunParams, genesis *core.Genesis, snapshot *benchmark.SnapshotDefinition, role string, datadirsConfig *benchmark.DatadirConfig) (*config.InternalClientOptions, error) {
 	err := os.MkdirAll(testDir, 0755)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create test directory")
@@ -102,7 +102,7 @@ func (s *service) setupInternalDirectories(testDir string, params types.RunParam
 	isSnapshot := snapshot != nil && snapshot.Command != ""
 	if isSnapshot {
 		// if we have a snapshot, restore it if needed or reuse from a previous test
-		snapshotDir, err := s.dataDirState.EnsureSnapshot(*snapshot, params.NodeType, role)
+		snapshotDir, err := s.dataDirState.EnsureSnapshot(datadirsConfig, *snapshot, params.NodeType, role)
 		if err != nil {
 			return nil, errors.Wrap(err, "failed to ensure snapshot")
 		}
@@ -291,18 +291,18 @@ func (s *service) getGenesisForSnapshotConfig(snapshotConfig *benchmark.Snapshot
 	return genesis, nil
 }
 
-func (s *service) setupDataDirs(workingDir string, params types.RunParams, genesis *core.Genesis, snapshot *benchmark.SnapshotDefinition) (*config.InternalClientOptions, *config.InternalClientOptions, error) {
+func (s *service) setupDataDirs(workingDir string, params types.RunParams, genesis *core.Genesis, snapshot *benchmark.SnapshotDefinition, datadirsConfig *benchmark.DatadirConfig) (*config.InternalClientOptions, *config.InternalClientOptions, error) {
 	// create temp directory for this test
 	testName := fmt.Sprintf("%d-%s-test", time.Now().Unix(), params.NodeType)
 	sequencerTestDir := path.Join(workingDir, fmt.Sprintf("%s-sequencer", testName))
 	validatorTestDir := path.Join(workingDir, fmt.Sprintf("%s-validator", testName))
 
-	sequencerOptions, err := s.setupInternalDirectories(sequencerTestDir, params, genesis, snapshot, "sequencer")
+	sequencerOptions, err := s.setupInternalDirectories(sequencerTestDir, params, genesis, snapshot, "sequencer", datadirsConfig)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "failed to setup internal directories")
 	}
 
-	validatorOptions, err := s.setupInternalDirectories(validatorTestDir, params, genesis, snapshot, "validator")
+	validatorOptions, err := s.setupInternalDirectories(validatorTestDir, params, genesis, snapshot, "validator", datadirsConfig)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "failed to setup internal directories")
 	}
@@ -320,7 +320,7 @@ func (s *service) setupBlobsDir(workingDir string) error {
 	return nil
 }
 
-func (s *service) runTest(ctx context.Context, params types.RunParams, workingDir string, outputDir string, snapshotConfig *benchmark.SnapshotDefinition, proofConfig *benchmark.ProofProgramOptions, transactionPayload payload.Definition) (*benchmark.RunResult, error) {
+func (s *service) runTest(ctx context.Context, params types.RunParams, workingDir string, outputDir string, snapshotConfig *benchmark.SnapshotDefinition, proofConfig *benchmark.ProofProgramOptions, transactionPayload payload.Definition, datadirsConfig *benchmark.DatadirConfig) (*benchmark.RunResult, error) {
 
 	s.log.Info(fmt.Sprintf("Running benchmark with params: %+v", params))
 
@@ -336,7 +336,7 @@ func (s *service) runTest(ctx context.Context, params types.RunParams, workingDi
 	validatorTestDir := path.Join(workingDir, fmt.Sprintf("%s-validator", testName))
 
 	// setup data directories (restore from snapshot if needed)
-	sequencerOptions, validatorOptions, err := s.setupDataDirs(workingDir, params, genesis, snapshotConfig)
+	sequencerOptions, validatorOptions, err := s.setupDataDirs(workingDir, params, genesis, snapshotConfig, datadirsConfig)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to setup data dirs")
 	}
@@ -566,7 +566,7 @@ outerLoop:
 				return errors.Wrap(err, "failed to create output directory")
 			}
 
-			metricSummary, err := s.runTest(ctx, c.Params, s.config.DataDir(), outputDir, testPlan.Snapshot, testPlan.ProofProgram, transactionPayloads[c.Params.PayloadID])
+			metricSummary, err := s.runTest(ctx, c.Params, s.config.DataDir(), outputDir, testPlan.Snapshot, testPlan.ProofProgram, transactionPayloads[c.Params.PayloadID], testPlan.Datadir)
 			if err != nil {
 				log.Error("Failed to run test", "err", err)
 				metricSummary = &benchmark.RunResult{