perf(store): lazy CacheMultiStore — defer cachekv alloc until access

sei-cosmos/store/cachemulti/store.go

@@ -3,6 +3,7 @@
 import (
 	"fmt"
 	"io"
+	"sync"
 
 	dbm "github.com/tendermint/tm-db"
 
@@ -21,11 +22,14 @@
 // NOTE: a Store (and MultiStores in general) should never expose the
 // keys for the substores.
 type Store struct {
-	db         types.CacheKVStore
-	stores     map[types.StoreKey]types.CacheWrap
-	keys       map[string]types.StoreKey
-	gigaStores map[types.StoreKey]types.KVStore
-	gigaKeys   []types.StoreKey
+	mu               *sync.RWMutex
+	db               types.CacheKVStore
+	stores           map[types.StoreKey]types.CacheWrap
+	storeParents     map[types.StoreKey]types.CacheWrapper // lazy: not yet wrapped in cachekv
+	keys             map[string]types.StoreKey
+	gigaStores       map[types.StoreKey]types.KVStore
+	gigaStoreParents map[types.StoreKey]types.KVStore // lazy: not yet wrapped in giga cachekv
+	gigaKeys         []types.StoreKey
 
 	traceWriter  io.Writer
 	traceContext types.TraceContext
@@ -37,23 +41,22 @@
 
 // NewFromKVStore 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.
+// is a branched store. Store creation is deferred until access time.
 func NewFromKVStore(
 	store types.KVStore, stores map[types.StoreKey]types.CacheWrapper,
 	gigaStores map[types.StoreKey]types.KVStore,
 	keys map[string]types.StoreKey, gigaKeys []types.StoreKey, traceWriter io.Writer, traceContext types.TraceContext,
 ) Store {
 	cms := newStoreWithoutGiga(store, stores, keys, gigaKeys, traceWriter, traceContext)
 
+	// Defer giga store creation: store parents, create cachekv wrappers on demand
 	cms.gigaStores = make(map[types.StoreKey]types.KVStore, len(gigaKeys))
+	cms.gigaStoreParents = make(map[types.StoreKey]types.KVStore, len(gigaKeys))
 	for _, key := range gigaKeys {
 		if gigaStore, ok := gigaStores[key]; ok {
-			// if key is in gigaStores, use it as the parent store
-			cms.gigaStores[key] = gigacachekv.NewStore(gigaStore, key, types.DefaultCacheSizeLimit)
-		} else {
-			// if not, use regular store as the parent store
-			parent := stores[key].(types.KVStore)
-			cms.gigaStores[key] = gigacachekv.NewStore(parent, key, types.DefaultCacheSizeLimit)
+			cms.gigaStoreParents[key] = gigaStore
+		} else if parent, ok := cms.storeParents[key]; ok {
+			cms.gigaStoreParents[key] = parent.(types.KVStore)
 		}
 	}
 
@@ -62,20 +65,20 @@
 
 func newStoreWithoutGiga(store types.KVStore, stores map[types.StoreKey]types.CacheWrapper, keys map[string]types.StoreKey, gigaKeys []types.StoreKey, traceWriter io.Writer, traceContext types.TraceContext) Store {
 	cms := Store{
+		mu:           &sync.RWMutex{},
 		db:           cachekv.NewStore(store, nil, types.DefaultCacheSizeLimit),
 		stores:       make(map[types.StoreKey]types.CacheWrap, len(stores)),
+		storeParents: make(map[types.StoreKey]types.CacheWrapper, len(stores)),
 		keys:         keys,
 		gigaKeys:     gigaKeys,
 		traceWriter:  traceWriter,
 		traceContext: traceContext,
 		closers:      []io.Closer{},
 	}
 
+	// Defer store creation: store parents, create cachekv wrappers on demand
 	for key, store := range stores {
-		if cms.TracingEnabled() {
-			store = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, cms.traceContext)
-		}
-		cms.stores[key] = cachekv.NewStore(store.(types.KVStore), key, types.DefaultCacheSizeLimit)
+		cms.storeParents[key] = store
 	}
 	return cms
 }
@@ -91,18 +94,108 @@
 }
 
 func newCacheMultiStoreFromCMS(cms Store) Store {
-	stores := make(map[types.StoreKey]types.CacheWrapper)
+	cms.mu.Lock()
+
+	// Materialize all lazy parents on the parent CMS first.
+	// This ensures the child CMS wraps cachekv stores (not raw stores),
+	// so child.Write() writes to the parent's cachekv — not directly to the
+	// underlying commit store, which would bypass the parent's caching layer.
+	for k := range cms.storeParents {
+		cms.ensureStoreLocked(k)
+	}
+	for k := range cms.gigaStoreParents {
+		cms.ensureGigaStoreLocked(k)
+	}
+
+	stores := make(map[types.StoreKey]types.CacheWrapper, len(cms.stores))
 	for k, v := range cms.stores {
 		stores[k] = v
 	}
+
 	gigaStores := make(map[types.StoreKey]types.KVStore, len(cms.gigaStores))
 	for k, v := range cms.gigaStores {
 		gigaStores[k] = v
 	}
 
+	cms.mu.Unlock()
+
 	return NewFromKVStore(cms.db, stores, gigaStores, cms.keys, cms.gigaKeys, cms.traceWriter, cms.traceContext)
 }
 
+// ensureStore lazily creates a cachekv wrapper for a store key on first access.
+// Thread-safe: uses double-checked locking so the common path (already materialized)
+// only takes an RLock.
+func (cms Store) ensureStore(key types.StoreKey) types.CacheWrap {
+	if cms.mu == nil {
+		// No lazy initialization possible (zero-value or eager store).
+		s := cms.stores[key]
+		return s
+	}
+	cms.mu.RLock()
+	if s, ok := cms.stores[key]; ok {
+		cms.mu.RUnlock()
+		return s
+	}
+	cms.mu.RUnlock()
+
+	cms.mu.Lock()
+	defer cms.mu.Unlock()
+	return cms.ensureStoreLocked(key)
+}
+
+// ensureStoreLocked materializes a store. Caller must hold cms.mu write lock.
+func (cms Store) ensureStoreLocked(key types.StoreKey) types.CacheWrap {
+	if s, ok := cms.stores[key]; ok {
+		return s
+	}
+	parent, ok := cms.storeParents[key]
+	if !ok {
+		return nil
+	}
+	var kvParent types.KVStore = parent.(types.KVStore)
+	if cms.TracingEnabled() {
+		kvParent = tracekv.NewStore(kvParent, cms.traceWriter, cms.traceContext)
+	}
+	s := cachekv.NewStore(kvParent, key, types.DefaultCacheSizeLimit)
+	cms.stores[key] = s
+	delete(cms.storeParents, key)
+	return s
+}
+
+// ensureGigaStore lazily creates a giga cachekv wrapper on first access.
+// Thread-safe: uses double-checked locking.
+func (cms Store) ensureGigaStore(key types.StoreKey) types.KVStore {
+	if cms.mu == nil {
+		s := cms.gigaStores[key]
+		return s
+	}
+	cms.mu.RLock()
+	if s, ok := cms.gigaStores[key]; ok {
+		cms.mu.RUnlock()
+		return s
+	}
+	cms.mu.RUnlock()
+
+	cms.mu.Lock()
+	defer cms.mu.Unlock()
+	return cms.ensureGigaStoreLocked(key)
+}
+
+// ensureGigaStoreLocked materializes a giga store. Caller must hold cms.mu write lock.
+func (cms Store) ensureGigaStoreLocked(key types.StoreKey) types.KVStore {
+	if s, ok := cms.gigaStores[key]; ok {
+		return s
+	}
+	parent, ok := cms.gigaStoreParents[key]
+	if !ok {
+		return nil
+	}
+	s := gigacachekv.NewStore(parent, key, types.DefaultCacheSizeLimit)
+	cms.gigaStores[key] = s
+	delete(cms.gigaStoreParents, key)
+	return s
+}
+
 // SetTracer sets the tracer for the MultiStore that the underlying
 // stores will utilize to trace operations. A MultiStore is returned.
 func (cms Store) SetTracer(w io.Writer) types.MultiStore {
@@ -137,17 +230,22 @@
 }
 
 // Write calls Write on each underlying store.
+// Only materialized (accessed) stores need writing — lazy stores have no dirty data.
 func (cms Store) Write() {
+	cms.mu.RLock()
 	cms.db.Write()
 	for _, store := range cms.stores {
 		store.Write()
 	}
+	cms.mu.RUnlock()
 }
 
 func (cms Store) WriteGiga() {
+	cms.mu.RLock()
 	for _, store := range cms.gigaStores {
 		store.(types.CacheKVStore).Write()
 	}
+	cms.mu.RUnlock()
 }
 
 // Implements CacheWrapper.
@@ -176,7 +274,7 @@
 
 // GetStore returns an underlying Store by key.
 func (cms Store) GetStore(key types.StoreKey) types.Store {
-	s := cms.stores[key]
+	s := cms.ensureStore(key)
 	if key == nil || s == nil {
 		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
 	}
@@ -185,23 +283,28 @@
 
 // GetKVStore returns an underlying KVStore by key.
 func (cms Store) GetKVStore(key types.StoreKey) types.KVStore {
-	store := cms.stores[key]
+	store := cms.ensureStore(key)
 	if key == nil || store == nil {
 		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
 	}
 	return store.(types.KVStore)
 }
 
 func (cms Store) GetGigaKVStore(key types.StoreKey) types.KVStore {
-	store := cms.gigaStores[key]
+	store := cms.ensureGigaStore(key)
 	if key == nil || store == nil {
 		panic(fmt.Sprintf("giga kv store with key %v has not been registered in stores", key))
 	}
 	return store
 }
 
 func (cms Store) IsStoreGiga(key types.StoreKey) bool {
-	_, ok := cms.gigaStores[key]
+	cms.mu.RLock()
+	defer cms.mu.RUnlock()
+	if _, ok := cms.gigaStores[key]; ok {
+		return true
+	}
+	_, ok := cms.gigaStoreParents[key]
 	return ok
 }
 
@@ -220,16 +323,42 @@
 
 // SetKVStores sets the underlying KVStores via a handler for each key
 func (cms Store) SetKVStores(handler func(sk types.StoreKey, s types.KVStore) types.CacheWrap) types.MultiStore {
+	cms.mu.Lock()
+	defer cms.mu.Unlock()
+	// Process already-materialized stores
 	for k, s := range cms.stores {
 		cms.stores[k] = handler(k, s.(types.KVStore))
 	}
+	// Process lazy parents (pass the parent KVStore directly to handler)
+	for k, parent := range cms.storeParents {
+		if _, exists := cms.stores[k]; !exists {
+			cms.stores[k] = handler(k, parent.(types.KVStore))
+		}
+	}
+	// All parents have been processed
+	for k := range cms.storeParents {
+		delete(cms.storeParents, k)
+	}
 	return cms
 }
 
 func (cms Store) SetGigaKVStores(handler func(sk types.StoreKey, s types.KVStore) types.KVStore) types.MultiStore {
+	cms.mu.Lock()
+	defer cms.mu.Unlock()
+	// Process already-materialized giga stores
 	for k, s := range cms.gigaStores {
 		cms.gigaStores[k] = handler(k, s)
 	}
+	// Process lazy giga parents
+	for k, parent := range cms.gigaStoreParents {
+		if _, exists := cms.gigaStores[k]; !exists {
+			cms.gigaStores[k] = handler(k, parent)
+		}
+	}
+	// All parents have been processed
+	for k := range cms.gigaStoreParents {
+		delete(cms.gigaStoreParents, k)
+	}
 	return cms
 }