feat(blockstm): cache pre-state to optimize value-based validation

skip cache oversized

Author: mmsqe

CHANGELOG.md

@@ -99,6 +99,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (blockstm) [25883](https://github.com/cosmos/cosmos-sdk/pull/25883) Re-use decoded tx object in pre-estimates.
 * (blockstm) [#25788](https://github.com/cosmos/cosmos-sdk/pull/25788) Only validate transactions that's executed at lease once.
 * (blockstm) [#25767](https://github.com/cosmos/cosmos-sdk/pull/25767) Optimize block-stm MVMemory with bitmap index.
+* (blockstm) [#25909](https://github.com/cosmos/cosmos-sdk/pull/25909) Cache pre-state to optimize value-based validation.
 
 ### Bug Fixes
 

internal/blockstm/mvmemory.go

@@ -1,6 +1,7 @@
 package blockstm
 
 import (
+	"sync"
 	"sync/atomic"
 
 	storetypes "cosmossdk.io/store/types"
@@ -17,9 +18,26 @@ type MVMemory struct {
 	scheduler   *Scheduler
 	stores      map[storetypes.StoreKey]int
 	data        []MVStore
+	preState    []preStateCache
 	lastReadSet []atomic.Pointer[MultiReadSet]
 }
 
+type preStateCache struct {
+	data sync.Map
+}
+
+func (c *preStateCache) Load(key Key) (any, bool) {
+	v, ok := c.data.Load(string(key))
+	if !ok {
+		return nil, false
+	}
+	return v, true
+}
+
+func (c *preStateCache) Store(key Key, value any) {
+	c.data.Store(string(key), value)
+}
+
 func NewMVMemory(
 	block_size int, stores map[storetypes.StoreKey]int,
 	storage MultiStore, scheduler *Scheduler,
@@ -41,6 +59,7 @@ func NewMVMemoryWithEstimates(
 		scheduler:   scheduler,
 		stores:      stores,
 		data:        data,
+		preState:    make([]preStateCache, len(stores)),
 		lastReadSet: make([]atomic.Pointer[MultiReadSet], block_size),
 	}
 
@@ -97,7 +116,7 @@ func (mv *MVMemory) View(txn TxnIndex) *MultiMVMemoryView {
 
 func (mv *MVMemory) newMVView(name storetypes.StoreKey, txn TxnIndex) MVView {
 	i := mv.stores[name]
-	return NewMVView(i, mv.storage.GetStore(name), mv.GetMVStore(i), mv.scheduler, txn)
+	return NewMVView(i, mv.storage.GetStore(name), mv.GetMVStore(i), mv.scheduler, txn, &mv.preState[i])
 }
 
 func (mv *MVMemory) GetMVStore(i int) MVStore {

internal/blockstm/mvview.go

@@ -11,6 +11,8 @@ import (
 	"github.com/cosmos/cosmos-sdk/telemetry"
 )
 
+const preStateMaxCacheValueBytes = 4096
+
 var (
 	_ storetypes.KVStore    = (*GMVMemoryView[[]byte])(nil)
 	_ storetypes.ObjKVStore = (*GMVMemoryView[any])(nil)
@@ -23,35 +25,57 @@ type GMVMemoryView[V any] struct {
 	storage   storetypes.GKVStore[V]
 	mvData    *GMVData[V]
 	scheduler *Scheduler
+	preState  *preStateCache
 	store     int
 
 	txn      TxnIndex
 	readSet  *ReadSet
 	writeSet *GMemDB[V]
 }
 
-func NewMVView(store int, storage storetypes.Store, mvData MVStore, scheduler *Scheduler, txn TxnIndex) MVView {
+func NewMVView(store int, storage storetypes.Store, mvData MVStore, scheduler *Scheduler, txn TxnIndex, preState *preStateCache) MVView {
 	switch data := mvData.(type) {
 	case *GMVData[any]:
-		return NewGMVMemoryView(store, storage.(storetypes.ObjKVStore), data, scheduler, txn)
+		return NewGMVMemoryView(store, storage.(storetypes.ObjKVStore), data, scheduler, txn, preState)
 	case *GMVData[[]byte]:
-		return NewGMVMemoryView(store, storage.(storetypes.KVStore), data, scheduler, txn)
+		return NewGMVMemoryView(store, storage.(storetypes.KVStore), data, scheduler, txn, preState)
 	default:
 		panic("unsupported value type")
 	}
 }
 
-func NewGMVMemoryView[V any](store int, storage storetypes.GKVStore[V], mvData *GMVData[V], scheduler *Scheduler, txn TxnIndex) *GMVMemoryView[V] {
+func NewGMVMemoryView[V any](store int, storage storetypes.GKVStore[V], mvData *GMVData[V], scheduler *Scheduler, txn TxnIndex, preState *preStateCache) *GMVMemoryView[V] {
 	return &GMVMemoryView[V]{
 		store:     store,
 		storage:   storage,
 		mvData:    mvData,
 		scheduler: scheduler,
+		preState:  preState,
 		txn:       txn,
 		readSet:   new(ReadSet),
 	}
 }
 
+func (s *GMVMemoryView[V]) getFromPreState(key []byte) V {
+	if s.preState != nil {
+		if cached, ok := s.preState.Load(key); ok {
+			return cached.(V)
+		}
+	}
+
+	result := s.storage.Get(key)
+	if s.preState != nil {
+		if bytesValue, ok := any(result).([]byte); ok {
+			if len(bytesValue) <= preStateMaxCacheValueBytes {
+				s.preState.Store(key, any(result))
+			}
+		} else {
+			s.preState.Store(key, any(result))
+		}
+	}
+	return result
+}
+
 func (s *GMVMemoryView[V]) init() {
 	if s.writeSet == nil {
 		s.writeSet = NewWriteSet(s.mvData.isZero, s.mvData.valueLen)
@@ -105,7 +129,7 @@ func (s *GMVMemoryView[V]) Get(key []byte) V {
 		// if not found, record version ⊥ when reading from storage.
 		s.readSet.Reads = append(s.readSet.Reads, ReadDescriptor{key, version})
 		if !version.Valid() {
-			result := s.storage.Get(key)
+			result := s.getFromPreState(key)
 			telemetry.MeasureSince(start, TelemetrySubsystem, KeyMVViewReadStorage) //nolint:staticcheck // TODO: switch to OpenTelemetry
 			return result
 		}