From 7d01116b7b767231bf20f573d011f72fc7fa15cb Mon Sep 17 00:00:00 2001 From: yihuang Date: Mon, 1 Apr 2024 14:46:42 +0800 Subject: [PATCH] Problem: no API to use the new CoW branch store (#243) * Support RunAtomic API * add unit test --- CHANGELOG.md | 1 + store/CHANGELOG.md | 1 + store/cachekv/store.go | 7 ++--- store/cachekv/store_test.go | 14 +++++++++ store/cachemulti/store.go | 53 ++++++++++++++++++++++++++++++---- store/cachemulti/store_test.go | 30 +++++++++++++++++++ store/internal/btreeadaptor.go | 5 +++- store/types/store.go | 2 ++ types/context.go | 21 ++++++++++++++ 9 files changed, 123 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 464d497b4525..20da25bf0bfd 100644 --- a/CHANGELOG.md +++ b/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] diff --git a/store/CHANGELOG.md b/store/CHANGELOG.md index df94d0bc57f3..850f6a98f19d 100644 --- a/store/CHANGELOG.md +++ b/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) diff --git a/store/cachekv/store.go b/store/cachekv/store.go index df47dfc3e6ae..c11ddd9e2b21 100644 --- a/store/cachekv/store.go +++ b/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. diff --git a/store/cachekv/store_test.go b/store/cachekv/store_test.go index 3c5622355403..140df66e3b94 100644 --- a/store/cachekv/store_test.go +++ b/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 diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go index 6b6aba8dc6bf..3c1164018072 100644 --- a/store/cachemulti/store.go +++ b/store/cachemulti/store.go @@ -25,6 +25,8 @@ type Store struct { traceWriter io.Writer traceContext types.TraceContext parentStore func(types.StoreKey) types.CacheWrap + + branched bool } var _ types.CacheMultiStore = Store{} @@ -43,7 +45,7 @@ func NewFromKVStore( } for key, store := range stores { - cms.initStore(key, store) + cms.stores[key] = cms.initStore(key, store) } return cms @@ -78,9 +80,7 @@ func (cms Store) initStore(key types.StoreKey, store types.CacheWrapper) types.C 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 @@ func (cms Store) GetStoreType() types.StoreType { // 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) CacheMultiStore() types.CacheMultiStore { 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 @@ func (cms Store) GetObjKVStore(key types.StoreKey) types.ObjKVStore { } 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 +} diff --git a/store/cachemulti/store_test.go b/store/cachemulti/store_test.go index 0ea7785bff93..8029282eafc4 100644 --- a/store/cachemulti/store_test.go +++ b/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)) +} diff --git a/store/internal/btreeadaptor.go b/store/internal/btreeadaptor.go index 318c67bfda89..dbc212779177 100644 --- a/store/internal/btreeadaptor.go +++ b/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 { diff --git a/store/types/store.go b/store/types/store.go index 19e45dbda43c..e63ad6d9a79b 100644 --- a/store/types/store.go +++ b/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. diff --git a/types/context.go b/types/context.go index 27df031ada9b..27f9b588842d 100644 --- a/types/context.go +++ b/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{}