Extend DB benchmark to support state store (SS) backends

sei-db/state_db/bench/bench_ss_test.go

@@ -0,0 +1,45 @@
+package bench
+
+import (
+	"testing"
+
+	"github.com/sei-protocol/sei-chain/sei-db/state_db/bench/wrappers"
+)
+
+func BenchmarkSSCompositeWrite(b *testing.B) {
+	const totalKeys int64 = 10_000
+
+	scenarios := []TestScenario{
+		{
+			Name:      "100_keys_per_block",
+			TotalKeys: totalKeys,
+			NumBlocks: totalKeys / 100,
+			Backend:   wrappers.SSComposite,
+		},
+	}
+
+	for _, scenario := range scenarios {
+		b.Run(scenario.Name, func(b *testing.B) {
+			runBenchmark(b, scenario, false)
+		})
+	}
+}
+
+func BenchmarkCombinedCompositeDualSSCompositeWrite(b *testing.B) {
+	const totalKeys int64 = 10_000
+
+	scenarios := []TestScenario{
+		{
+			Name:      "100_keys_per_block",
+			TotalKeys: totalKeys,
+			NumBlocks: totalKeys / 100,
+			Backend:   wrappers.CompositeDual_SSComposite,
+		},
+	}
+
+	for _, scenario := range scenarios {
+		b.Run(scenario.Name, func(b *testing.B) {
+			runBenchmark(b, scenario, false)
+		})
+	}
+}

sei-db/state_db/bench/cryptosim/config/ss-composite-config.json

@@ -0,0 +1,5 @@
+{
+  "Comment": "State store benchmark using composite SS with EVM sub-stores (dual write + EVM-first read).",
+  "Backend": "SSComposite",
+  "DataDir": "data"
+}

sei-db/state_db/bench/wrappers/combined_wrapper.go

@@ -0,0 +1,66 @@
+package wrappers
+
+import (
+	"sync/atomic"
+
+	dbTypes "github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	scTypes "github.com/sei-protocol/sei-chain/sei-db/state_db/sc/types"
+)
+
+var _ DBWrapper = (*combinedWrapper)(nil)
+
+// combinedWrapper drives both a State Commit (SC) and State Store (SS) backend
+// from the same changeset stream, mirroring production where SC and SS receive
+// identical writes.
+type combinedWrapper struct {
+	sc        DBWrapper
+	ss        dbTypes.StateStore
+	ssVersion atomic.Int64
+}
+
+func NewCombinedWrapper(sc DBWrapper, ss dbTypes.StateStore) DBWrapper {
+	w := &combinedWrapper{sc: sc, ss: ss}
+	w.ssVersion.Store(ss.GetLatestVersion())
+	return w
+}
+
+func (c *combinedWrapper) ApplyChangeSets(cs []*proto.NamedChangeSet) error {
+	if err := c.sc.ApplyChangeSets(cs); err != nil {
+		return err
+	}
+	nextVersion := c.ssVersion.Add(1)
+	return c.ss.ApplyChangesetSync(nextVersion, cs)
+}
+
+func (c *combinedWrapper) Read(key []byte) (data []byte, found bool, err error) {
+	return c.sc.Read(key)
+}
+
+func (c *combinedWrapper) Commit() (int64, error) {
+	if _, err := c.sc.Commit(); err != nil {
+		return 0, err
+	}
+	return c.ssVersion.Load(), nil
+}
+
+func (c *combinedWrapper) Close() error {
+	scErr := c.sc.Close()
+	ssErr := c.ss.Close()
+	if scErr != nil {
+		return scErr
+	}
+	return ssErr
+}
+
+func (c *combinedWrapper) Version() int64 {
+	return c.ssVersion.Load()
+}
+
+func (c *combinedWrapper) LoadVersion(version int64) error {
+	return c.sc.LoadVersion(version)
+}
+
+func (c *combinedWrapper) Importer(version int64) (scTypes.Importer, error) {
+	return c.sc.Importer(version)
+}

sei-db/state_db/bench/wrappers/composite_wrapper.go

@@ -34,7 +34,7 @@ func (c *compositeWrapper) LoadVersion(version int64) error {
 }
 
 func (c *compositeWrapper) Version() int64 {
-	return c.base.WorkingCommitInfo().Version
+	return c.base.Version()
 }
 
 func (c *compositeWrapper) Importer(version int64) (types.Importer, error) {

sei-db/state_db/bench/wrappers/db_implementations.go

@@ -2,12 +2,14 @@ package wrappers
 
 import (
 	"fmt"
+	"path/filepath"
 
 	"github.com/sei-protocol/sei-chain/sei-db/common/logger"
 	"github.com/sei-protocol/sei-chain/sei-db/config"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/composite"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/flatkv"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/memiavl"
+	ssComposite "github.com/sei-protocol/sei-chain/sei-db/state_db/ss/composite"
 )
 
 const EVMStoreName = "evm"
@@ -20,6 +22,9 @@ const (
 	CompositeDual   DBType = "CompositeDual"
 	CompositeSplit  DBType = "CompositeSplit"
 	CompositeCosmos DBType = "CompositeCosmos"
+
+	SSComposite               DBType = "SSComposite"
+	CompositeDual_SSComposite DBType = "CompositeDual+SSComposite"
 )
 
 func newMemIAVLCommitStore(dbDir string) (DBWrapper, error) {
@@ -76,6 +81,38 @@ func newCompositeCommitStore(dbDir string, writeMode config.WriteMode) (DBWrappe
 	return NewCompositeWrapper(loadedStore), nil
 }
 
+func openSSComposite(dir string) (*ssComposite.CompositeStateStore, error) {
+	cfg := config.DefaultStateStoreConfig()
+	cfg.Backend = config.PebbleDBBackend
+	cfg.AsyncWriteBuffer = 0
+	cfg.WriteMode = config.DualWrite
+	cfg.ReadMode = config.EVMFirstRead
+	return ssComposite.NewCompositeStateStore(cfg, dir, logger.NewNopLogger())
+}
+
+func newSSCompositeStateStore(dbDir string) (DBWrapper, error) {
+	fmt.Printf("Opening composite state store from directory %s\n", dbDir)
+	store, err := openSSComposite(dbDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open composite state store: %w", err)
+	}
+	return NewStateStoreWrapper(store), nil
+}
+
+func newCombinedCompositeDualSSComposite(dbDir string) (DBWrapper, error) {
+	fmt.Printf("Opening CompositeDual (SC) + Composite (SS) from directory %s\n", dbDir)
+	sc, err := newCompositeCommitStore(filepath.Join(dbDir, "sc"), config.DualWrite)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create SC store: %w", err)
+	}
+	ss, err := openSSComposite(filepath.Join(dbDir, "ss"))
+	if err != nil {
+		_ = sc.Close()
+		return nil, fmt.Errorf("failed to create SS store: %w", err)
+	}
+	return NewCombinedWrapper(sc, ss), nil
+}
+
 // NewDBImpl instantiates a new empty DBWrapper based on the given DBType.
 func NewDBImpl(dbType DBType, dataDir string) (DBWrapper, error) {
 	switch dbType {
@@ -89,6 +126,10 @@ func NewDBImpl(dbType DBType, dataDir string) (DBWrapper, error) {
 		return newCompositeCommitStore(dataDir, config.SplitWrite)
 	case CompositeCosmos:
 		return newCompositeCommitStore(dataDir, config.CosmosOnlyWrite)
+	case SSComposite:
+		return newSSCompositeStateStore(dataDir)
+	case CompositeDual_SSComposite:
+		return newCombinedCompositeDualSSComposite(dataDir)
 	default:
 		return nil, fmt.Errorf("unsupported DB type: %s", dbType)
 	}

sei-db/state_db/bench/wrappers/state_store_wrapper.go

@@ -0,0 +1,66 @@
+package wrappers
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	dbTypes "github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
+	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	scTypes "github.com/sei-protocol/sei-chain/sei-db/state_db/sc/types"
+)
+
+var _ DBWrapper = (*stateStoreWrapper)(nil)
+
+// stateStoreWrapper adapts a versioned StateStore (SS layer) to the DBWrapper
+// interface used by the cryptosim benchmark. Each ApplyChangeSets call maps to
+// a single ApplyChangesetSync at an incrementing version. The SS layer persists
+// on every apply, so Commit is a no-op.
+type stateStoreWrapper struct {
+	base    dbTypes.StateStore
+	version atomic.Int64
+}
+
+func NewStateStoreWrapper(store dbTypes.StateStore) DBWrapper {
+	w := &stateStoreWrapper{
+		base: store,
+	}
+	w.version.Store(store.GetLatestVersion())
+	return w
+}
+
+func (s *stateStoreWrapper) ApplyChangeSets(cs []*proto.NamedChangeSet) error {
+	nextVersion := s.version.Add(1)
+	return s.base.ApplyChangesetSync(nextVersion, cs)
+}
+
+func (s *stateStoreWrapper) Read(key []byte) (data []byte, found bool, err error) {
+	version := s.version.Load()
+	if version == 0 {
+		return nil, false, nil
+	}
+	val, err := s.base.Get(EVMStoreName, version, key)
+	if err != nil {
+		return nil, false, err
+	}
+	return val, val != nil, nil
+}
+
+func (s *stateStoreWrapper) Commit() (int64, error) {
+	return s.version.Load(), nil
+}
+
+func (s *stateStoreWrapper) Close() error {
+	return s.base.Close()
+}
+
+func (s *stateStoreWrapper) Version() int64 {
+	return s.version.Load()
+}
+
+func (s *stateStoreWrapper) LoadVersion(_ int64) error {
+	return nil
+}
+
+func (s *stateStoreWrapper) Importer(_ int64) (scTypes.Importer, error) {
+	return nil, fmt.Errorf("import not supported for state store wrapper")
+}