# This is a combination of 5 commits.

# This is the 1st commit message:

Enable locking kv store

# The commit message #2 will be skipped:

# Fix some lock orderings

# The commit message #3 will be skipped:

# Fix minor typo

# The commit message #4 will be skipped:

# Ensure that writes happen in a deterministic order.
#
# Ensure that reads are also done all the time, remove this if it doesn't impact gas.

# The commit message #5 will be skipped:

# Remove locking for now for lockingkv.Get/Has

baseapp/baseapp.go

@@ -19,6 +19,7 @@ import (
 	protov2 "google.golang.org/protobuf/proto"
 
 	"cosmossdk.io/core/header"
+	"cosmossdk.io/store/rootmulti"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
 	"cosmossdk.io/store"
@@ -191,6 +192,7 @@ type BaseApp struct {
 
 	// Used to synchronize the application when using an unsynchronized ABCI++ client.
 	mtx sync.RWMutex
+
 	// Used to synchronize CacheMultistoreWithVersion since the multistore mutates version
 	// information internally during first time loads leading to data races.
 	cacheMsWithVersionMtx sync.Mutex
@@ -478,7 +480,12 @@ func (app *BaseApp) IsSealed() bool { return app.sealed }
 // multi-store (i.e. a CacheMultiStore) and a new Context with the same
 // multi-store branch, and provided header.
 func (app *BaseApp) setState(mode execMode, h cmtproto.Header) {
-	ms := app.cms.CacheMultiStore()
+	var ms storetypes.CacheMultiStore
+	if mode == execModeCheck {
+		ms = app.cms.(*rootmulti.Store).LockingCacheMultiStore()
+	} else {
+		ms = app.cms.CacheMultiStore()
+	}
 	headerInfo := header.Info{
 		Height:  h.Height,
 		Time:    h.Time,

go.mod

@@ -199,3 +199,5 @@ retract (
 )
 
 replace github.com/cometbft/cometbft => github.com/dydxprotocol/cometbft v0.38.6-0.20240220185844-e704122c8540
+
+replace cosmossdk.io/store => github.com/dydxprotocol/cosmos-sdk/store v1.0.3-0.20240226162728-392d75d8a6da

go.sum

@@ -14,8 +14,6 @@ cosmossdk.io/log v1.3.1 h1:UZx8nWIkfbbNEWusZqzAx3ZGvu54TZacWib3EzUYmGI=
 cosmossdk.io/log v1.3.1/go.mod h1:2/dIomt8mKdk6vl3OWJcPk2be3pGOS8OQaLUM/3/tCM=
 cosmossdk.io/math v1.2.0 h1:8gudhTkkD3NxOP2YyyJIYYmt6dQ55ZfJkDOaxXpy7Ig=
 cosmossdk.io/math v1.2.0/go.mod h1:l2Gnda87F0su8a/7FEKJfFdJrM0JZRXQaohlgJeyQh0=
-cosmossdk.io/store v1.0.2 h1:lSg5BTvJBHUDwswNNyeh4K/CbqiHER73VU4nDNb8uk0=
-cosmossdk.io/store v1.0.2/go.mod h1:EFtENTqVTuWwitGW1VwaBct+yDagk7oG/axBMPH+FXs=
 cosmossdk.io/x/tx v0.13.0 h1:8lzyOh3zONPpZv2uTcUmsv0WTXy6T1/aCVDCqShmpzU=
 cosmossdk.io/x/tx v0.13.0/go.mod h1:CpNQtmoqbXa33/DVxWQNx5Dcnbkv2xGUhL7tYQ5wUsY=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
@@ -201,6 +199,8 @@ github.com/dvsekhvalnov/jose2go v1.6.0 h1:Y9gnSnP4qEI0+/uQkHvFXeD2PLPJeXEL+ySMEA
 github.com/dvsekhvalnov/jose2go v1.6.0/go.mod h1:QsHjhyTlD/lAVqn/NSbVZmSCGeDehTB/mPZadG+mhXU=
 github.com/dydxprotocol/cometbft v0.38.6-0.20240220185844-e704122c8540 h1:pkYQbAdOAAoZBSId9kLupCgZHj8YvA9LzM31fVYpjlw=
 github.com/dydxprotocol/cometbft v0.38.6-0.20240220185844-e704122c8540/go.mod h1:REQN+ObgfYxi39TcYR/Hv95C9bPxY3sYJCvghryj7vY=
+github.com/dydxprotocol/cosmos-sdk/store v1.0.3-0.20240226162728-392d75d8a6da h1:Ts59JvgcwDLjLh4a/8YNAI+i9LYfyebFuiIpRsOJcMY=
+github.com/dydxprotocol/cosmos-sdk/store v1.0.3-0.20240226162728-392d75d8a6da/go.mod h1:M4hH9rTCui49tnujKeM+BWCiXGQwZ//3QCvl0HbR6Cs=
 github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
 github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=
 github.com/eapache/queue v1.1.0/go.mod h1:6eCeP0CKFpHLu8blIFXhExK/dRa7WDZfr6jVFPTqq+I=

store/cachemulti/store.go

@@ -8,6 +8,7 @@ import (
 
 	"cosmossdk.io/store/cachekv"
 	"cosmossdk.io/store/dbadapter"
+	"cosmossdk.io/store/lockingkv"
 	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
@@ -33,6 +34,7 @@ type Store struct {
 }
 
 var _ types.CacheMultiStore = Store{}
+var _ types.LockingStore = Store{}
 
 // NewFromKVStore creates a new Store object from a mapping of store keys to
 // CacheWrapper objects and a KVStore as the database. Each CacheWrapper store
@@ -63,6 +65,39 @@ func NewFromKVStore(
 	return cms
 }
 
+// NewLockingFromKVStore creates a new Store object from a mapping of store keys to
+// CacheWrapper objects and a KVStore as the database. Each CacheWrapper store
+// is a branched store.
+func NewLockingFromKVStore(
+	store types.KVStore, stores map[types.StoreKey]types.CacheWrapper,
+	keys map[string]types.StoreKey, traceWriter io.Writer, traceContext types.TraceContext,
+) Store {
+	cms := Store{
+		db:           cachekv.NewStore(store),
+		stores:       make(map[types.StoreKey]types.CacheWrap, len(stores)),
+		keys:         keys,
+		traceWriter:  traceWriter,
+		traceContext: traceContext,
+	}
+
+	for key, store := range stores {
+		if cms.TracingEnabled() {
+			tctx := cms.traceContext.Clone().Merge(types.TraceContext{
+				storeNameCtxKey: key.Name(),
+			})
+
+			store = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, tctx)
+		}
+		if kvStoreKey, ok := key.(*types.KVStoreKey); ok && kvStoreKey.IsLocking() {
+			cms.stores[key] = lockingkv.NewStore(store.(types.KVStore))
+		} else {
+			cms.stores[key] = cachekv.NewStore(store.(types.KVStore))
+		}
+	}
+
+	return cms
+}
+
 // NewStore creates a new Store object from a mapping of store keys to
 // CacheWrapper objects. Each CacheWrapper store is a branched store.
 func NewStore(
@@ -72,6 +107,15 @@ func NewStore(
 	return NewFromKVStore(dbadapter.Store{DB: db}, stores, keys, traceWriter, traceContext)
 }
 
+// NewLockingStore creates a new Store object from a mapping of store keys to
+// CacheWrapper objects. Each CacheWrapper store is a branched store.
+func NewLockingStore(
+	db dbm.DB, stores map[types.StoreKey]types.CacheWrapper, keys map[string]types.StoreKey,
+	traceWriter io.Writer, traceContext types.TraceContext,
+) Store {
+	return NewLockingFromKVStore(dbadapter.Store{DB: db}, stores, keys, traceWriter, traceContext)
+}
+
 func newCacheMultiStoreFromCMS(cms Store) Store {
 	stores := make(map[types.StoreKey]types.CacheWrapper)
 	for k, v := range cms.stores {
@@ -127,6 +171,14 @@ func (cms Store) Write() {
 	}
 }
 
+func (cms Store) Unlock() {
+	for _, store := range cms.stores {
+		if s, ok := store.(types.LockingStore); ok {
+			s.Unlock()
+		}
+	}
+}
+
 // Implements CacheWrapper.
 func (cms Store) CacheWrap() types.CacheWrap {
 	return cms.CacheMultiStore().(types.CacheWrap)
@@ -142,6 +194,39 @@ func (cms Store) CacheMultiStore() types.CacheMultiStore {
 	return newCacheMultiStoreFromCMS(cms)
 }
 
+// Implements MultiStore.
+func (cms Store) CacheMultiStoreWithLocking(storeLocks map[types.StoreKey][][]byte) types.CacheMultiStore {
+	stores := make(map[types.StoreKey]types.CacheWrapper)
+	for k, v := range cms.stores {
+		stores[k] = v
+	}
+
+	cms2 := Store{
+		db:           cachekv.NewStore(cms.db),
+		stores:       make(map[types.StoreKey]types.CacheWrap, len(stores)),
+		keys:         cms.keys,
+		traceWriter:  cms.traceWriter,
+		traceContext: cms.traceContext,
+	}
+
+	for key, store := range stores {
+		if rowLocks, ok := storeLocks[key]; ok {
+			cms2.stores[key] = store.(types.LockingCacheWrapper).CacheWrapWithLocks(rowLocks)
+		} else {
+			if cms.TracingEnabled() {
+				tctx := cms.traceContext.Clone().Merge(types.TraceContext{
+					storeNameCtxKey: key.Name(),
+				})
+
+				store = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, tctx)
+			}
+			cms2.stores[key] = cachekv.NewStore(store.(types.KVStore))
+		}
+	}
+
+	return cms2
+}
+
 // CacheMultiStoreWithVersion implements the MultiStore interface. It will panic
 // as an already cached multi-store cannot load previous versions.
 //