store.go

package cachemulti

import (
	"fmt"
	"io"
	"maps"
	"sync"

	dbm "github.com/cosmos/cosmos-db"

	"cosmossdk.io/store/cachekv"
	"cosmossdk.io/store/dbadapter"
	"cosmossdk.io/store/tracekv"
	"cosmossdk.io/store/types"
)

// storeNameCtxKey is the TraceContext metadata key that identifies
// the store which emitted a given trace.
const storeNameCtxKey = "store_name"

//----------------------------------------
// Store

// Store holds many branched stores.
// Implements MultiStore.
// 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

	traceWriter  io.Writer
	traceContext types.TraceContext
}

// PooledStore is a wrapper around Store that implements the PooledCacheKVStore interface.
// It's used to avoid allocating new Store instances .
type PooledStore struct {
	Store
}

var (
	_ types.CacheMultiStore       = &Store{}
	_ types.PooledCacheMultiStore = &PooledStore{}
)

// 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.
func NewFromKVStore(
	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)
		}
		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(
	db dbm.DB, stores map[types.StoreKey]types.CacheWrapper, keys map[string]types.StoreKey,
	traceWriter io.Writer, traceContext types.TraceContext,
) *Store {
	return NewFromKVStore(dbadapter.Store{DB: db}, stores, keys, traceWriter, traceContext)
}

var storePool = sync.Pool{
	New: func() any {
		return &PooledStore{
			Store: Store{
				stores: make(map[types.StoreKey]types.CacheWrap),
				keys:   make(map[string]types.StoreKey),
			},
		}
	},
}

// newFromKVStorePooled returns a PooledStore object, populated with a mapping of store keys to
// CacheWrapper objects and a KVStore as the database.
func newFromKVStorePooled(
	store types.KVStore, stores map[types.StoreKey]types.CacheWrap,
	traceWriter io.Writer, traceContext types.TraceContext,
) *PooledStore {
	cms := storePool.Get().(*PooledStore)
	cms.traceWriter = traceWriter
	cms.traceContext = traceContext
	for key, store := range stores {
		var cwStore types.CacheWrapper = store
		if cms.TracingEnabled() {
			tctx := cms.traceContext.Clone().Merge(types.TraceContext{
				storeNameCtxKey: key.Name(),
			})

			cwStore = tracekv.NewStore(store.(types.KVStore), cms.traceWriter, tctx)
		}
		cms.stores[key] = cachekv.NewPooledStore(cwStore.(types.KVStore))
	}
	cms.db = cachekv.NewPooledStore(store)
	return cms
}

// Release releases the PooledStore object back to the pool.
func (cms *PooledStore) Release() {
	// clear the stores map
	for k, v := range cms.stores {
		if pStore, ok := v.(*cachekv.PooledStore); ok {
			pStore.Release()
		}
		delete(cms.stores, k)
	}
	for k := range cms.keys {
		delete(cms.keys, k)
	}
	if pStoreDb, ok := cms.db.(*cachekv.PooledStore); ok {
		pStoreDb.Release()
	}
	cms.db = nil
	cms.traceContext = nil
	cms.traceWriter = nil
	storePool.Put(cms)
}

func newCacheMultiStoreFromCMS(cms *Store) *Store {
	stores := make(map[types.StoreKey]types.CacheWrapper)
	for k, v := range cms.stores {
		stores[k] = v
	}

	return NewFromKVStore(cms.db, stores, nil, cms.traceWriter, cms.traceContext)
}

// 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 {
	cms.traceWriter = w
	return cms
}

// SetTracingContext updates the tracing context for the MultiStore by merging
// the given context with the existing context by key. Any existing keys will
// be overwritten. It is implied that the caller should update the context when
// necessary between tracing operations. It returns a modified MultiStore.
func (cms *Store) SetTracingContext(tc types.TraceContext) types.MultiStore {
	if cms.traceContext != nil {
		maps.Copy(cms.traceContext, tc)
	} else {
		cms.traceContext = tc
	}

	return cms
}

// TracingEnabled returns if tracing is enabled for the MultiStore.
func (cms *Store) TracingEnabled() bool {
	return cms.traceWriter != nil
}

// LatestVersion returns the branch version of the store
func (cms *Store) LatestVersion() int64 {
	panic("cannot get latest version from branch cached multi-store")
}

// GetStoreType returns the type of the store.
func (cms *Store) GetStoreType() types.StoreType {
	return types.StoreTypeMulti
}

// Write calls Write on each underlying store.
func (cms *Store) Write() {
	cms.db.Write()
	for _, store := range cms.stores {
		store.Write()
	}
}

// CacheWrap implements CacheWrapper, returns the cache multi-store as a CacheWrap.
func (cms *Store) CacheWrap() types.CacheWrap {
	return cms.CacheMultiStore().(types.CacheWrap)
}

// CacheWrapWithTrace implements the CacheWrapper interface.
func (cms *Store) CacheWrapWithTrace(_ io.Writer, _ types.TraceContext) types.CacheWrap {
	return cms.CacheWrap()
}

// CacheMultiStore implements MultiStore, returns a new CacheMultiStore from the
// underlying CacheMultiStore.
func (cms *Store) CacheMultiStore() types.CacheMultiStore {
	return newCacheMultiStoreFromCMS(cms)
}

// CacheMultiStorePooled returns a PooledCacheMultiStore object from a pool.
func (cms *Store) CacheMultiStorePooled() types.PooledCacheMultiStore {
	return newFromKVStorePooled(cms.db, cms.stores, cms.traceWriter, cms.traceContext)
}

// CacheMultiStoreWithVersion implements the MultiStore interface. It will panic
// as an already cached multi-store cannot load previous versions.
//
// TODO: The store implementation can possibly be modified to support this as it
// seems safe to load previous versions (heights).
func (cms *Store) CacheMultiStoreWithVersion(_ int64) (types.CacheMultiStore, error) {
	panic("cannot branch cached multi-store with a version")
}

// 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 {
		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
	}
	return s.(types.Store)
}

// GetKVStore returns an underlying KVStore by key.
func (cms *Store) GetKVStore(key types.StoreKey) types.KVStore {
	store := cms.stores[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)
}