Problem: object cache not enabled for accounts

collections/collections.go

@@ -16,6 +16,8 @@ var (
 	ErrEncoding = codec.ErrEncoding
 	// ErrConflict is returned when there are conflicts, for example in UniqueIndex.
 	ErrConflict = errors.New("collections: conflict")
+	// ErrTypeMismatch is returned when there is a type mismatch.
+	ErrTypeMismatch = errors.New("collections: type mismatch")
 )
 
 // KEYS

collections/map.go

@@ -17,7 +17,9 @@ type Map[K, V any] struct {
 	vc codec.ValueCodec[V]
 
 	// store accessor
-	sa     func(context.Context) store.KVStore
+	sa func(context.Context) store.KVStore
+	// cache accessor
+	ca     func(context.Context) store.ObjKVStore
 	prefix []byte
 	name   string
 }
@@ -36,6 +38,7 @@ func NewMap[K, V any](
 		kc:     keyCodec,
 		vc:     valueCodec,
 		sa:     schemaBuilder.schema.storeAccessor,
+		ca:     schemaBuilder.schema.cacheAccessor,
 		prefix: prefix.Bytes(),
 		name:   name,
 	}
@@ -65,7 +68,14 @@ func (m Map[K, V]) Set(ctx context.Context, key K, value V) error {
 	}
 
 	kvStore := m.sa(ctx)
-	return kvStore.Set(bytesKey, valueBytes)
+	if err := kvStore.Set(bytesKey, valueBytes); err != nil {
+		return err
+	}
+
+	if m.ca != nil {
+		m.ca(ctx).Set(bytesKey, value)
+	}
+	return nil
 }
 
 // Get returns the value associated with the provided key,
@@ -77,6 +87,15 @@ func (m Map[K, V]) Get(ctx context.Context, key K) (v V, err error) {
 		return v, err
 	}
 
+	if m.ca != nil {
+		genericValue := m.ca(ctx).Get(bytesKey)
+		v, ok := genericValue.(V)
+		if !ok {
+			return v, fmt.Errorf("%w: key '%s' has type %T, expect %s", ErrTypeMismatch, m.kc.Stringify(key), genericValue, m.vc.ValueType())
+		}
+		return v, nil
+	}
+
 	kvStore := m.sa(ctx)
 	valueBytes, err := kvStore.Get(bytesKey)
 	if err != nil {
@@ -90,6 +109,11 @@ func (m Map[K, V]) Get(ctx context.Context, key K) (v V, err error) {
 	if err != nil {
 		return v, fmt.Errorf("%w: value decode: %s", ErrEncoding, err) // TODO: use multi err wrapping in go1.20: https://github.com/golang/go/issues/53435
 	}
+
+	if m.ca != nil {
+		m.ca(ctx).Set(bytesKey, v)
+	}
+
 	return v, nil
 }
 
@@ -100,6 +124,13 @@ func (m Map[K, V]) Has(ctx context.Context, key K) (bool, error) {
 	if err != nil {
 		return false, err
 	}
+
+	if m.ca != nil {
+		if m.ca(ctx).Has(bytesKey) {
+			return true, nil
+		}
+	}
+
 	kvStore := m.sa(ctx)
 	return kvStore.Has(bytesKey)
 }
@@ -112,6 +143,11 @@ func (m Map[K, V]) Remove(ctx context.Context, key K) error {
 	if err != nil {
 		return err
 	}
+
+	if m.ca != nil {
+		m.ca(ctx).Delete(bytesKey)
+	}
+
 	kvStore := m.sa(ctx)
 	return kvStore.Delete(bytesKey)
 }

collections/schema.go

@@ -21,9 +21,17 @@ type SchemaBuilder struct {
 
 // NewSchemaBuilderFromAccessor creates a new schema builder from the provided store accessor function.
 func NewSchemaBuilderFromAccessor(accessorFunc func(ctx context.Context) store.KVStore) *SchemaBuilder {
+	return NewSchemaBuilderFromAccessorWithCache(accessorFunc, nil)
+}
+
+func NewSchemaBuilderFromAccessorWithCache(
+	accessorFunc func(ctx context.Context) store.KVStore,
+	cacheFunc func(ctx context.Context) store.ObjKVStore,
+) *SchemaBuilder {
 	return &SchemaBuilder{
 		schema: &Schema{
 			storeAccessor:       accessorFunc,
+			cacheAccessor:       cacheFunc,
 			collectionsByName:   map[string]Collection{},
 			collectionsByPrefix: map[string]Collection{},
 		},
@@ -37,6 +45,14 @@ func NewSchemaBuilder(service store.KVStoreService) *SchemaBuilder {
 	return NewSchemaBuilderFromAccessor(service.OpenKVStore)
 }
 
+func NewSchemaBuilderWithCache(service store.KVStoreService, cacheService store.ObjectStoreService) *SchemaBuilder {
+	var cacheAccessor func(context.Context) store.ObjKVStore
+	if cacheService != nil {
+		cacheAccessor = cacheService.OpenObjKVStore
+	}
+	return NewSchemaBuilderFromAccessorWithCache(service.OpenKVStore, cacheAccessor)
+}
+
 // Build should be called after all collections that are part of the schema
 // have been initialized in order to get a reference to the Schema. It is
 // important to check the returned error for any initialization errors.
@@ -126,6 +142,7 @@ var nameRegex = regexp.MustCompile("^" + NameRegex + "$")
 // clients.
 type Schema struct {
 	storeAccessor       func(context.Context) store.KVStore
+	cacheAccessor       func(context.Context) store.ObjKVStore
 	collectionsOrdered  []string
 	collectionsByPrefix map[string]Collection
 	collectionsByName   map[string]Collection
@@ -154,8 +171,16 @@ func NewMemoryStoreSchema(service store.MemoryStoreService) Schema {
 //			return sdk.UnwrapSDKContext(ctx).KVStore(kvStoreKey)
 //	}
 func NewSchemaFromAccessor(accessor func(context.Context) store.KVStore) Schema {
+	return NewSchemaFromAccessorWithCache(accessor, nil)
+}
+
+func NewSchemaFromAccessorWithCache(
+	accessor func(context.Context) store.KVStore,
+	cacheAccessor func(context.Context) store.ObjKVStore,
+) Schema {
 	return Schema{
 		storeAccessor:       accessor,
+		cacheAccessor:       cacheAccessor,
 		collectionsByName:   map[string]Collection{},
 		collectionsByPrefix: map[string]Collection{},
 	}

core/store/service.go

@@ -25,3 +25,7 @@ type TransientStoreService interface {
 	// OpenTransientStore retrieves the transient store from the context.
 	OpenTransientStore(context.Context) KVStore
 }
+type ObjectStoreService interface {
+	// OpenObjKVStore retrieves the ObjKVStore from the context.
+	OpenObjKVStore(context.Context) ObjKVStore
+}

core/store/store.go

@@ -34,3 +34,18 @@ type KVStore interface {
 
 // Iterator is an alias db's Iterator for convenience.
 type Iterator = dbm.Iterator
+
+// ObjKVStore don't support iteration yet to keep it simple.
+type ObjKVStore interface {
+	// Get returns nil iff key doesn't exist. Errors on nil key.
+	Get(key []byte) any
+
+	// Has checks if a key exists. Errors on nil key.
+	Has(key []byte) bool
+
+	// Set sets the key. Errors on nil key or value.
+	Set(key []byte, value any)
+
+	// Delete deletes the key. Errors on nil key.
+	Delete(key []byte)
+}

runtime/store.go

@@ -21,6 +21,18 @@ func (k kvStoreService) OpenKVStore(ctx context.Context) store.KVStore {
 	return newKVStore(sdk.UnwrapSDKContext(ctx).KVStore(k.key))
 }
 
+func NewObjectStoreService(storeKey *storetypes.ObjectStoreKey) store.ObjectStoreService {
+	return &objKVStoreService{key: storeKey}
+}
+
+type objKVStoreService struct {
+	key *storetypes.ObjectStoreKey
+}
+
+func (k objKVStoreService) OpenObjKVStore(ctx context.Context) store.ObjKVStore {
+	return sdk.UnwrapSDKContext(ctx).ObjectStore(k.key)
+}
+
 type memStoreService struct {
 	key *storetypes.MemoryStoreKey
 }

simapp/app.go

@@ -267,7 +267,7 @@ func NewSimApp(
 	}
 
 	tkeys := storetypes.NewTransientStoreKeys(paramstypes.TStoreKey)
-	okeys := storetypes.NewObjectStoreKeys(banktypes.ObjectStoreKey)
+	okeys := storetypes.NewObjectStoreKeys(banktypes.ObjectStoreKey, authtypes.StoreKey)
 	app := &SimApp{
 		BaseApp:           bApp,
 		legacyAmino:       legacyAmino,
@@ -286,7 +286,7 @@ func NewSimApp(
 	bApp.SetParamStore(app.ConsensusParamsKeeper.ParamsStore)
 
 	// add keepers
-	app.AccountKeeper = authkeeper.NewAccountKeeper(appCodec, runtime.NewKVStoreService(keys[authtypes.StoreKey]), authtypes.ProtoBaseAccount, maccPerms, authcodec.NewBech32Codec(sdk.Bech32MainPrefix), sdk.Bech32MainPrefix, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.AccountKeeper = authkeeper.NewAccountKeeperWithCache(appCodec, runtime.NewKVStoreService(keys[authtypes.StoreKey]), runtime.NewObjectStoreService(okeys[authtypes.StoreKey]), authtypes.ProtoBaseAccount, maccPerms, authcodec.NewBech32Codec(sdk.Bech32MainPrefix), sdk.Bech32MainPrefix, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	app.BankKeeper = bankkeeper.NewBaseKeeper(
 		appCodec,

store/cachemulti/store.go

@@ -12,6 +12,9 @@ import (
 // the store which emitted a given trace.
 const storeNameCtxKey = "store_name"
 
+// storesPreAllocation is the pre-allocate size for the stores map
+const storesPreAllocation = 8
+
 //----------------------------------------
 // Store
 
@@ -67,7 +70,7 @@ func NewFromParent(
 	traceWriter io.Writer, traceContext types.TraceContext,
 ) Store {
 	return Store{
-		stores:       make(map[types.StoreKey]types.CacheWrap),
+		stores:       make(map[types.StoreKey]types.CacheWrap, storesPreAllocation),
 		traceWriter:  traceWriter,
 		traceContext: traceContext,
 		parentStore:  parentStore,

x/auth/keeper/keeper.go

@@ -115,13 +115,20 @@ var _ AccountKeeperI = &AccountKeeper{}
 func NewAccountKeeper(
 	cdc codec.BinaryCodec, storeService store.KVStoreService, proto func() sdk.AccountI,
 	maccPerms map[string][]string, ac address.Codec, bech32Prefix, authority string,
+) AccountKeeper {
+	return NewAccountKeeperWithCache(cdc, storeService, nil, proto, maccPerms, ac, bech32Prefix, authority)
+}
+
+func NewAccountKeeperWithCache(
+	cdc codec.BinaryCodec, storeService store.KVStoreService, cacheService store.ObjectStoreService, proto func() sdk.AccountI,
+	maccPerms map[string][]string, ac address.Codec, bech32Prefix, authority string,
 ) AccountKeeper {
 	permAddrs := make(map[string]types.PermissionsForAddress)
 	for name, perms := range maccPerms {
 		permAddrs[name] = types.NewPermissionsForAddress(name, perms)
 	}
 
-	sb := collections.NewSchemaBuilder(storeService)
+	sb := collections.NewSchemaBuilderWithCache(storeService, cacheService)
 
 	ak := AccountKeeper{
 		addressCodec:  ac,