Problem: many memory allocs in cache store wrapping

store/CHANGELOG.md

@@ -31,6 +31,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#205](https://github.com/crypto-org-chain/cosmos-sdk/pull/205) Support object store.
 * [#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.
 
 ## v1.1.0 (March 20, 2024)
 

store/cachemulti/store.go

@@ -4,9 +4,6 @@ import (
 	"fmt"
 	"io"
 
-	dbm "github.com/cosmos/cosmos-db"
-
-	"cosmossdk.io/store/dbadapter"
 	"cosmossdk.io/store/tracekv"
 	"cosmossdk.io/store/types"
 )
@@ -23,12 +20,11 @@ const storeNameCtxKey = "store_name"
 // NOTE: a Store (and MultiStores in general) should never expose the
 // keys for the substores.
 type Store struct {
-	db     types.CacheWrap
 	stores map[types.StoreKey]types.CacheWrap
-	keys   map[string]types.StoreKey
 
 	traceWriter  io.Writer
 	traceContext types.TraceContext
+	parentStore  func(types.StoreKey) types.CacheWrap
 }
 
 var _ types.CacheMultiStore = Store{}
@@ -37,29 +33,17 @@ var _ types.CacheMultiStore = Store{}
 // CacheWrapper objects and a KVStore as the database. Each CacheWrapper store
 // is a branched store.
 func NewFromKVStore(
-	store types.CacheWrapper, stores map[types.StoreKey]types.CacheWrapper,
-	keys map[string]types.StoreKey, traceWriter io.Writer, traceContext types.TraceContext,
+	stores map[types.StoreKey]types.CacheWrapper,
+	traceWriter io.Writer, traceContext types.TraceContext,
 ) Store {
 	cms := Store{
-		db:           store.CacheWrap(),
 		stores:       make(map[types.StoreKey]types.CacheWrap, len(stores)),
-		keys:         keys,
 		traceWriter:  traceWriter,
 		traceContext: traceContext,
 	}
 
 	for key, store := range stores {
-		if cms.TracingEnabled() {
-			// only support tracing on KVStore.
-			if kvstore, ok := store.(types.KVStore); ok {
-				tctx := cms.traceContext.Clone().Merge(types.TraceContext{
-					storeNameCtxKey: key.Name(),
-				})
-
-				store = tracekv.NewStore(kvstore, cms.traceWriter, tctx)
-			}
-		}
-		cms.stores[key] = store.CacheWrap()
+		cms.initStore(key, store)
 	}
 
 	return cms
@@ -68,19 +52,35 @@ func NewFromKVStore(
 // NewStore creates a new Store object from a mapping of store keys to
 // CacheWrapper objects. Each CacheWrapper store is a branched store.
 func NewStore(
-	db dbm.DB, stores map[types.StoreKey]types.CacheWrapper, keys map[string]types.StoreKey,
+	stores map[types.StoreKey]types.CacheWrapper,
 	traceWriter io.Writer, traceContext types.TraceContext,
 ) Store {
-	return NewFromKVStore(dbadapter.Store{DB: db}, stores, keys, traceWriter, traceContext)
+	return NewFromKVStore(stores, traceWriter, traceContext)
 }
 
 func newCacheMultiStoreFromCMS(cms Store) Store {
-	stores := make(map[types.StoreKey]types.CacheWrapper)
-	for k, v := range cms.stores {
-		stores[k] = v
+	return Store{
+		stores:       make(map[types.StoreKey]types.CacheWrap),
+		traceWriter:  cms.traceWriter,
+		traceContext: cms.traceContext,
+		parentStore:  cms.getCacheWrap,
 	}
+}
+
+func (cms Store) initStore(key types.StoreKey, store types.CacheWrapper) types.CacheWrap {
+	if cms.TracingEnabled() {
+		// only support tracing on KVStore.
+		if kvstore, ok := store.(types.KVStore); ok {
+			tctx := cms.traceContext.Clone().Merge(types.TraceContext{
+				storeNameCtxKey: key.Name(),
+			})
 
-	return NewFromKVStore(cms.db, stores, nil, cms.traceWriter, cms.traceContext)
+			store = tracekv.NewStore(kvstore, cms.traceWriter, tctx)
+		}
+	}
+	cache := store.CacheWrap()
+	cms.stores[key] = cache
+	return cache
 }
 
 // SetTracer sets the tracer for the MultiStore that the underlying
@@ -118,7 +118,6 @@ func (cms Store) GetStoreType() types.StoreType {
 
 // Write calls Write on each underlying store.
 func (cms Store) Write() {
-	cms.db.Write()
 	for _, store := range cms.stores {
 		store.Write()
 	}
@@ -134,19 +133,23 @@ func (cms Store) CacheMultiStore() types.CacheMultiStore {
 	return newCacheMultiStoreFromCMS(cms)
 }
 
-// GetStore returns an underlying Store by key.
-func (cms Store) GetStore(key types.StoreKey) types.Store {
-	s := cms.stores[key]
-	if key == nil || s == nil {
+func (cms Store) getCacheWrap(key types.StoreKey) types.CacheWrap {
+	store, ok := cms.stores[key]
+	if !ok && cms.parentStore != nil {
+		// load on demand
+		store = cms.initStore(key, cms.parentStore(key))
+	}
+	if key == nil || store == nil {
 		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
 	}
-	return s.(types.Store)
+	return store
 }
 
-func (cms Store) getCacheWrap(key types.StoreKey) types.CacheWrap {
-	store := cms.stores[key]
-	if key == nil || store == nil {
-		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
+// GetStore returns an underlying Store by key.
+func (cms Store) GetStore(key types.StoreKey) types.Store {
+	store, ok := cms.getCacheWrap(key).(types.Store)
+	if !ok {
+		panic(fmt.Sprintf("store with key %v is not Store", key))
 	}
 	return store
 }

store/rootmulti/store.go

@@ -565,7 +565,7 @@ func (rs *Store) CacheMultiStore() types.CacheMultiStore {
 		}
 		stores[k] = store
 	}
-	return cachemulti.NewStore(rs.db, stores, rs.keysByName, rs.traceWriter, rs.getTracingContext())
+	return cachemulti.NewStore(stores, rs.traceWriter, rs.getTracingContext())
 }
 
 // CacheMultiStoreWithVersion is analogous to CacheMultiStore except that it
@@ -627,7 +627,7 @@ func (rs *Store) CacheMultiStoreWithVersion(version int64) (types.CacheMultiStor
 		cachedStores[key] = cacheStore
 	}
 
-	return cachemulti.NewStore(rs.db, cachedStores, rs.keysByName, rs.traceWriter, rs.getTracingContext()), nil
+	return cachemulti.NewStore(cachedStores, rs.traceWriter, rs.getTracingContext()), nil
 }
 
 // GetStore returns a mounted Store for a given StoreKey. If the StoreKey does