perf(context): avoid unnecessary copies in KVStore, TransientStore and CacheContext methods (#14354)

Co-authored-by: testinginprod <testinginprod@somewhere.idk>
Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>

CHANGELOG.md

@@ -58,7 +58,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#13473](https://github.com/cosmos/cosmos-sdk/pull/13473) ADR-038: Go plugin system proposal
 
 ### Improvements
-
+* (types) [#14354](https://github.com/cosmos/cosmos-sdk/pull/14354) - improve performance on Context.KVStore and Context.TransientStore by 40%
 * (crypto/keyring) [#14151](https://github.com/cosmos/cosmos-sdk/pull/14151) Move keys presentation from `crypto/keyring` to `client/keys`
 * (types) [#14163](https://github.com/cosmos/cosmos-sdk/pull/14163) Refactor `(coins Coins) Validate()` to avoid unnecessary map.
 * (signing) [#14087](https://github.com/cosmos/cosmos-sdk/pull/14087) Add SignModeHandlerWithContext interface with a new `GetSignBytesWithContext` to get the sign bytes using `context.Context` as an argument to access state.

types/context.go

@@ -278,12 +278,12 @@ func (c Context) Value(key interface{}) interface{} {
 
 // KVStore fetches a KVStore from the MultiStore.
 func (c Context) KVStore(key storetypes.StoreKey) KVStore {
-	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), c.kvGasConfig)
+	return gaskv.NewStore(c.ms.GetKVStore(key), c.gasMeter, c.kvGasConfig)
 }
 
 // TransientStore fetches a TransientStore from the MultiStore.
 func (c Context) TransientStore(key storetypes.StoreKey) KVStore {
-	return gaskv.NewStore(c.MultiStore().GetKVStore(key), c.GasMeter(), c.transientKVGasConfig)
+	return gaskv.NewStore(c.ms.GetKVStore(key), c.gasMeter, c.transientKVGasConfig)
 }
 
 // CacheContext returns a new Context with the multi-store cached and a new
@@ -291,7 +291,7 @@ func (c Context) TransientStore(key storetypes.StoreKey) KVStore {
 // is called. Note, events are automatically emitted on the parent context's
 // EventManager when the caller executes the write.
 func (c Context) CacheContext() (cc Context, writeCache func()) {
-	cms := c.MultiStore().CacheMultiStore()
+	cms := c.ms.CacheMultiStore()
 	cc = c.WithMultiStore(cms).WithEventManager(NewEventManager())
 
 	writeCache = func() {

types/context_bench_test.go

@@ -0,0 +1,41 @@
+package types_test
+
+import (
+	"testing"
+
+	"github.com/cosmos/cosmos-sdk/store/types"
+	"github.com/cosmos/cosmos-sdk/testutil"
+)
+
+func BenchmarkContext_KVStore(b *testing.B) {
+	key := types.NewKVStoreKey(b.Name() + "_TestCacheContext")
+
+	ctx := testutil.DefaultContext(key, types.NewTransientStoreKey("transient_"+b.Name()))
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = ctx.KVStore(key)
+	}
+}
+
+func BenchmarkContext_TransientStore(b *testing.B) {
+	key := types.NewKVStoreKey(b.Name() + "_TestCacheContext")
+
+	ctx := testutil.DefaultContext(key, types.NewTransientStoreKey("transient_"+b.Name()))
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = ctx.TransientStore(key)
+	}
+}
+
+func BenchmarkContext_CacheContext(b *testing.B) {
+	key := types.NewKVStoreKey(b.Name() + "_TestCacheContext")
+
+	ctx := testutil.DefaultContext(key, types.NewTransientStoreKey("transient_"+b.Name()))
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _ = ctx.CacheContext()
+	}
+}