Problem: no API to use the new CoW branch store

CHANGELOG.md

@@ -44,6 +44,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (baseapp) [#206](https://github.com/crypto-org-chain/cosmos-sdk/pull/206) Support mount object store in baseapp, add `ObjectStore` api in context.
 * (bank) [#237](https://github.com/crypto-org-chain/cosmos-sdk/pull/237) Support virtual accounts in sending coins.
 * (x/bank) [#239](https://github.com/crypto-org-chain/cosmos-sdk/pull/239) Add low level `AddBalance`,`SubBalance` APIs to bank keeper.
+* [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API in `Context` to use new CoW branched cache store.
 
 ## [Unreleased-Upstream]
 

store/CHANGELOG.md

@@ -32,6 +32,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#240](https://github.com/crypto-org-chain/cosmos-sdk/pull/240) Split methods from `MultiStore` into specialized `RootMultiStore`, keep `MultiStore` generic.
 * [#241](https://github.com/crypto-org-chain/cosmos-sdk/pull/241) Refactor the cache store to be btree backed, prepare to support copy-on-write atomic branching.
 * [#242](https://github.com/crypto-org-chain/cosmos-sdk/pull/242) Init cache on cache lazily, save memory allocations.
+* [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API to use new CoW cache store.
 
 ## v1.1.0 (March 20, 2024)
 

store/cachekv/store.go

@@ -53,10 +53,9 @@ func (store *GStore[V]) GetStoreType() types.StoreType {
 // Clone creates a copy-on-write snapshot of the cache store,
 // it only performs a shallow copy so is very fast.
 func (store *GStore[V]) Clone() types.BranchStore {
-	return &GStore[V]{
-		writeSet: store.writeSet.Copy(),
-		parent:   store.parent,
-	}
+	v := *store
+	v.writeSet = store.writeSet.Copy()
+	return &v
 }
 
 // swapCache swap out the internal cache store and leave the current store unusable.

store/cachekv/store_test.go

@@ -447,6 +447,20 @@ func TestIteratorDeadlock(t *testing.T) {
 	defer it2.Close()
 }
 
+func TestBranchStore(t *testing.T) {
+	mem := dbadapter.Store{DB: dbm.NewMemDB()}
+	store := cachekv.NewStore(mem)
+
+	store.Set([]byte("key1"), []byte("value1"))
+
+	branch := store.Clone().(types.CacheKVStore)
+	branch.Set([]byte("key1"), []byte("value2"))
+
+	require.Equal(t, []byte("value1"), store.Get([]byte("key1")))
+	store.Restore(branch.(types.BranchStore))
+	require.Equal(t, []byte("value2"), store.Get([]byte("key1")))
+}
+
 //-------------------------------------------------------------------------------------------
 // do some random ops
 

store/cachemulti/store.go

@@ -25,6 +25,8 @@
 	traceWriter  io.Writer
 	traceContext types.TraceContext
 	parentStore  func(types.StoreKey) types.CacheWrap
+
+	branched bool
 }
 
 var _ types.CacheMultiStore = Store{}
@@ -43,7 +45,7 @@
 	}
 
 	for key, store := range stores {
-		cms.initStore(key, store)
+		cms.stores[key] = cms.initStore(key, store)
 	}
 
 	return cms
@@ -78,9 +80,7 @@
 			store = tracekv.NewStore(kvstore, cms.traceWriter, tctx)
 		}
 	}
-	cache := store.CacheWrap()
-	cms.stores[key] = cache
-	return cache
+	return store.CacheWrap()
 }
 
 // SetTracer sets the tracer for the MultiStore that the underlying
@@ -118,6 +118,9 @@
 
 // Write calls Write on each underlying store.
 func (cms Store) Write() {
+	if cms.branched {
+		panic("cannot Write on branched store")
+	}
 	for _, store := range cms.stores {
 		store.Write()
 	}
@@ -135,9 +138,14 @@
 
 func (cms Store) getCacheWrap(key types.StoreKey) types.CacheWrap {
 	store, ok := cms.stores[key]
-	if !ok && cms.parentStore != nil {
+	if !ok {
 		// load on demand
-		store = cms.initStore(key, cms.parentStore(key))
+		if cms.branched {
+			store = cms.parentStore(key).(types.BranchStore).Clone().(types.CacheWrap)
+		} else if cms.parentStore != nil {
+			store = cms.initStore(key, cms.parentStore(key))
+		}
+		cms.stores[key] = store
 	}
 	if key == nil || store == nil {
 		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
@@ -171,3 +179,36 @@
 	}
 	return store
 }
+
+func (cms Store) Clone() Store {
+	return Store{
+		stores: make(map[types.StoreKey]types.CacheWrap),
+
+		traceWriter:  cms.traceWriter,
+		traceContext: cms.traceContext,
+		parentStore:  cms.getCacheWrap,
+
+		branched: true,
+	}
+}
+
+func (cms Store) Restore(other Store) {
+	if !other.branched {
+		panic("cannot restore from non-branched store")
+	}
+
+	// restore the stores
+	for k, v := range other.stores {
+		cms.stores[k].(types.BranchStore).Restore(v.(types.BranchStore))
+	}
+}
+
+func (cms Store) RunAtomic(cb func(types.CacheMultiStore) error) error {
+	branch := cms.Clone()
+	if err := cb(branch); err != nil {
+		return err
+	}
+
+	cms.Restore(branch)
+	return nil
+}

store/cachemulti/store_test.go

@@ -4,8 +4,12 @@ import (
 	"fmt"
 	"testing"
 
+	dbm "github.com/cosmos/cosmos-db"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/store/dbadapter"
+	"cosmossdk.io/store/internal"
+	"cosmossdk.io/store/internal/btree"
 	"cosmossdk.io/store/types"
 )
 
@@ -22,3 +26,29 @@ func TestStoreGetKVStore(t *testing.T) {
 	require.PanicsWithValue(errMsg,
 		func() { s.GetKVStore(key) })
 }
+
+func TestRunAtomic(t *testing.T) {
+	store := dbadapter.Store{DB: dbm.NewMemDB()}
+	objStore := internal.NewBTreeStore(btree.NewBTree[any](),
+		func(v any) bool { return v == nil },
+		func(v any) int { return 1 },
+	)
+	keys := map[string]types.StoreKey{
+		"abc":  types.NewKVStoreKey("abc"),
+		"obj":  types.NewObjectStoreKey("obj"),
+		"lazy": types.NewKVStoreKey("lazy"),
+	}
+	s := Store{stores: map[types.StoreKey]types.CacheWrap{
+		keys["abc"]:  store.CacheWrap(),
+		keys["obj"]:  objStore.CacheWrap(),
+		keys["lazy"]: nil,
+	}}
+
+	s.RunAtomic(func(ms types.CacheMultiStore) error {
+		ms.GetKVStore(keys["abc"]).Set([]byte("key"), []byte("value"))
+		ms.GetObjKVStore(keys["obj"]).Set([]byte("key"), "value")
+		return nil
+	})
+	require.Equal(t, []byte("value"), s.GetKVStore(keys["abc"]).Get([]byte("key")))
+	require.Equal(t, "value", s.GetObjKVStore(keys["obj"]).Get([]byte("key")).(string))
+}

store/internal/btreeadaptor.go

@@ -6,7 +6,10 @@ import (
 	"cosmossdk.io/store/types"
 )
 
-var _ types.KVStore = (*BTreeStore[[]byte])(nil)
+var (
+	_ types.KVStore    = (*BTreeStore[[]byte])(nil)
+	_ types.ObjKVStore = (*BTreeStore[any])(nil)
+)
 
 // BTreeStore is a wrapper for a BTree with GKVStore[V] implementation
 type BTreeStore[V any] struct {

store/types/store.go

@@ -163,6 +163,8 @@ type RootMultiStore interface {
 type CacheMultiStore interface {
 	MultiStore
 	Write() // Writes operations to underlying KVStore
+
+	RunAtomic(func(CacheMultiStore) error) error
 }
 
 // CommitMultiStore is an interface for a MultiStore without cache capabilities.

types/context.go

@@ -2,6 +2,7 @@ package types
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	abci "github.com/cometbft/cometbft/abci/types"
@@ -398,6 +399,26 @@ func (c Context) CacheContext() (cc Context, writeCache func()) {
 	return cc, writeCache
 }
 
+// RunAtomic execute the callback function atomically, i.e. the state and event changes are
+// only persisted if the callback returns no error, or discarded as a whole.
+// It uses an efficient approach than CacheContext, without wrapping stores.
+func (c Context) RunAtomic(cb func(Context) error) error {
+	evtManager := NewEventManager()
+	cacheMS, ok := c.ms.(storetypes.CacheMultiStore)
+	if !ok {
+		return errors.New("multistore is not a CacheMultiStore")
+	}
+	if err := cacheMS.RunAtomic(func(ms storetypes.CacheMultiStore) error {
+		ctx := c.WithMultiStore(ms).WithEventManager(evtManager)
+		return cb(ctx)
+	}); err != nil {
+		return err
+	}
+
+	c.EventManager().EmitEvents(evtManager.Events())
+	return nil
+}
+
 var (
 	_ context.Context    = Context{}
 	_ storetypes.Context = Context{}